Flotando flechas con diferentes colores en matplotlib.

Tengo una matriz bidimensional con 5 columnas y un número de filas. Las diferentes columnas tienen las siguientes entradas x1 y1 x2 y2 z Quiero trazar una flecha de (x1, y1) a (x2, y2) y el color de la flecha debe tomarse de la columna z correspondiente a algunos mapas de colores incorporados. ¿Cómo puedo hacer esto matplotlib / python?

Puedes hacerlo:

 import numpy as np import matplotlib.pyplot as plt import matplotlib.colors as colors import matplotlib.cm as cmx DATA = np.random.rand(5,5) cmap = plt.cm.jet cNorm = colors.Normalize(vmin=np.min(DATA[:,4]), vmax=np.max(DATA[:,4])) scalarMap = cmx.ScalarMappable(norm=cNorm,cmap=cmap) for idx in range(0,len(DATA[:,1])): colorVal = scalarMap.to_rgba(DATA[idx,4]) plt.arrow(DATA[idx,0], #x1 DATA[idx,1], # y1 DATA[idx,2]-DATA[idx,0], # x2 - x1 DATA[idx,3]-DATA[idx,1], # y2 - y1 color=colorVal) plt.show() 

Desea usar scalarMap.to_rgba para convertir su valor z en un argumento para pasar a la opción de color del comando de arrow .
Su resultado debe verse algo como esto:

introduzca la descripción de la imagen aquí

EDITAR
Si también quieres ver la barra de colores, las cosas son un poco más complicadas. Aquí hay un ejemplo mínimo actualizado:

 import numpy as np import matplotlib.pyplot as plt import matplotlib.colors as colors import matplotlib.cm as cmx import matplotlib as mpl DATA = np.random.rand(5,5) cmap = plt.cm.jet cNorm = colors.Normalize(vmin=np.min(DATA[:,4]), vmax=np.max(DATA[:,4])) scalarMap = cmx.ScalarMappable(norm=cNorm,cmap=cmap) fig = plt.figure() ax = fig.add_axes([0.1, 0.1, 0.7, 0.85]) # [left, bottom, width, height] axc = fig.add_axes([0.85, 0.10, 0.05, 0.85]) for idx in range(0,len(DATA[:,1])): colorVal = scalarMap.to_rgba(DATA[idx,4]) ax.arrow(DATA[idx,0], # x1 DATA[idx,1], # y1 DATA[idx,2]-DATA[idx,0], # x2 - x1 DATA[idx,3]-DATA[idx,1], # y2 - y1 color=colorVal) cb1 = mpl.colorbar.ColorbarBase(axc, cmap=cmap, norm=cNorm,orientation='vertical') plt.show() 

Cosas a tener en cuenta:

  • La import matplotlib as mpl para tener acceso a ColorbarBase
  • Ahora, hay una necesidad explícita de especificar dos ejes, uno para las flechas y otro para la barra de colores. Este segundo conjunto de ejes debe tener un tamaño razonable para la barra de colores.
    El comando add_axes toma [left, botton, width, height] en unidades relativas como entrada. Así que el lado derecho está dado por la left + width .
  • Traza las flechas en el primer conjunto de ejes, ax , tu figura inicial.
  • Grafique el colobar en el segundo conjunto de ejes, axc . Pase el cmap , la normalización, cNorm y una orientación como argumentos.

Tu figura debe verse algo como esto:

introduzca la descripción de la imagen aquí

Editar 2

Si desea un borde de color diferente en las flechas, cambie el color a facecolor (o fc ) y especifique un edgecolor ( ec ). Además, es posible que ahora desee controlar el ancho de la flecha ( predeterminado = 0.001 ) y el ancho de la cabeza ( predeterminado = 3x width ).

 plt.arrow(DATA[idx,0], #x1 DATA[idx,1], # y1 DATA[idx,2]-DATA[idx,0], # x2 - x1 DATA[idx,3]-DATA[idx,1], # y2 - y1 facecolor=colorVal, edgecolor='k', width=0.005, head_width=0.01)