python - 在 matplotlib 中绘制不同颜色的箭头

Question

我有一个包含 5 列和一些行的二维数组。不同的列具有以下条目x1 y1 x2 y2 z我想绘制一个从 (x1,y1) 到 (x2,y2) 的箭头，箭头的颜色应该取自与某些内置颜色图相对应的 z 列。我该怎么做 matplotlib/python？

Answer 1

你可以这样做:

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()  

您想使用 scalarMap.to_rgba 将您的 z 值转换为参数以传递给 的 color 选项箭头命令。
您的结果应如下所示:

编辑
如果你也想看到颜色条，事情就有点棘手了。这是一个更新的最小示例:

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() 

注意事项:

• 额外的 import matplotlib as mpl 可以访问 ColorbarBase
• 现在，明确需要指定两个 轴，一个用于箭头，一个用于颜色条。第二组轴应该有一个合理的颜色条大小。
  add_axes 命令采用相对单位的 [left, botton, width, height] 作为输入。所以右侧由 left + width 给出。
• 在第一组轴上绘制箭头，ax，您的初始图形。
• 在第二组轴 axc 上绘制颜色条。将 cmap、规范化、cNorm 和方向作为参数传递。

你的图应该是这样的:

编辑 2

如果你想在箭头上使用不同颜色的边缘，请将 color 更改为 facecolor(或 fc)并指定 edgecolor (EC)。此外，您现在可能想要控制箭头的宽度 ( default = 0.001 ) 和头部的宽度 ( default = 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)