Making a trail animation for the N body problem-CodePudding

I solved the N body problem with Python. My code works perfectly but I have to write manually line animation for each body. I would like to make a loop instead but I did not succeed. Here is the code :

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from scipy.integrate import odeint
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D



## Constantes dimensionnées du système
N = 4
G = 6.674e-11
M = np.random.uniform(low=1e30, high=3e30, size=N)

## Constantes adimensionnées du système
G = G * (6e24 * (365*24*60*60)**2) / (1.496e11)**3
M = M/6e24

## Conditions initiales
ti = 0
tf = 10
n = 500
dt = (tf-ti)/n
T = np.linspace(ti, tf, n)

R0 = np.random.uniform(low=-4, high=4, size=3*N)
V0 = np.random.uniform(low=-1, high=1, size=3*N)
Y0 = np.append(R0,V0)

def NCorps(Yk, t):

    #Création tableau contenant les positions de tous les corps
    Rk = np.zeros((3,N))
    a = 0
    b = 0
    c = 0
    for i in range(0,3*N):
        if i%3 == 0:
            Rk[0,a] = Yk[i]
            a = a 1
        if i%3 == 1:
            Rk[1,b] = Yk[i]
            b = b 1
        if i%3 == 2:
            Rk[2,c] = Yk[i]
            c = c 1

    #Création tableau contenant les vitesses et les accélérations de tous les corps
    Sk = []
    
    for i in range(3*N, 6*N):
            Sk.append(Yk[i])

    for i in range(0, N):
        axk = 0
        ayk = 0
        azk = 0
        for j in range(0, N):
            if j != i:
                axk = axk - G*M[j]*(Rk[0,i]-Rk[0,j])/((Rk[0,i]-Rk[0,j])**2 (Rk[1,i]-Rk[1,j])**2 (Rk[2,i]-Rk[2,j])**2)**(3/2)
                ayk = ayk - G*M[j]*(Rk[1,i]-Rk[1,j])/((Rk[0,i]-Rk[0,j])**2 (Rk[1,i]-Rk[1,j])**2 (Rk[2,i]-Rk[2,j])**2)**(3/2)
                azk = azk - G*M[j]*(Rk[2,i]-Rk[2,j])/((Rk[0,i]-Rk[0,j])**2 (Rk[1,i]-Rk[1,j])**2 (Rk[2,i]-Rk[2,j])**2)**(3/2)
        Sk.append(axk)
        Sk.append(ayk)
        Sk.append(azk)
    return Sk


Y = odeint(NCorps, Y0, T)


## Plot 3D image

fig = plt.figure()
ax = plt.axes(projection = '3d')

a=0
for i in range(0, 3*N):
    if i%3 == 0:
        ax.plot3D(Y[:,i], Y[:,i 1], Y[:,i 2], label = "Corps "   str(a))
        ax.scatter(Y[-1,i], Y[-1,i 1], Y[-1,i 2], marker = "o", s = 75)
        a = a 1

ax.set_xlim3d(-7, 7)
ax.set_ylim3d(-7, 7)
ax.set_zlim3d(-7, 7)
ax.set_xlabel("x (UA)")
ax.set_ylabel("y (UA)")
ax.set_zlabel("z (UA)")
plt.legend()


## Plot 3D animation

fig = plt.figure()
ax = plt.axes(projection = '3d')

ax.set_xlim3d(-7, 7)
ax.set_ylim3d(-7, 7)
ax.set_zlim3d(-7, 7)
ax.set_xlabel("x (UA)")
ax.set_ylabel("y (UA)")
ax.set_zlabel("z (UA)")

trail = [50]*N

###########################################################################################
########## HERE IS MY PROBLEM : HOW TO MAKE Animate3D FUNCTION WORK FOR N BODY ? ##########
###########################################################################################

def Animate3D(k):
    ligne1.set_data(Y[k:max(1,k-trail[0]):-1, 0], Y[k:max(1,k-trail[0]):-1, 1])
    ligne1.set_3d_properties(Y[k:max(1,k-trail[0]):-1, 2])
    
    ligne2.set_data(Y[k:max(1,k-trail[1]):-1, 3], Y[k:max(1,k-trail[1]):-1, 4])
    ligne2.set_3d_properties(Y[k:max(1,k-trail[1]):-1, 5])

    ligne3.set_data(Y[k:max(1,k-trail[2]):-1, 6], Y[k:max(1,k-trail[2]):-1, 7])
    ligne3.set_3d_properties(Y[k:max(1,k-trail[2]):-1, 8])
    
    return (ligne1, ligne2, ligne3)

ligne1, = ax.plot([], [], "o-", markevery = 10000)
ligne2, = ax.plot([], [], "o-", markevery = 10000)
ligne3, = ax.plot([], [], "o-", markevery = 10000)

anim3D = animation.FuncAnimation(fig, Animate3D, frames = n, interval = 30, blit = False)

plt.show()

As you can see in the last part, I have no idea how to generalize the Animate3D function for the N body. I wanted to create a loop like the one bellow :

j=0
lignes = []
for i in range(0, N):
    lignes.append(ligne   i)

def Animate3D(k):
    for i in range(0, 3*N):
        if i%3 == 0:
            ligne   i   .set_data(Y[k:max(1,k-trail[j]):-1, i], Y[k:max(1,k-trail[j]):-1, i 1])
            ligne   i   .set_3d_properties(Y[k:max(1,k-trail[j]):-1, i 2])
            j=j 1
    
    return lignes

for i in range(0, N):
    ligne   i   , = ax.plot([], [], "o-", markevery = 10000)

Obviously it does not work.

CodePudding user response：

You should collect lignes in a list and loop over them. In order to update counters within Animate3D, you can replace range with enumerate, as in the code below:

def Animate3D(k):
    for i, ligne in enumerate(lignes, 0):
        ligne.set_data(Y[k:max(1, k - trail[0]):-1, 3*i], Y[k:max(1, k - trail[0]):-1, 3*i   1])
        ligne.set_3d_properties(Y[k:max(1, k - trail[0]):-1, 2])

    return lignes


lignes = [ax.plot([], [], "o-", markevery = 10000)[0] for _ in range(N)]

anim3D = animation.FuncAnimation(fig, Animate3D, frames = n, interval = 30, blit = False)

plt.show()