使用Matplotlib动态绘制多个图形

April 22, 2020 (最后修改: August 22, 2021)

上一篇文章《GRIB笔记：使用Matplotlib绘制剖面图》介绍如何绘制剖面图。

如果想要绘制多个剖面图进行对比，比如比较不同试验的差异，可以逐个绘制并使用 ImageMagick 等工具拼接到一块。或者使用 matplotlib.pyplot.subplots 在同一个画布中绘制多个图片。

本文介绍如何使用 Matplotlib 绘制共享色表的多个图形，并根据输入数据的数量动态调整绘图个数。

准备数据

为了后续函数调用，对数据进行简单的封装，包含名称、时间范围等描述信息。

grapes_1 = {
    "name": "GRAPES",
    "data": grapes_1_data,
    "var": "gh",
    "length": 24,
    "starttime": starttime,
    "endtime": endtime,
}

一共生成 5 个数据

dataset = (
    grapes_1, 
    grapes_2, 
    grapes_3, 
    grapes_4, 
    test_1,
)

为了使用统一的色表，需要为色表生成数值列表。取数据集中的最大值和最小值，分成 10 等分。

max_value = np.ceil(np.max([m.max() for m in (d["data"] for d in dataset)]))
min_value = np.ceil(np.min([m.min() for m in (d["data"] for d in dataset)]))
levels = np.linspace(min_value, max_value, 10, dtype=int)
levels

array([-56, -45, -34, -23, -12,  -2,   8,  19,  30,  41])

绘制单个剖面图

将《GRIB笔记：使用Matplotlib绘制剖面图》中的绘图方法封装到一个函数中。

其中 levels 是填充图的层次，所有图像都使用相同的层次，就可以共用一个色表。 ax 是 subplots 返回的绘图对象。

def draw_plot(
    data,
    xi,
    yi,
    levels,
    title,
    ax,
):
    xi = np.flip(np.linspace(-90, 90, 121))
    yi = plevs
    z = data
    
    cf = ax.contourf(xi, yi, z, levels=levels, cmap='bwr')
    
    ax.set_xticks([-90, -60, -30, 0, 30, 60, 90])
    ax.xaxis.set_major_formatter(mticker.ScalarFormatter())
    
    ax.invert_yaxis()
    ax.set_yscale("log")
    ax.set_yticks([10, 50, 100, 250, 500, 850])
    ax.yaxis.set_major_formatter(mticker.ScalarFormatter())
    ax.yaxis.set_minor_locator(mticker.NullLocator())
    
    ax.set_title(title)
    return cf

绘制多个图形

根据数据集大小，创建绘图对象。每行固定 2 个图形，计算需要的行数。

ncols = 2
nrows = int(np.ceil(len(dataset)/ncols))
fig, axes = plt.subplots( 
    ncols=ncols, 
    nrows=nrows,
    figsize=(ncols * 8, nrows * 5),
)

为了动态绘制图像，使用 np.nditer 创建 axes 对象的迭代器

it = np.nditer(axes, flags=["refs_ok"], op_flags=["readwrite"])

为数据集中的每个数据绘制剖面图

for d in dataset:
    cf = draw_plot(
        data=d["data"],
        xi=xi,
        yi=yi,
        levels=levels,
        title=f'Zonal Mean BIAS: {var:>5}  {d["length"]}hr  {d["starttime"].strftime("%Y%m%d")}-{d["endtime"].strftime("%Y%m%d")}{d["name"]:>9}',
        ax=it[0].item(0)
    )
    it.iternext()

如果有剩余的绘图位置，Matplotlib 会默认显示一个空的坐标轴图形，使用下面的代码关闭图形显示。

for i in it:
    i.item(0).axis('off')

绘制统一的色表

fig.colorbar(
    cf, 
    ax=axes, 
    location='bottom', 
    shrink=0.6,
)

上面的图形还有许多问题，比如色表位置不太合适，多个图形可能会重叠等等，还需进一步优化。

参考

《GRIB笔记：使用Matplotlib绘制剖面图》