N corps (animation 3D)

Bonsoir,
Je précise avant toute chose que je suis débutant en python (et nouveau sur le forum ... :) )
Je travail depuis longtemps sur un programme qui calcule la position de n corps soumis à leur attractions gravitationnelles mutuelles (le fameux problème à n corps)
J'ai beaucoup de versions de ce programme (ça se compte en dizaines, j'en fais quand j'ai le temps depuis un an)
Il y en a en 2D, en 2D animé(avec VPython) en 3D, avec 2 corps, avec 3, avec N...
Dans l'ensemble ils marchent assez bien
là par exemple : https://www.youtube.com/watch?v=LpkZ2NJhKa0

Mais je n'ai jamais réussi en 3D animé. J'aimerais surtout utiliser matplotlib pour cela !
Si vous pouviez m'apporter votre aide, je vous en serez reconnaissant car ça fait bien longtemps que je bloque là dessus, et je ne vois pas du tout comment faire.
Voici le code actuel :

Code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from math import*
import numpy as np
import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.pyplot as plt
import random 
from matplotlib import animation
 
#définition des paramètres
G=1
tf=120
dt=0.01
n=10
m=1
L=100
 
###Nbody sim using Eulerian method###
 
def nuage():
    #creation d un nuage de points random
    position_initiale=np.zeros((n,3))
    vitesse_initiale=np.zeros((n,3))
    for i in range(n):
        position_initiale[i,0]=random.random()*L
        position_initiale[i,1]=random.random()*L
        position_initiale[i,2]=random.random()*L
        vitesse_initiale[i]=[0,0,0]
    return position_initiale,vitesse_initiale
 
###recuperation des donnees
A=nuage()
pos=A[0]
vitesse=A[1]
 
#preparation du graph
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1], projection='3d')
pts = ax.plot([], [], [], 'o', c='k', ms=4)
 
 
def init():
    pts[0].set_data(pos[:,0], pos[:,1])
    pts[0].set_3d_properties(pos[:,2])
    return pts
 
def distance(position): #renvoi une matrice symétrique correspondant à toutes les distances à l'instant fixé car l'application distance 
    D=np.zeros([n, n])  #est symétrique(d(k,i)=d(i,k)), ce qui économise la moitié des calculs de distance
    for i in range(n):
        for k in range(n):
            if k==i:
                D[i,k]=0
 
            else:
                dx = position[k,0]-position[i,0]
                dy = position[k,1]-position[i,1]
                dz = position[k,2]-position[i,2]
                d2 = dx*dx + dy*dy +dz*dz
                D[k,i]=D[i,k]=d2*sqrt(d2)
    return D
 
 
 
def nbody(i):
    global pos
    #Calcul des distances
    D=distance(pos)
    Lx=pos[:,0]     #listes temporaires
    Ly=pos[:,1]     #pour calculer les 
    Lz=pos[:,2]     #accélérations à 
    vx=vitesse[:,0] #partir du même instant
    vy=vitesse[:,1]
    vz=vitesse[:,2]
    for i in range(n):
        Ax=Ay=Az=0
        for k in range(n):
            if D[i,k]>0: #elimine le cas k==i qu'on ne veut pas !
                Ax+=G*m*(Lx[k]-Lx[i])/D[i,k]
                Ay+=G*m*(Ly[k]-Ly[i])/D[i,k]
                Az+=G*m*(Lz[k]-Lz[i])/D[i,k]
 
        #Mouvement de l'objet i
        vitesse[i,0]=vx[i]+Ax*dt
        vitesse[i,1]=vy[i]+Ay*dt
        vitesse[i,2]=vz[i]+Az*dt
 
 
        pos[i,0]=Lx[i]+vx[i]*dt
        pos[i,1]=Ly[i]+vy[i]*dt
        pos[i,2]=Lz[i]+vz[i]*dt
 
    pts[0].set_data(pos[:,0], pos[:,1])
    pts[0].set_3d_properties(pos[:,2])
    fig.canvas.draw()
 
 
'''
def animate():
    pos=nbody()
    l.set_ydata(pos[:,1])
    l.set_xdata(pos[:,0])
    l.set_zdata(pos[:,2])
    fig.canvas.draw()
    return pts
'''    
 
 
 
anim = animation.FuncAnimation(fig, nbody, init_func=init, interval=30, blit=True)
 
plt.show()