Matplotlib中绘制不同颜色的箭头

9

我有一个二维数组,其中有5列和一些行。不同的列具有以下条目x1 y1 x2 y2 z。我想要从(x1,y1)到(x2,y2)绘制箭头,并且箭头的颜色应该来自于z列对应于一些内置的颜色映射。如何在matplotlib/python中实现这个功能?

1个回答

16
你可以这样做:
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_rgbaz 值转化为传递给 arrow 命令的 color 选项的参数。
你的结果应该类似于这样:

enter image description here

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

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

需要注意的事项:

  • 为了使用ColorbarBase,需要额外导入import matplotlib as mpl
  • 现在需要明确指定两个轴,一个用于箭头,一个用于颜色条。第二组轴应该有一个合理的大小以适应颜色条。
    add_axes命令以相对单位[left, botton, width, height]作为输入。因此右侧由left + width给出。
  • 在第一组轴ax中绘制箭头。
  • 在第二组轴axc上绘制颜色条,将、cNorm和方向作为参数传递。

您的图应该长这样:

enter image description here

编辑2

如果你希望箭头的边缘有不同的颜色,将color改为facecolor(或fc),并指定一个edgecolorec)。此外,您可能需要控制箭头的宽度(默认值=0.001)和箭头头部的宽度(默认值=3x宽度)。

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)

如果您能展示如何同时绘制色条,那将非常有帮助,谢谢。 - lovespeed
有没有办法在箭头周围放置边缘颜色?例如,黄色的箭头除非有深色边缘,否则很难看到。另外,白色箭头怎么办? - zyxue

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接