[FIXED] How can I animate a python pcolormesh using an existing 3D array where the third axis is the time steps?

Issue

I create a pcolormesh for the 2D diffusion equation creating a 3D array of x, y and t. A static 2D plot for a specific t is straightforward. How can I animate it over all the time steps?

I create my 3D array using the following:

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)

I can plot the data using pcolormesh (without the labels and stuff) for specific values of t:

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

I tried this, but it doesn't "animate" anything. What am I missing?

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()

Solution

Three things are going wrong:

• In animate(), the mesh should be updated. mesh.set_array() sets new values. As the internal data structure needs a 1D array, the 2D array should be "raveled": mesh.set_array(rho[:, :, i].ravel()).
• animate() should return a list of the changed elements. The trailing comma in return mesh, is Python's way to make a "tuple" of just one element.
• The most tricky problem here, is that preferably all images use the same color mapping. vmin tells which value maps to the "lowest" color (dark purple in the default viridis map), while vmax corresponds to the value for the "highest" color (yellow in viridis). If they are net set explicitly, matplotlib will calculate them as the minimum and maximum of the first image, 0 and 1.2 in this case. These values don't go well for the other images. Usually, the overall minimum and maximum give suitable values. But, in this case, the "interesting" part of the image is in a much narrower range. I experimented using the 1st and 99th percentiles, which seem to work well, but you'll probably want to adapt these values.

The updated code:

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()

Answered By - JohanC