1

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!

6
  • Use %matplotlib notebook instead of inline Commented Apr 20, 2024 at 15:15
  • @TinoD Woops I didn't include that in my post, sorry, however I did try using %matplotlib notebook each time, and sadly it didn't work either, the issue must be something else :/ Commented Apr 20, 2024 at 16:27
  • There is an important distinction when making matplotlib animations these days. You need to keep track of the specific Jupyter tech you are using. And you cannot necessarily assume your teacher is using the same as you? (Or can you? Did they set you up on Jupyter?) A lot of it boils down to if you are using current JupyterLab or Jupyter Notebook 7+, then you need to install ipympl and use %matplotlib ipympl to signal the backend. See here and ... Commented Apr 20, 2024 at 17:26
  • <continued> my repo about modern Jupyter animations. And then if you are using the old Jupyter versions (NbClassic) and only then can you use %matplotlib notebook these days to do interactive plotting, which animations happen to be, although you may not know it because the ipympl documentation curiously doesn't seem to contain an animation example at this time. I only figured it out when I kept getting the same error I see in cases where ipympl isn't installed. Commented Apr 20, 2024 at 17:29
  • For the record, VPython is another nice library for this type of animation. They even have an example animation with 3D orbits here. I linked to the one that runs in your browser; however, VPython can be installed and run in Jupyter, too, see here. Commented Apr 20, 2024 at 17:33

1 Answer 1

0

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

enter image description here

Sign up to request clarification or add additional context in comments.

5 Comments

Thank you very much! I still have an issue, but it may not be the code's fault! In fact, one of my draft was very similar to this but didn't work (the animation would play, and only the second frame would show before stopping). Thanks to you, it seems the issue may be with my Jupyter/Python setup, as I encounter almost the same issue with your code, although it works on your side. With your code, the animation does play, but it only shows some frames, and stops after a few frames. [...]
[...] There's more, I tried playing it again, and it began showing only the second frame, like my draft did (maybe my code was good at one point but my pc tricked me into thinking it wasn't...). With my teacher's example, I had the same kind of issue, although it was smoother and more reliable: it stopped working after the first run, and I had to restart the whole kernel to make it work again. I'll try on other instances real quick.
Alright, it works well on my school's VMs, it's fine. I'll take a closer look on my setup, maybe I did something wrong. Anyway, thank you for your help! I'll also see if I can do something interesting with VPython, as @Wayne suggested in a comment. Thank you!
Glad it helped! Happy coding and good luck with your studies!
Thank you! Also, thanks for the heads-up on arrays sum, yours is faster I think, I'll change that too! And your Colab code works on DeepNote too, I'll use it for my shared notebook :)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.