使用matplotlib的quiver绘制二维箭头图

使用ax.quiver绘制二维箭头图

1. matplotlib的quiver函数的调用方式

quiver函数是axes类的成员函数，其基本调用方式为：

quiver([X, Y], U, V, [C], **kwargs)
[X,Y]是箭头的位置，U,V是箭头的方向，C是箭头颜色。
具体而言，C是一个与X,Y的尺寸相同的数组，每个元素记录与箭头颜色相关的值。这些值通过norm和cmap参数映射到真正的颜色值。设置了C参数之后，color参数不生效。
color是直接设置箭头颜色的参数，若它为一个值，则表示所有箭头都是同一种色彩；也可以是与X,Y的尺寸相同的数组，记录每个箭头的颜色。
箭头方向参数angles可以取'uv'（默认）或'xy'。'uv’表示箭头指向采用屏幕坐标系中的方向，与X,Y所在坐标系的定义无关；'xy’表示箭头方向在数据坐标系中，箭头从(x,y)指向(x+u,y+v)。
箭头的尺寸参数包括：
scale_unit箭头长度的单位，可以取'width', 'height', 'dots', 'inches', 'x', 'y', 'xy'。例如，(u,v)=(1,0)，scale取1.0，当scale_unit取’width’时，箭头的长度是轴的宽度；当scale_unit取’height’时，箭头长度是轴的高度；
长度，用scale参数来调节。若不设置scale参数，则matplotlib自动调整所有箭头的长度；若设置scale的值，值越小，箭头的显示越长。scale取1时，箭头的长度等于U,V的计算值×scale_unit

通过下面的例子认识scale_units：

import matplotlib.pyplot as plt

ax.quiver(0,25,0.5,0,scale_units='width',scale=1.0, color='red')
ax.quiver(0,50,1,0,scale_units='width',scale=1.0, color='red')
ax.quiver(0,75,1.5,0,scale_units='width',scale=1.0, color='red')

ax.quiver(25,0,0,0.5,scale_units='height',scale=1.0, color='green')
ax.quiver(50,0,0,1,scale_units='height',scale=1.0, color='green')
ax.quiver(75,0,0,1.5,scale_units='height',scale=1.0, color='green')

ax.set_xlim(0,100)
ax.set_ylim(0,200)
ax.set_aspect(1)
plt.show()

运行上面的代码，会得到三根红色的水平箭头和三根绿色的垂直箭头。他们的长度分别是axes（坐标轴）的x范围和y范围的0.5倍、1倍和1.5倍。

通过下面的例子认识angles：

ax.quiver(50,50,0.5,0,scale_units='width',scale=1.0, color='red', label='Angles=uv')
ax.quiver(50,50,0.5,0,scale_units='width',scale=1.0, color='blue', angles='xy', label='Angles=xy')
ax.quiver(50,0,0,1,scale_units='height',scale=1.0, color='green')
ax.set_xlim(0,100)
ax.set_ylim(0,200)
ax.invert_xaxis()
ax.legend()
ax.set_aspect(1)
plt.show()

运行结果如图所示，其中红色箭头、绿色箭头为默认的angles参数的绘制效果；蓝色箭头为angles='xy’时的绘制效果，即箭头方向受数据坐标系影响；

1.1循环绘制与批量绘制

当我们有大量箭头想绘制时，一般会选择在for循环中不断调用quiver完成绘制。这种方式虽然便于理解，但是效率太低。当待绘制的箭头数量太多时（1,000,000量级以上），显示绘图结果的速度极慢。

循环绘制的代码如下：

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

grid_cols=80
grid_rows=100
grid_u = np.random.randn(grid_rows, grid_cols)
grid_v = np.random.randn(grid_rows, grid_cols)
grid_val = np.sqrt(grid_u**2+grid_v**2)

jet_colormap = mpl.colormaps['jet']# 0 is blue and 1 is red.

# call quiver in the loop
ax = plt.axes()
for i in range(grid_rows):
    for j in range(grid_cols):
        ax.quiver(i,j,grid_u[i,j],grid_v[i,j],grid_val[i,j],scale=1, scale_units='xy',cmap=jet_colormap, norm=plt.Normalize(0.0,1.0))
            

ax.set_aspect(1)
plt.show()

所幸，quiver函数可以仅调用一次，批量绘制所有箭头数据，这样的速度会远远快于循环绘制。使用这种方式之前，需要实现把箭头数据整理到数组中。一般使用np.meshgrid来组织X,Y数据。
批量绘制的代码如下：

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

grid_cols=80
grid_rows=100
grid_u = np.random.randn(grid_rows, grid_cols)
grid_v = np.random.randn(grid_rows, grid_cols)
grid_val = np.sqrt(grid_u**2+grid_v**2)

jet_colormap = mpl.colormaps['jet']# 0 is blue and 1 is red.

grid_x, grid_y = np.meshgrid(np.arange(grid_cols),np.arange(grid_rows))
ax.quiver(grid_x,grid_y,grid_u,grid_v,grid_val,scale=1,scale_units='xy', cmap=jet_colormap, norm=plt.Normalize(0.0,1.0))
ax.set_aspect(1)
plt.show()

分别运行它们，可以感受到绘制速度快慢。

