0

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,

calmymail
  • 127
  • 8
  • Please see https://stackoverflow.com/questions/18797175/animation-with-pcolormesh-routine-in-matplotlib-how-do-i-initialize-the-data – Jody Klymak Jan 09 '21 at 16:02

2 Answers2

1

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.

pyano
  • 1,885
  • 10
  • 28
  • This is great information and will make use of many of these comments. My basic problem is in your statement (i), which is also in my code within my "k" loop, `cm = plt.pcolormesh(u[k], cmap=plt.cm.jet, vmin=0, vmax=100)`. My issue is I don't want to have to close a figure then have the new figure pop up. I want one figure that updates with the new values of "u". For line plots, I would simply state `set_xdata` or `set_ydata` and my plot would get updated with the current data. How do I do this with `pcolormesh`? – calmymail Jan 09 '21 at 14:59
  • 1
    I hve not tried myself: What happens when you replace in this example (https://www.delftstack.com/howto/matplotlib/how-to-automate-plot-updates-in-matplotlib/ ) line1, = ax.plot(x, y) with your cm = plt.pcolormesh(u[k], cmap=plt.cm.jet, vmin=0, vmax=100) ? – pyano Jan 11 '21 at 13:53
  • Unfortunately, this only works for line plots as it updates the data with the set_xdata and set_ydata calls. Would be good if pcolormesh had something like this. – calmymail Jan 11 '21 at 18:30
0

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()
calmymail
  • 127
  • 8