python python-3.x matplotlib jupyter-notebook matplotlib-animation

Matplotlib animation script runs fine first time but gets stuck on re-runs until kernel is restarted

I'm trying to run an animation where a scatter plot is updated every timestep. I am working in a jupyter notebook in vscode (if relevant). The animation script, given below, runs great the first time but it doesn't execute as intended on subsequent runs. The problem is only fixed when I restart the kernel. My work requires me to constantly modify this animation code and hence need to rerun it multiple times. Restarting the kernel every time is cumbersome.

In following code (edited to make executable), I am plotting object locations (block_centroids) using a scatter plot and updating them using the displacements at time timepoints[i] (given by fields[i]).

import numpy as np
import matplotlib
%matplotlib ipympl
%matplotlib widget
import matplotlib.pyplot as plt
from matplotlib import animation

# Toy example

timepoints = np.arange(0,1,0.02)    # 100 frames

xs = np.arange(1,11)
ys = np.arange(1,11)
block_centroids = np.dstack(np.meshgrid(xs, ys)).reshape(-1, 2) 
# (N, 2) initial location array for N = 100 objects

rand_amp = np.random.rand(100, 2)/10
fields = [rand_amp * np.sin(i/(20*np.pi)) for i in range(len(timepoints))] 
# sinusoidal displacements for each object with random amplitudes

fig, axes = plt.subplots(figsize=(9,5), tight_layout = True)

scat = axes.scatter(xs, ys)
axes.set_aspect('equal', adjustable='box')

def update(val):
    if val == len(timepoints) - 1:
        print('closing!')
        exit()

    scat.set_offsets(block_centroids + fields[val])
    fig.suptitle(fr'time $t = {timepoints[val]:.2f}$s')
    fig.canvas.draw_idle()
    return scat

def init():
    pass

ani = animation.FuncAnimation(fig=fig, func=update, init_func=init, frames=len(timepoints), blit = True, interval=100)

After the first time I run this code block, the re-run either gets stuck at some midway time (between 0s and 1s) and prematurely prints 'closing!', or the code keeps running seemingly forever with no output and no error until I try to stop it and the kernel times out.

I have tried multiple things without completely understanding why they might work. My guess is that the animation-related objects are not reset and somehow the update() function gets called with the wrong (old) value and my attempts have mostly tried to close/delete the old instances.

I added the if statement inside the update() function in an attempt to close the animation after a run. Earlier, instead of the exit() function, I was using plt.close(fig) which led to the animation getting stuck at a midway time and printing 'closing!' output repeatedly.
I tried adding plt.cla(), plt.clf() and plt.close() at the start of the code block to close any open instances of fig and axes but that didn't help.
I tried deleting the fig and ani objects using del command to remove their previous instances in hopes of getting a fresh start for every run but that didn't work.
Earlier even the first runs of the animation were having problems. I added Agg backend after reading a stackoverflow answer recommending GTKAgg for handling animations better. I couldn't use GTKAgg because I don't have cairo, so I settled for the closest thing and that helped is making the first run successful.

What is not working here as intended? What can I try to debug this? I tried printing val at every function call, to check what is happening during reruns but then the code kept running without any output.

Solution

This code should let you see if VSCode is capable of what I am describing with making a controller that you can use instead of re-rerunning all the calculations. (This is no guarantee it will work with your data because you still need to save all the frames as part of the notebook.) However, it will allow you to understand maybe how you could in theory approach this in VSCode if the same code works.

It makes three points in a scatter plot.

