Блог

Я создаю pcolormesh для 2D-уравнения диффузии, создающего 3D-массив x, y и t. Статический 2D-график для конкретного t прост. Как я могу анимировать его на всех временных этапах?

Я создаю свой 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()
 

Три вещи идут не так:

• В animate() , сетка должна быть обновлена. mesh.set_array() устанавливает новые значения. Поскольку внутренняя структура данных нуждается в массиве 1D, массив 2D должен быть «выровнен»: mesh.set_array(rho[:, :, i].ravel()) .
• animate() должен вернуться список измененных элементов. Запятая в return mesh, конце-это способ Python создать «кортеж» только из одного элемента.
• Самая сложная проблема здесь заключается в том, что предпочтительно все изображения используют одно и то же цветовое отображение. vmin указывает, какое значение соответствует «самому низкому» цвету (темно-фиолетовый на карте viridis по умолчанию), в то время vmax как соответствует значению «самого высокого» цвета (желтый в viridis). Если они заданы явно, matplotlib вычислит их как минимум и максимум первого изображения, 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()