Jupyter Notebook - Can't update matplotlib plot in a for loop (keeping previous plots), nor animate the plot using animation subfunction

Question

For a project, I need to plot the orbit of a particle around another in a Jupyter Notebook.

I succesfully managed to do this in a for loop, printing the result only when the for loop is finished. Here is the code with comments for context:

# Import necessary Python libraries
import matplotlib.pyplot as plt
import numpy as np
%matplotlib notebook
from matplotlib import animation

# Define constants and initialize variables:
# All quantities are directly in SI units
# (To begin, I will take the Earth-Moon system as a reference)

G = 6.67e-11
R_12 = 384399 * 1000 + 3474 / 2 * 1000 + 6371 * 1000

# Attracting particle 1 (P1), kept fixed.
m1 = 5.97e24
x1 = 0; y1 = 0

# Attracted particle 2 (P2), moving around the fixed particle.
m2 = 7.35e22
x2_0 = x1 + R_12; y2_0 = y1 + 0
Vx2_0 = 0; Vy2_0 = 1000

# Simulation parameters
dt = 3600  # Time step
N = 24 * 30  # Number of iterations

# Define functions necessary for the simulation:
# The chosen position vector structure is Xi = np.array([xi, yi]).

def calculate_A2(X2_t):  # Calculate the acceleration of P2
    F12_t = -G * m1 * m2 * (X2_t - X1) / (
        ((X2_t[0] - X1[0]) ** 2 + (X2_t[1] - X1[1]) ** 2) ** (3/2))
    A2_t = F12_t / m2
    return A2_t


def verlet(X2_t, V2_t):
    A2_t = calculate_A2(X2_t)
    X2_t_dt = X2_t + dt * V2_t + (dt ** 2) / 2 * A2_t
    A2_t_dt = calculate_A2(X2_t_dt)
    V2_t_dt = V2_t + dt / 2 * (A2_t + A2_t_dt)
    return X2_t_dt, V2_t_dt

# Simulation:
X1 = np.array([x1, y1])
X2 = np.array([x2_0, y2_0])
V2 = np.array([Vx2_0, Vy2_0])
# Here, we will not keep the position of P2 after plotting its position.
# This choice is due to the desire not to "waste" memory unnecessarily.

fig, ax = plt.subplots()
ax.plot(X1[0], X1[1], 'o', markersize=10, color='blue', label='P1')  # P1

ax.plot(X2[0], X2[1], 'o', markersize=10, color='red', label='P2')  # P2, initial position

for i in range(N):
    X2, V2 = verlet(X2, V2)
    ax.plot(X2[0], X2[1], 'o', markersize=10, color='red')

ax.legend()
plt.axis('scaled')
plt.show()

The problem is that I would like matplotlib to plot while the loop is still going, updating each time with the new position of the P2 particle (the Moon here). But I just can't make it work, although I searched here, read and watched some tutorials, and even used an example from my teacher (using the animation sublibrary of matplotlib).

First, I wanted to keep the for loop, as I didn't see the point of changing it to an animation function. After reading some tutorials, I figured set_data would suit my needs in the for loop. Thus, I tried modifying my code to make it work, based on this example from GeeksForGeeks:

import time
fig, ax = plt.subplots()
ax.plot(X1[0],X1[1],'o',markersize=10,color='blue', label = 'P1') # P1

posP2, = ax.plot(X2[0],X2[1],'o',markersize=10,color='red', label = 'P2') # P2, position initiale

for i in range(N):
    X2, V2 = verlet(X2, V2)
    x2 = X2[0]
    y2 = X2[1]
    posP2.set_xdata(x2)
    posP2.set_ydata(y2)
    fig.canvas.draw()
    fig.canvas.flush_events()
    time.sleep(0.1)

This doesn't work, even though it does run. I had another version that managed to update and plot a single P2 dot each time (no P1 plot), but P2 would get out of the frame after a few seconds so I tried changing the xy limits and the scale (I would like P1 to be the center of the plot at all time, and the scale to be the same for both axis), and it stopped working after that. I unfortunately lost this version in my numerous tries...

I also tried using animation from matplotlib, using a notebook from my teacher as a reference (aswell as online examples), but it didn't work either. Here is one of my drafts :

%matplotlib notebook from matplotlib import animation

# Simulation :

X1 = np.array([x1,y1])
X2 = np.array([x2_0,y2_0])
V2 = np.array([Vx2_0,Vy2_0])

plt.ion()
fig, ax = plt.subplots()

# Plot de P1 initial
ax.plot(x1,y1,'o',markersize=10,color='blue', label = 'P1') # P1

ax.set_xlim(X1[0]-1.1*R_12, X1[0]+1.1*R_12)  # Limites de l'axe x
ax.set_ylim(X1[1]-1.1*R_12, X1[1]+1.1*R_12)  # Limites de l'axe y
ax.set_aspect('equal')  # Même échelle en x et y