A Jupyter notebook with related code already run is here. (You'll note that even though that is a 'static' rendering of the code of the notebook, the controller at the bottom of the nbviewer-rendered page works, see more about this below the second code block.)
Demonstrate making these two animations in active Jupyter: Click here to launch that notebook in an active, temporary Jupyter session served via MyBinder.org with the environment already having Python everything needed to make the animation with a widget work in the notebook installed and working.

1. This first code block corresponds to top plot in the linked notebooks and it just runs if you run it and then stops

Plot animation that you can run to repeat is below. Make this code block a new cell in a new notebook in your VSCode all its own and try it:

%matplotlib ipympl
# based on OP's code & code currently next to 'In [6]' at https://nbviewer.org/github/fomightez/animated_matplotlib-binder/blob/main/index.ipynb , adapted from https://discourse.jupyter.org/t/matplotlib-animation-not-appearing-in-jupyter-notebook/24938/3?u=fomightez
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

tstep = 1
#xs, ys = block_centroids.T
xs, ys = (1,4,6),(2,5,9)
timepoints = list(zip(xs,ys))
points = list(zip(xs,ys))
num_steps = len(timepoints)

#print(points)
#print(np.arange(0, num_steps))

fig = plt.figure(figsize=(9,5))
ax = fig.add_subplot()
fig.subplots_adjust(left=0.1, right=0.85)


def animate(frame):
    step_num = frame % (num_steps)
    ax.scatter(points[step_num][0],points[step_num][1], c='skyblue', s=60)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(f"frame: {frame} ; step_num: {step_num} ; {points[step_num][0]}; {points[step_num][1]}")  
    ax.grid(True)
    
    return ax

fig.suptitle(f"animation with {len(timepoints)} timepoints and setting of tsteps = {tstep}.")
#ani = FuncAnimation(fig, animate, frames= np.arange(0, num_steps), interval = 1000, repeat = False)
ani = FuncAnimation(fig, animate, frames= 3, interval = 1000, repeat = False)
ani;

That code block above is essentially our starting point and is something along the lines of what OP was posting about trying to implement.
However, the matplotlib FuncAnimation() has ability to make frames that can run via a controller. This makes replay of the animation easier! The next cell in the associated notebook and the code below is to show how just slight changes to the above code can enable this ability.

2. Add playback widget controller: the block below corresponds to second cell in linked notebook

So you can make a controller, you can just add a couple lines and remove the ending ; (Note the warning at the top of ignore the extra static frame or re-run), make this a cell in a notebook all on its own and try running it:

# If upon first running, it shows a non-interactive, single static shot of the plot below the interactive one with the widget controller,
# JUST RE-RUN TWICE. Re-run usually fixes that display quirk.
%matplotlib ipympl
# based on OP's code & code currently next to 'In [6]' at https://nbviewer.org/github/fomightez/animated_matplotlib-binder/blob/main/index.ipynb , adapted from https://discourse.jupyter.org/t/matplotlib-animation-not-appearing-in-jupyter-notebook/24938/3?u=fomightez
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
plt.rcParams["animation.html"] = "jshtml"
plt.ioff() #needed so the second time you run it you get only single plot

tstep = 1
#xs, ys = block_centroids.T
xs, ys = (1,4,6),(2,5,9)
timepoints = list(zip(xs,ys))
points = list(zip(xs,ys))
num_steps = len(timepoints)

#print(points)
#print(np.arange(0, num_steps))

fig = plt.figure(figsize=(9,5))
ax = fig.add_subplot()
fig.subplots_adjust(left=0.1, right=0.85)


def animate(frame):
    step_num = frame % (num_steps)
    ax.scatter(points[step_num][0],points[step_num][1], c='skyblue', s=60)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(f"frame: {frame} ; step_num: {step_num} ; {points[step_num][0]}; {points[step_num][1]}")  
    ax.grid(True)
    
    return ax

fig.suptitle(f"animation with {len(timepoints)} timepoints and setting of tsteps = {tstep}.")
#ani = FuncAnimation(fig, animate, frames= np.arange(0, num_steps), interval = 1000, repeat = False)
ani = FuncAnimation(fig, animate, frames= 3, interval = 1000, repeat = False)
ani

Two background references for output of FuncAnimation()-generated animation to frames with javascript based controller widget or to HTML5 that better explain it and were where I learned about it:

Those animations frames and associated controller can be viewed in nbviewer using the controller widget, too!

See for yourself:
A Jupyter notebook with related code already run is here. Hit the play button there at that link to play the pre-made animation!
You'll note that even though that is a 'static' rendering of the code of the notebook, the controller at the bottom of the nbviewer-rendered page works because when the cell is run in a Jupyter session the data is recorded in the notebook file. When the notebook is saved in that session and then downloaed and placed online, you can view the notebook and its output subsequently with the Jupyter Community's nbviewer and the javascript-based playback controller widget added by Matplotlib is still functional when the notebook is rendered that way and you can playback the animation without even being in Jupyter . The notebook is actually stored on GitHub, yet GitHub's preview rendering won't allow playback.