Realtime update of pcolormesh plot in matplotlib after each iteration of computation

Question

I have the following python code for solving the 2D heat equation. I can run this code for some number of iterations and save all the computations, then use FuncAnimation to animate the iterations. However, what I would prefer to do, is create an initial plot based on the initial conditions, then update this figure after each iteration. I am somewhat able to do this in the below code. However, I must close each figure in order for it to update and what I am seeking is the full animation of the code in the same figure without the need to close each figure after each time integration. I would then like the last figure to remain (not close.)

My current code is:

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

plate_length = 50
max_iter = 25 

alpha = 2
delta_x = 1

delta_t = (delta_x ** 2)/(4 * alpha)
gamma = (alpha * delta_t) / (delta_x ** 2)

# Initialize solution: the grid of u(k, i, j)
u = np.empty((max_iter, plate_length, plate_length))

# Boundary conditions
u_top = 100.0
u_left = 0.0
u_bottom = 0.0
u_right = 0.0

# Set the initial condition
u.fill(0.)

# Set the boundary conditions
u[:, (plate_length-1):, :] = u_top
u[:, :, :1] = u_left
u[:, :1, 1:] = u_bottom
u[:, :, (plate_length-1):] = u_right

# Plot initial condition
cm = plt.pcolormesh(u[0], cmap=plt.cm.jet, vmin=0, vmax=100)
plt.show()

# Do the calculation here, update u for each k iteration
for k in range(0, max_iter, 1):
    for i in range(1, plate_length-1, delta_x):
        for j in range(1, plate_length-1, delta_x):
            u[k + 1, i, j] = u[k][i][j] + \
                             gamma * (u[k][i+1][j] + u[k][i-1][j] + \
                                      u[k][i][j+1] + u[k][i][j-1] - 4*u[k][i][j]) 

# Plot current value of u
    cm = plt.pcolormesh(u[k], cmap=plt.cm.jet, vmin=0, vmax=100)
    plt.show()

Any help you could provide would be greatly appreciated.

Thanks,

Answer 1

I was able to accomplish this by replacing the above code starting with

# Plot initial condition

with the following

# Plot initial condition
fig = plt.figure()

plt.pcolormesh(u[0], cmap=plt.cm.jet, vmin=0, vmax=100)
fig.canvas.draw()
plt.pause(0.01)

# Do the calculation here, update u for each k iteration
for k in range(0, max_iter-1, 1):
    for i in range(1, plate_length-1, delta_x):
        for j in range(1, plate_length-1, delta_x):
            u[k + 1, i, j] = u[k][i][j] + \
                             gamma * (u[k][i+1][j] + u[k][i-1][j] + \
                                      u[k][i][j+1] + u[k][i][j-1] - 4*u[k][i][j]) 

    plt.pcolormesh(u[k], cmap=plt.cm.jet, vmin=0, vmax=100)
    fig.canvas.draw()
    plt.pause(0.01)
    
plt.show()

Answer 2

A few remarks:

• for computational efficiency avoid the for-loops when computing the diffusion terms (very slow). Vectorized forms are by far faster: see the SO question here
• find a code using that here
• note: for most practical applications the timestep 'delta_t' must be smaller and the number of iterations 'max_iter' must be much larger than assumed here when using the explicite time stepping scheme. (This changes when using implicite time stepping.). So, holding the results of each time step in the k-component of the u-array is fine for the very first trials only.

For plotting intermediate results there might be different approaches

(i) plot every timestep mT only, e.g.:

mT = 10
for k in range(max_iter):
   DU = ......
   u = u + DU
   if k % mT == 0:   # modulo-operator
       plot_uField(u)

(ii) display several timesteps in 1 figure: as here or here or here

(iii) think about animation with matplotlib or with mayavi (for 3D and animation, mayavi is more efficient). But computing the final video can be time consuming.

I hope this gives you some ideas.