1.2用colormap和Normalize函数来调节箭头颜色

Matplotlib库内置了一系列的颜色映射，包括亮度渐变色、冷暖过渡色、循环色等、如图所示：

如果想获得其中的'jet'颜色映射，则：

import matplotlib as mpl

jet_colormap = mpl.colormaps['jet']

或者直接在需要颜色映射的地方调用，例如：

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt


grid_cols=80
grid_rows=100
grid_u = np.random.randn(grid_rows, grid_cols)*2
grid_v = np.random.randn(grid_rows, grid_cols)*2
grid_val = np.sqrt(grid_u**2+grid_v**2)

jet_colormap = mpl.colormaps['jet']

grid_x, grid_y = np.meshgrid(np.arange(grid_cols),np.arange(grid_rows))
ax = plt.axes()
ax.quiver(grid_x,grid_y,grid_u,grid_v,grid_val,scale=1,scale_units='xy', cmap='jet',#或cmap=jet_colormap，都对
           norm=plt.Normalize(0.0,1.0))# 可以试试norm=plt.Normalize(0.0,4.0)  norm=plt.Normalize(0.0,8.0)
plt.show()

上面的代码中，plt.Normalize(vmin, vmax,clip)函数将数据按照[vmin,vmax]线性地归化到[0,1]区间。对于超出[vmin,vmax]的数据，若clip值为False，则依然将这些数据按[vmin,vmax]->[0,1]的线性映射求得值；若clip值为True，则大于vmax的统一映射为1，小于vmin的值统一映射为0。
vmin, vmax缺省时，则根据给定数据计算最小值和最大值来初始化这两个值。代码如下：

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

grid_cols=80
grid_rows=100
grid_u = np.random.randn(grid_rows, grid_cols)*2
grid_v = np.random.randn(grid_rows, grid_cols)*2
grid_val = np.sqrt(grid_u**2+grid_v**2)

jet_colormap = mpl.colormaps['jet']

grid_x, grid_y = np.meshgrid(np.arange(grid_cols),np.arange(grid_rows))
ax = plt.axes()
ax.quiver(grid_x,grid_y,grid_u,grid_v,grid_val,scale=1,scale_units='xy', cmap='jet',#或cmap=jet_colormap，都对
           norm=plt.Normalize())#或norm=mpl.colors.LogNorm(1.0,10.0)，对数标准化
plt.show()

了解更多可参考：
选择颜色映射
创建颜色映射
Matplotlib中的各种Normalize方法

2.用numpy的ma.masked_where过滤掉不希望绘制的箭头

有时，格网数据中存在我们不想绘制的箭头，例如值大于0.5的箭头，应该怎么做呢？第一反应写下来的代码如下所示：

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

grid_cols=40
grid_rows=60
grid_u = np.random.randn(grid_rows, grid_cols)*2
grid_v = np.random.randn(grid_rows, grid_cols)*2
grid_val = np.sqrt(grid_u**2+grid_v**2)

jet_colormap = mpl.colormaps['jet']# 0 is blue and 1 is red.

# call quiver in the loop
ax = plt.axes()
for i in range(grid_rows):
    for j in range(grid_cols):
        if grid_val[i,j] < 0.5:# 仅绘制值小于0.5的箭头
            ax.quiver(i,j,grid_u[i,j],grid_v[i,j],scale_units='xy',cmap=jet_colormap, width=0.003, headwidth = 0.003)
            
ax.set_aspect(1)# 设置x\y轴比例一致
plt.show()

这又回到逐个箭头调用quiver的龟速模式了！幸好，quiver函数支持输入masked_array类型的数据，masked_array中除了数据数组之外，还有一个标记是否mask的布尔类型数组，来标记某些数据元素。
如此一来，我们只需要在调用quiver之前，按照我们的标准调用np.ma.masked_where()来标记不想绘制的箭头即可。
np.ma.masked_where(condition, a)返回的是标记mask之后的数组，第一个参数condition是条件数组，大小与a一致，即满足这个条件的a的元素被mask；第二个参数a是输入的待标记数组。第三个参数copy默认为Ture，返回a的一个复制品，若为False，则返回a的一个视图(view)。
改进的代码如下：

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

grid_cols=40
grid_rows=60
grid_u = np.random.randn(grid_rows, grid_cols)*2
grid_v = np.random.randn(grid_rows, grid_cols)*2
grid_val = np.sqrt(grid_u**2+grid_v**2)

jet_colormap = mpl.colormaps['jet']

# 过滤值大于0.5的箭头
grid_val = np.ma.masked_where(grid_val>0.5, grid_val)

# plus 0.5 to make the quiver start from center of a grid.
grid_x, grid_y = np.meshgrid(np.arange(grid_cols),np.arange(grid_rows))
# 一次性绘制所有箭头
ax = plt.axes()
ax.quiver(grid_x,grid_y,grid_u,grid_v,grid_val,scale=1,scale_units='xy', cmap=jet_colormap)
# 设置x\y轴比例一致
ax.set_aspect(1)
plt.show()

参考

使用quiver函数的其他高阶参数：

quiver绘制箭头进阶

matplotlib给出的箭头绘制示例代码：
quiver官方示例

关于使用numpy的masked_array来表达有标记的数组，可参考：
masked_where
Masked array operations