posP2, = ax.plot([],[],'o',markersize=10,color='red', label = 'P2') # P2, initial position

ax.legend() # On affiche la légende

hist_P2 = np.zeros((2,N+1))

def animMouvement(i):
    x_posP2 = [X2[0],X2[0]]
    y_posP2 = [X2[1],X2[1]]
    hist_P2[0,i] = X2[0]
    hist_P2[1,i] = X2[1]
    X2, V2 = verlet(X2,V2)
    posP2.set_data(x_posP2,y_posP2)
    print(f"X2_{i} : {X2}")
    return posP2

anim = animation.FuncAnimation(fig, animMouvement, frames=N+1, interval=20, blit=True)

plt.show()  # Call plt.show() at the end

A figure is created and updated (hovering my mouse over it shows each load), but it is still not working as expected.

I don't really know what to do, there is surely something I am missing but I don't know where to look or what could be the issue.

If you know what could cause those issues, preferably keeping the for loop, let me know! Thank you very much for your time!

Answer 1

You should do it this way when using matplotlib:

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

%matplotlib notebook

G = 6.67e-11
R_12 = 384399 * 1000 + 3474 / 2 * 1000 + 6371 * 1000
m1 = 5.97e24
m2 = 7.35e22
x1, y1 = 0, 0
x2_0, y2_0 = x1 + R_12, y1
Vx2_0, Vy2_0 = 0, 1000
dt = 3600
N = 24 * 30

def calculate_A2(X2_t, X1):
    F12_t = -G * m1 * m2 * (X2_t - X1) / (np.sum((X2_t - X1)**2) ** 1.5)
    A2_t = F12_t / m2
    return A2_t

def verlet(X2_t, V2_t, X1):
    A2_t = calculate_A2(X2_t, X1)
    X2_t_dt = X2_t + dt * V2_t + 0.5 * dt**2 * A2_t
    A2_t_dt = calculate_A2(X2_t_dt, X1)
    V2_t_dt = V2_t + 0.5 * dt * (A2_t + A2_t_dt)
    return X2_t_dt, V2_t_dt

X1 = np.array([x1, y1])
X2 = np.array([x2_0, y2_0])
V2 = np.array([Vx2_0, Vy2_0])

fig, ax = plt.subplots()
ax.set_xlim(X1[0] - 1.1 * R_12, X1[0] + 1.1 * R_12)
ax.set_ylim(X1[1] - 1.1 * R_12, X1[1] + 1.1 * R_12)
ax.set_aspect('equal')
ax.plot(X1[0], X1[1], 'o', markersize=10, color='blue', label='P1')
posP2, = ax.plot([], [], 'o', markersize=10, color='red', label='P2')
ax.legend()

def animMouvement(i):
    global X2, V2
    X2, V2 = verlet(X2, V2, X1)
    posP2.set_data(X2[0], X2[1])
    return posP2,

anim = animation.FuncAnimation(fig, animMouvement, frames=N, interval=20, blit=True)
plt.show()

I used google.colab so I modified a little (%matplotlib is not supported:

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import animation
from IPython.display import HTML

G = 6.67e-11
R_12 = 384399 * 1000 + 3474 / 2 * 1000 + 6371 * 1000
m1 = 5.97e24
m2 = 7.35e22
x1, y1 = 0, 0
x2_0, y2_0 = x1 + R_12, y1
Vx2_0, Vy2_0 = 0, 1000
dt = 3600
N = 24 * 30

def calculate_A2(X2_t, X1):
    F12_t = -G * m1 * m2 * (X2_t - X1) / (np.sum((X2_t - X1)**2) ** 1.5)
    A2_t = F12_t / m2
    return A2_t

def verlet(X2_t, V2_t, X1):
    A2_t = calculate_A2(X2_t, X1)
    X2_t_dt = X2_t + dt * V2_t + 0.5 * dt**2 * A2_t
    A2_t_dt = calculate_A2(X2_t_dt, X1)
    V2_t_dt = V2_t + 0.5 * dt * (A2_t + A2_t_dt)
    return X2_t_dt, V2_t_dt

X1 = np.array([x1, y1])
X2 = np.array([x2_0, y2_0])
V2 = np.array([Vx2_0, Vy2_0])

fig, ax = plt.subplots()
ax.set_xlim(X1[0] - 1.1 * R_12, X1[0] + 1.1 * R_12)
ax.set_ylim(X1[1] - 1.1 * R_12, X1[1] + 1.1 * R_12)
ax.set_aspect('equal')
ax.plot(X1[0], X1[1], 'o', markersize=10, color='blue', label='P1')
posP2, = ax.plot([], [], 'o', markersize=10, color='red', label='P2')
ax.legend()

def animMouvement(i):
    global X2, V2
    X2, V2 = verlet(X2, V2, X1)
    posP2.set_data(X2[0], X2[1])
    return posP2,

anim = animation.FuncAnimation(fig, animMouvement, frames=N, interval=20, blit=True)

HTML(anim.to_html5_video())

which gives