How to create a dynamic plot in python jupyter notebook using a for loop?

Question

I am following this lesson on solving the heat equation in python. The lesson states that after solving the heat equation, we can visualize the solution by simply calling pyplot.plot() inside the loop for an animated plot, where the code below will dynamically plot the temperature at every point at every time, resulting in an animated plot (an example of the animated plot is provided in the lesson post).

import numpy
from matplotlib import pyplot

length = 2
k = .466
temp_left = 200
temp_right = 200

total_time = 4

dx = .1
x_vec = numpy.linspace(0, length, int(length/dx))

dt = .0001
t_vec = numpy.linspace(0, total_time, int(total_time/dt))

u = numpy.zeros([len(t_vec), len(x_vec)])

u[:, 0] = temp_left
u[:, -1] = temp_right

for t in range(1, len(t_vec)-1):
    for x in range(1, len(x_vec)-1):
        u[t+1, x] = k * (dt / dx**2) * (u[t, x+1] - 2*u[t, x] +
                                        u[t, x-1]) + u[t, x]

    pyplot.plot(x_vec, u[t], 'black')
    pyplot.pause(.001)
    pyplot.cla()


pyplot.plot(x_vec, u[0])
pyplot.ylabel("Temperature (C˚)")
pyplot.xlabel("Distance Along Rod (m)")
pyplot.show()

I am just learning to code, and this is all new to me. I am working in a jupyter notebook version 7.1.2. I placed the code into the jupyter notebook cell, ran the cell, and was expecting a single dynamic plot. Instead, what I got was many static plots.

I've looked at the examples found here to see if I can modify the code to work for my application. In trying every example I would get a blank figure, like so:

Answer 1

Below is a working example of your code. I based it off of a very similar question here. Specifically see BlackHC's answer since I could not get the other answers to work. You will also need to have ipympl installed. %matplotlib widget (magic command) only needs to be run once per notebook. The magic command sets the plotting backend. It takes a very long time to run with the current number of iterations and sleep time.

Edit: Per the suggestion in the comment. This will work in JupyterLab and JupyterNotebook 6+ (6 was released in 2020 and BlackHC's answer here was posted in 2020). You can %matplotlib widget or %matplotlib ipympl per the ipyml website

# install ipympl first
%matplotlib widget
import numpy as np
import matplotlib.pyplot as plt
from IPython import display
import time


length = 2
k = 0.466
temp_left = 200
temp_right = 200

total_time = 4

dx = 0.1
dt = 0.0001
x_vec = np.linspace(0, length, int(length / dx))

t_vec = np.linspace(0, total_time, int(total_time / dt))

u = np.zeros([len(t_vec), len(x_vec)])

u[:, 0] = temp_left
u[:, -1] = temp_right


fig, ax = plt.subplots(1, 1)
hdisplay = display.display("", display_id=True)

ax.set_xlabel("Temperature (C˚)")
ax.set_ylabel("Distance Along Rod (m)")
for t in range(1, len(t_vec) - 1):
    for x in range(1, len(x_vec) - 1):
        u[t + 1, x] = (
            k * (dt / dx**2) * (u[t, x + 1] - 2 * u[t, x] + u[t, x - 1]) + u[t, x]
        )
    plt.cla()
    ax.plot(x_vec, u[t], "black")
    hdisplay.update(fig)
    time.sleep(0.00001)

plt.close(fig)