如何使用现有的 3D 数组为 python pcolormesh 设置动画，其中第三轴是时间步长？

Posted 2023-05-08

【中文标题】如何使用现有的 3D 数组为 python pcolormesh 设置动画，其中第三轴是时间步长？

【英文标题】：How can I animate a python pcolormesh using an existing 3D array where the third axis is the time steps?

【发布时间】：2021-12-10 10:18:36

【问题描述】：

我为 2D 扩散方程创建了一个 pcolormesh，创建了一个由 x、y 和 t 组成的 3D 数组。特定 t 的静态 2D 图很简单。如何在所有时间步上为其设置动画？

我使用以下内容创建我的 3D 数组：

set array for x and y
grid_size = 100
t_iter = 1000

D = .01
length = 1.0
tmax = 1.0

dx = dy = length/grid_size
dt = tmax/t_iter

rho = np.zeros((grid_size, grid_size, t_iter))
#rho[:] = 1.2
rho[grid_size//2, grid_size//2, 0] = 1.2  # set the initial configuration

for n in range(t_iter-1):
    for j in range(grid_size-1):
        for i in range(grid_size-1):
            pxx = rho[i+1,j,n] + rho[i-1,j,n] - 2*rho[i,j,n]
            pyy = rho[i,j+1,n] + rho[i,j-1,n] - 2*rho[i,j,n]
            rho[i,j,n+1] = D*dt*(pxx/dx**2+pyy/dy**2)

我可以使用 pcolormesh（没有标签和东西）为特定的 t 值绘制数据：

plt.pcolormesh(rho[:,:,500])

我试过这个，但它没有“动画”任何东西。我错过了什么？

from matplotlib.animation import FuncAnimation
fig, ax = plt.subplots()

mesh = ax.pcolormesh(rho[:,:,0])

def animate(i):
    ax.pcolormesh(rho[:,:,i])
    
anim = FuncAnimation(fig, animate, interval=10, frames=t_iter-1, repeat=False)

plt.draw()
plt.show()

【参考方案1】：

出了三件事：

在animate() 中，应该更新网格。 mesh.set_array() 设置新值。由于内部数据结构需要一维数组，所以二维数组应该是“raveled”：mesh.set_array(rho[:, :, i].ravel())。 animate() 应返回已更改元素的列表。 return mesh, 中的尾随逗号是 Python 制作仅包含一个元素的“元组”的方式。 这里最棘手的问题是所有图像最好使用相同的颜色映射。 vmin 告诉哪个值映射到“最低”颜色（默认 viridis 映射中的深紫色），而 vmax 对应于“最高”颜色的值（viridis 中的黄色）。如果它们是明确设置的，matplotlib 会将它们计算为

first

图像的最小值和最大值，在这种情况下为 0 和 1.2。这些值不适用于其他图像。通常，整体最小值和最大值给出合适的值。但是，在这种情况下，图像的“有趣”部分的范围要窄得多。我尝试使用第 1 个和第 99 个百分位数，看起来效果不错，但您可能需要调整这些值。

更新后的代码：

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np

grid_size = 100
t_iter = 1000

D = .01
length = 1.0
tmax = 1.0

dx = dy = length / grid_size
dt = tmax / t_iter

rho = np.zeros((grid_size, grid_size, t_iter))
# rho[:] = 1.2
rho[grid_size // 2, grid_size // 2, 0] = 1.2  # set the initial configuration

for n in range(t_iter - 1):
    for j in range(grid_size - 1):
        for i in range(grid_size - 1):
            pxx = rho[i + 1, j, n] + rho[i - 1, j, n] - 2 * rho[i, j, n]
            pyy = rho[i, j + 1, n] + rho[i, j - 1, n] - 2 * rho[i, j, n]
            rho[i, j, n + 1] = D * dt * (pxx / dx ** 2 + pyy / dy ** 2)

fig, ax = plt.subplots()

mesh = ax.pcolormesh(rho[:, :, 0], vmin=np.percentile(rho.ravel(), 1), vmax=np.percentile(rho.ravel(), 99))

def animate(i):
    mesh.set_array(rho[:, :, i].ravel())
    return mesh,

anim = FuncAnimation(fig, animate, interval=10, frames=t_iter, repeat=False)
plt.show()

【讨论】：

感谢您帮助破译动画代码。尽管如此，没有什么是动画的。我得到的只是初始颜色网格的静态框架。

我在 PyCharm 环境中运行它，它确实会产生动画。如果您使用的是 Jupyter，您可能需要像 %matplotlib widget 或 %matplotlib qt 这样的东西来拥有一个非静态的 matplotlib 窗口。 tutorial examples 在您的环境中工作吗？

我正在运行 Jupyter 笔记本环境。更新 conda 包后一切正常。谢谢你。我什至添加了一个颜色条供参考。