如何让热方程模拟循环中的Colormap强度保持恒定？

使用有限差分法模拟二维矩形空间热方程时，每次清除绘图并重绘温度曲面后，Matplotlib会自动校准色图范围。即使后期中心温度接近边界温度，色图仍会将中心显示为高温（如红色），无法基于初始温度分布进行对比。

原代码如下：

# Libraries
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib as mpl

def f(X,Y): # f(x,y,0)
    return 3 - (X**2+Y**2) * np.sin(Y)


def simulation_heat_eq(rect,hs,BC, c, frames, eps = 1e-10):

    """ Simulation of the heat equation on a
    2D rectangular grid using the Finite Difference Method """

    # Figure settings
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    ax.set_facecolor((0.8,0.8,0.8))

    X, Y = np.meshgrid(np.linspace(rect[0],rect[1],hs[0]),
                       np.linspace(rect[2],rect[3],hs[1])) # 2D meshgrid

    # Initial Conditions
    Z_init = f(X,Y) # f(x,y,0)
    zmax = 3

    # Boundary Conditions
    Z_init[0] = np.ones(len(Z_init[0])) * BC[0]
    Z_init[-1] = np.ones(len(Z_init[-1]))* BC[1]
    Z_init[:,0] =  np.ones(len(Z_init[:,0]))* BC[2]
    Z_init[:, -1] = np.ones(len(Z_init[:, -1])) * BC[3]

    # Infitesimals
    dx = (rect[1] - rect[0]) / (hs[0] - 1)
    dt = dx**2/(10*c**2)

    # Looping over Z:s
    for iteration in range(frames):

        # Clearing graph
        ax.cla()
        ax.axes.set_xlim3d(rect[0] + eps, rect[1] - eps)
        ax.axes.set_ylim3d(rect[2] + eps, rect[3] - eps)
        ax.axes.set_zlim3d(0, zmax)
        ax.set_axis_off()

        # Finite Difference Method
        for row in range(1, hs[0]-1):
            for col in range(1,hs[1]-1):
                Z_init[row,col] = Z_init[row,col] + c**2 * dt / (dx)**2  *(Z_init[row+1,col] - 4*Z_init[row,col]
                                               + Z_init[row-1,col] + Z_init[row,col+1]
                                               + Z_init[row,col-1])

        # Plotting surface with slight pause
        #if iteration > 0:
           # ax.view_init(10+5*np.sin(iteration*np.pi/180), iteration/3)
        ax.plot([rect[0], rect[1]], [rect[2], rect[2]], [BC[0], BC[0]], color='black', linewidth=2)
        ax.plot([rect[1], rect[1]], [rect[2], rect[3]], [BC[3], BC[3]], color='black', linewidth=2)
        ax.plot([rect[0], rect[0]], [rect[2], rect[3]], [BC[2], BC[2]], color='black', linewidth=2)
        ax.plot([rect[0], rect[1]], [rect[3], rect[3]], [BC[1], BC[1]], color='black', linewidth=2)

        ax.plot([rect[0], rect[0]], [rect[2], rect[2]], [BC[0], BC[2]], color='black', linewidth=2)
        ax.plot([rect[1], rect[1]], [rect[2], rect[2]], [BC[0], BC[3]], color='black', linewidth=2)
        ax.plot([rect[0], rect[0]], [rect[3], rect[3]], [BC[1], BC[2]], color='black', linewidth=2)
        ax.plot([rect[1], rect[1]], [rect[3], rect[3]], [BC[1], BC[3]], color='black', linewidth=2)
        surf = ax.plot_surface(X, Y, Z_init, alpha = 0.7, cmap =cm.hot)
        if iteration == 0:
            plt.pause(5)
        plt.pause(dt)

simulation_heat_eq(rect = [-5,1,-1,3], hs = [25,25], BC = [2,2,2,2], c = 5, frames = 1000)

核心问题 #

每次调用plot_surface时，Matplotlib会自动根据当前Z_init的数值范围调整色图的vmin和vmax，导致色标无法保持初始的参考基准。

修改步骤 #

1. 固定色图数值范围：在初始化阶段计算初始温度分布的最小值和最大值，后续绘图时强制使用该范围作为色图映射基准。
2. 优化边界线绘制：边界线是固定不变的，将其移到循环外，避免每次清除后重复绘制，提升运行效率。

