并行化pcolormesh绘制多AMR块至同一轴时出现空轴问题

我有AMR网格生成的非结构化数据，每个图表包含2个子图，需要把大量AMR块绘制到同一坐标轴上。因为绘图耗时太长，尝试用并行化让每个进程绘制一个块，但运行脚本后得到的是空坐标轴。我猜需要传递axes对象，但不确定具体方法，以下是返回空轴的代码：

import os,sys
import numpy as np
import function_library as fl
from multiprocessing import Pool
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import tqdm
import time as ttime
import pandas
import dask

# x_corner and y_corner have size [numgrid,17]
# var has size [numgrid,16,16]
# numgrid is the number of blocks

for i in range(indxStart,indxEnd+1):

    fig,axs   = plt.subplots(nrows=1, ncols=len(varPlots),figsize=(10,4),layout='constrained', sharey=True)

    def plotVars(igrid):
        # Rho with electron beams
        pcmRho = axs[0].pcolormesh(x_corner[igrid,:],y_corner[igrid,:], (var[igrid,0,:,:]/HMass).T, norm=colors.LogNorm(vmin=vLims[0][0],vmax=vLims[0][1]), cmap=cmaps[0], antialiased=False)
        # Electron beams
        pcmlQ = axs[0].pcolormesh(x_corner[igrid,:],y_corner[igrid,:], var[igrid,2,:,:].T, norm=colors.Normalize(vmin=vLims[2][0] ,vmax=vLims[2][1]), cmap=cmaps[2], antialiased=False)
        pcmcQ = axs[0].pcolormesh(x_corner[igrid,:],y_corner[igrid,:], var[igrid,3,:,:].T, norm=colors.Normalize(vmin=vLims[3][0],vmax=vLims[3][1]), cmap=cmaps[3], antialiased=False)

        # v_y
        pcmVy = axs[1].pcolormesh(x_corner[igrid,:], y_corner[igrid,:], var[igrid,1,:,:].T, norm=colors.Normalize(vmin=vLims[1][0],vmax=vLims[1][1]), cmap=cmaps[1], antialiased=False)

    pool = Pool(numProc)
    r = list(tqdm.tqdm(pool.imap_unordered(plotVars,range(numgrid)),total=numgrid))

    # axes and cbar rho 
    axs[0].set_xlim(xMin,xMax)
    axs[0].set_ylim(yMin,yMax)
    axs[0].set_xlabel('x (Mm)',fontsize=20,weight='bold')
    axs[0].set_ylabel('y (Mm)',fontsize=20,weight='bold')
    axs[0].xaxis.set_tick_params(labelsize=18)
    axs[0].yaxis.set_tick_params(labelsize=18)
    cbarRho = plt.colorbar(pcmRho,ax=axs[0])
    cbarRho.set_label(label=cbTxts[0],weight='bold',fontsize=15)

    cbarlQ = plt.colorbar(pcmlQ,ax=axs[0])
    #cbarlQ.set_label(label=cbTxts[0],weight='bold',fontsize=15)

    cbarcQ = plt.colorbar(pcmcQ,ax=axs[0])
    #cbarlQ.set_label(label=cbTxts[0],weight='bold',fontsize=15)

    #axes and cbar vy
    axs[1].set_xlim(xMin,xMax)
    axs[1].set_ylim(yMin,yMax)
    axs[1].set_xlabel('x (Mm)',fontsize=20,weight='bold')
    axs[1].set_ylabel('')
    axs[1].xaxis.set_tick_params(labelsize=18)
    cbarVy=plt.colorbar(pcmVy,ax=axs[1])
    #cbarVy.set_label(label=cbTxts[1],weight='bold',fontsize=15)  

    plt.suptitle('t = '+str(round(float(time),2))+'s',fontsize=20,weight='bold')

    figName = fileBase+str(i).rjust(4,'0')+'.png'
    plt.savefig(pathFig+figName,dpi=500)
    print(pathFig+figName)
    plt.close()

核心问题 #

Matplotlib的Axes对象不能跨进程共享，子进程里对主进程创建的axs进行绘图操作，不会同步到主进程的Figure对象上，这就是出现空坐标轴的直接原因。

正确思路 #

并行绘图的合理方式是：让每个子进程独立准备对应AMR块的绘图数据，主进程收集所有数据后，统一将所有块绘制到同一坐标轴上——所有绘图操作必须在同一个进程内完成。

修改后的代码 #

import os,sys
import numpy as np
import function_library as fl
from multiprocessing import Pool
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import tqdm
import time as ttime
import pandas
import dask

# x_corner and y_corner have size [numgrid,17]
# var has size [numgrid,16,16]
# numgrid is the number of blocks

