为 N body 问题制作拖尾动画

Question

我用 Python 解决了 N body 个问题。我的代码运行完美，但我必须为每个 body 手动编写线条动画。我想做一个循环，但我没有成功。这是代码：

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()

正如您在最后一部分中所看到的，我不知道如何为 N body 泛化 Animate3D 函数。我想创建一个如下所示的循环：

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)

显然不行。

Answer 1

您应该将 lignes 收集到一个列表中并循环遍历它们。为了在 Animate3D 内更新计数器，您可以将 range 替换为 enumerate，如下面的代码所示：

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()

为 N body 问题制作拖尾动画

Making a trail animation for the N body problem

python

animation

numpy

data-visualization

matplotlib