修改后的代码 #

# Libraries
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib as mpl

def f(X,Y): # f(x,y,0)
    return 3 - (X**2+Y**2) * np.sin(Y)


def simulation_heat_eq(rect,hs,BC, c, frames, eps = 1e-10):

    """ Simulation of the heat equation on a
    2D rectangular grid using the Finite Difference Method """

    # Figure settings
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    ax.set_facecolor((0.8,0.8,0.8))

    X, Y = np.meshgrid(np.linspace(rect[0],rect[1],hs[0]),
                       np.linspace(rect[2],rect[3],hs[1])) # 2D meshgrid

    # Initial Conditions
    Z_init = f(X,Y) # f(x,y,0)
    zmax = 3

    # Boundary Conditions
    Z_init[0] = np.ones(len(Z_init[0])) * BC[0]
    Z_init[-1] = np.ones(len(Z_init[-1]))* BC[1]
    Z_init[:,0] =  np.ones(len(Z_init[:,0]))* BC[2]
    Z_init[:, -1] = np.ones(len(Z_init[:, -1])) * BC[3]

    # 固定色图的数值范围：基于初始温度分布
    z_min = Z_init.min()
    z_max = Z_init.max()

    # Infitesimals
    dx = (rect[1] - rect[0]) / (hs[0] - 1)
    dt = dx**2/(10*c**2)

    # 绘制固定边界线（移到循环外，只绘制一次）
    ax.plot([rect[0], rect[1]], [rect[2], rect[2]], [BC[0], BC[0]], color='black', linewidth=2)
    ax.plot([rect[1], rect[1]], [rect[2], rect[3]], [BC[3], BC[3]], color='black', linewidth=2)
    ax.plot([rect[0], rect[0]], [rect[2], rect[3]], [BC[2], BC[2]], color='black', linewidth=2)
    ax.plot([rect[0], rect[1]], [rect[3], rect[3]], [BC[1], BC[1]], color='black', linewidth=2)

    ax.plot([rect[0], rect[0]], [rect[2], rect[2]], [BC[0], BC[2]], color='black', linewidth=2)
    ax.plot([rect[1], rect[1]], [rect[2], rect[2]], [BC[0], BC[3]], color='black', linewidth=2)
    ax.plot([rect[0], rect[0]], [rect[3], rect[3]], [BC[1], BC[2]], color='black', linewidth=2)
    ax.plot([rect[1], rect[1]], [rect[3], rect[3]], [BC[1], BC[3]], color='black', linewidth=2)

    # Looping over Z:s
    for iteration in range(frames):

        # 只清除曲面对象，保留边界线
        for artist in ax.get_children():
            if isinstance(artist, plt.Surface):
                artist.remove()

        ax.axes.set_xlim3d(rect[0] + eps, rect[1] - eps)
        ax.axes.set_ylim3d(rect[2] + eps, rect[3] - eps)
        ax.axes.set_zlim3d(0, zmax)
        ax.set_axis_off()

        # Finite Difference Method
        for row in range(1, hs[0]-1):
            for col in range(1,hs[1]-1):
                Z_init[row,col] = Z_init[row,col] + c**2 * dt / (dx)**2  *(Z_init[row+1,col] - 4*Z_init[row,col]
                                               + Z_init[row-1,col] + Z_init[row,col+1]
                                               + Z_init[row,col-1])

        # Plotting surface with fixed colormap range
        #if iteration > 0:
           # ax.view_init(10+5*np.sin(iteration*np.pi/180), iteration/3)
        surf = ax.plot_surface(X, Y, Z_init, alpha = 0.7, cmap=cm.hot, vmin=z_min, vmax=z_max)
        if iteration == 0:
            plt.pause(5)
        plt.pause(dt)

simulation_heat_eq(rect = [-5,1,-1,3], hs = [25,25], BC = [2,2,2,2], c = 5, frames = 1000)

关键修改说明 #

• 固定色图范围：添加vmin=z_min, vmax=z_max参数到plot_surface，强制色图始终使用初始温度分布的数值区间，确保不同时刻的温度颜色对比一致。
• 优化清除逻辑：不再使用ax.cla()清除整个坐标轴，而是只移除曲面对象，保留已绘制的边界线，避免重复绘制，提升性能。

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