# 子进程任务：仅准备当前块的绘图所需数据，不操作坐标轴
def prepare_plot_data(igrid):
    rho_data = (var[igrid,0,:,:]/HMass).T
    lQ_data = var[igrid,2,:,:].T
    cQ_data = var[igrid,3,:,:].T
    vy_data = var[igrid,1,:,:].T
    return (x_corner[igrid,:], y_corner[igrid,:], rho_data, lQ_data, cQ_data, vy_data)

for i in range(indxStart,indxEnd+1):

    fig,axs   = plt.subplots(nrows=1, ncols=len(varPlots),figsize=(10,4),layout='constrained', sharey=True)
    
    # 并行生成所有块的绘图数据
    pool = Pool(numProc)
    all_plot_data = list(tqdm.tqdm(pool.imap_unordered(prepare_plot_data, range(numgrid)), total=numgrid))
    pool.close()
    pool.join()

    # 主进程统一绘制所有AMR块到坐标轴
    pcmRho_list = []
    pcmlQ_list = []
    pcmcQ_list = []
    pcmVy_list = []

    for x_c, y_c, rho_d, lQ_d, cQ_d, vy_d in all_plot_data:
        # 绘制第一个子图的三层数据
        pcmRho = axs[0].pcolormesh(x_c, y_c, rho_d, norm=colors.LogNorm(vmin=vLims[0][0],vmax=vLims[0][1]), cmap=cmaps[0], antialiased=False)
        pcmlQ = axs[0].pcolormesh(x_c, y_c, lQ_d, norm=colors.Normalize(vmin=vLims[2][0] ,vmax=vLims[2][1]), cmap=cmaps[2], antialiased=False)
        pcmcQ = axs[0].pcolormesh(x_c, y_c, cQ_d, norm=colors.Normalize(vmin=vLims[3][0],vmax=vLims[3][1]), cmap=cmaps[3], antialiased=False)
        # 绘制第二个子图
        pcmVy = axs[1].pcolormesh(x_c, y_c, vy_d, norm=colors.Normalize(vmin=vLims[1][0],vmax=vLims[1][1]), cmap=cmaps[1], antialiased=False)
        # 保存绘图对象用于后续生成色条
        pcmRho_list.append(pcmRho)
        pcmlQ_list.append(pcmlQ)
        pcmcQ_list.append(pcmcQ)
        pcmVy_list.append(pcmVy)

    # 配置第一个子图的坐标轴与色条
    axs[0].set_xlim(xMin,xMax)
    axs[0].set_ylim(yMin,yMax)
    axs[0].set_xlabel('x (Mm)',fontsize=20,weight='bold')
    axs[0].set_ylabel('y (Mm)',fontsize=20,weight='bold')
    axs[0].xaxis.set_tick_params(labelsize=18)
    axs[0].yaxis.set_tick_params(labelsize=18)
    # 所有块用同一刻度范围，取最后一个块的绘图对象生成色条即可
    cbarRho = plt.colorbar(pcmRho_list[-1],ax=axs[0])
    cbarRho.set_label(label=cbTxts[0],weight='bold',fontsize=15)

    cbarlQ = plt.colorbar(pcmlQ_list[-1],ax=axs[0])
    #cbarlQ.set_label(label=cbTxts[0],weight='bold',fontsize=15)

    cbarcQ = plt.colorbar(pcmcQ_list[-1],ax=axs[0])
    #cbarlQ.set_label(label=cbTxts[0],weight='bold',fontsize=15)

    # 配置第二个子图的坐标轴与色条
    axs[1].set_xlim(xMin,xMax)
    axs[1].set_ylim(yMin,yMax)
    axs[1].set_xlabel('x (Mm)',fontsize=20,weight='bold')
    axs[1].set_ylabel('')
    axs[1].xaxis.set_tick_params(labelsize=18)
    cbarVy=plt.colorbar(pcmVy_list[-1],ax=axs[1])
    #cbarVy.set_label(label=cbTxts[1],weight='bold',fontsize=15)  

    plt.suptitle('t = '+str(round(float(time),2))+'s',fontsize=20,weight='bold')

    figName = fileBase+str(i).rjust(4,'0')+'.png'
    plt.savefig(pathFig+figName,dpi=500)
    print(pathFig+figName)
    plt.close()

关键说明 #

• 禁止跨进程操作Axes：子进程仅负责数据预处理，把绘图操作全部放在主进程完成，避免进程间的对象同步问题。
• 色条处理逻辑：因为所有AMR块使用同一套刻度范围，用最后一个块的绘图对象生成色条不会影响色条的准确性。
• 进程池管理：添加pool.close()和pool.join()确保所有子进程完成任务后再执行主进程的绘图逻辑。

内容的提问来源于stack exchange，提问作者Forged