Python n体问题的实现问题

Python implementation of n-body problem issue

我目前正在尝试使用欧拉求解微分方程的方法来实现 N 体问题。但是,图形输出似乎不正确,经过一段时间的测试后,我不确定我的代码中的问题出在哪里。我目前正在使用半人马座阿尔法星 A 和 B 的近似值进行测试。这是我的代码:

import numpy as np
import matplotlib.pyplot as plt
from math import floor

# gravitation constant
G = 6.67430e-11
# astronomical units
au = 1.496e11
sec_in_day = 60 * 60 * 24
dt = 1 * sec_in_day

class Body(object):
    def __init__(self, name, init_pos, init_vel, mass):
        self.name = name
        self.p = init_pos
        self.v = init_vel
        self.m = mass

def run_sim(bodies, t):
    mass = np.array([[b.m] for b in bodies], dtype=float) # (n, 1, 1)
    vel = np.array([b.v for b in bodies], dtype=float) # (n, 1, 3)
    pos = np.array([b.p for b in bodies], dtype=float) # (n, 1, 3)
    names = np.array([b.name for b in bodies], dtype=str)

    # save positions and velocities for plotting
    plt_pos = np.empty((floor(t/dt), pos.shape[0], pos.shape[1]))
    plt_vel = np.empty((floor(t/dt), pos.shape[0], pos.shape[1]))

    # center of mass
    com_p = np.sum(np.multiply(mass, pos),axis=0) / np.sum(mass,axis=0)

    curr = 0
    i = 0
    while curr < t:
        dr = np.nan_to_num(pos[None,:] - pos[:,None]) 
        r3 = np.nan_to_num(np.sum(np.abs(dr)**2, axis=-1)**(0.5)).reshape((pos.shape[0],pos.shape[0],1))
        a = G * np.sum((np.nan_to_num(np.divide(dr, r3)) * np.tile(mass,(pos.shape[0],1)).reshape(pos.shape[0],pos.shape[0],1)), axis=1)

        pos += vel * dt
        plt_pos[i] = pos
        vel += a * dt
        plt_vel[i] = vel
        curr += dt
        i += 1

    fig = plt.figure(figsize=(15,15))
    ax = fig.add_subplot()
    for i in list(range(plt_pos.shape[1])):
        ax.plot(plt_pos[:,i,0], plt_pos[:,i,1], alpha=0.5, label=names[i])
        ax.scatter(plt_pos[-1,i,0], plt_pos[-1,i,1], marker="o", label=f'{i}')
    plt.legend()
    plt.show()

run_sim(bodies = [ Body('Alpha Centauri A', [0, 0, 0], [0,22345,0], 1.989e30*1.1),
                  Body('Alpha Centauri B', [23 * au, 0, 0], [0,-18100,0], 1.989e30*0.907),
                 ],
        t = 100 * 365 * sec_in_day
        )

这是结果图。我希望它们的轨道变化更少,更圆,有点像维恩图式的形式。

绘制正确的绘图需要 3 个步骤。

首先也是最重要的一点,正确实施物理模型。 r3 应该包含距离的三次方,因此平方根的三次方有指数 1.5

        r3 = np.nan_to_num(np.sum(np.abs(dr)**2, axis=-1)**(1.5)).reshape((pos.shape[0],pos.shape[0],1))

然后给出清理后的图

注意比例的差异,必须水平压缩图像才能在两个方向上获得相同的比例。

其次,这意味着初速度过大,恒星相互逃逸。这些速度可能恰好位于恒星距离最近的位置。作为快速修复,将速度除以 10。这给出了图

通过评估和转换来自 https://towardsdatascience.com/modelling-the-three-body-problem-in-classical-mechanics-using-python-9dc270ad7767 or use the Kepler laws with the more general data from http://www.solstation.com/orbits/ac-absys.htm

的假定更真实的数据,可以获得更好的初始值

第三,质心不在原点,速度不为零。标准化

的初始值
    # center of mass
    com_p = np.sum(np.multiply(mass, pos),axis=0) / np.sum(mass,axis=0)
    com_v = np.sum(np.multiply(mass, vel),axis=0) / np.sum(mass,axis=0)
    for p in pos: p -= com_p
    for v in vel: v -= com_v

(或应用合适的广播魔法而不是最后两行)以获得您可能期待的情节。

欧拉方法的典型轨道是向外螺旋的,因为各个步骤沿着切线移动到精确解的凸椭圆。

仅使用具有 5 天时间步长的 RK4 得到完美的椭圆

对于 RK4 实现,最重要的一步是将非平凡的导数计算打包到一个单独的子过程中

def run_sim(bodies, t, dt, method = "RK4"):
    ...
    
    def acc(pos):
        dr = np.nan_to_num(pos[None,:] - pos[:,None]) 
        r3 = np.nan_to_num(np.sum(np.abs(dr)**2, axis=-1)**(1.5)).reshape((pos.shape[0],pos.shape[0],1))
        return G * np.sum((np.nan_to_num(np.divide(dr, r3)) * np.tile(mass,(pos.shape[0],1)).reshape(pos.shape[0],pos.shape[0],1)), axis=1)

然后可以将欧拉步从时间循环中取出来

    def Euler_step(pos, vel, dt):
        a = acc(pos);
        return pos+vel*dt, vel+a*dt

并类似地执行 RK4 步骤

    def RK4_step(pos, vel, dt):
        v1 = vel
        a1 = acc(pos) 
        v2 = vel + a1*0.5*dt
        a2 = acc(pos+v1*0.5*dt) 
        v3 = vel + a2*0.5*dt
        a3 = acc(pos+v2*0.5*dt)
        v4 = vel + a3*dt
        a4 = acc(pos+v3*dt) 
        return pos+(v1+2*v2+2*v3+v4)/6*dt, vel+(a1+2*a2+2*a3+a4)/6*dt

Select方法类似

    stepper = RK4_step if method == "RK4" else Euler_step

然后时间循环采用通用形式

    N = floor(t/dt)
    ...
    for i in range(1,N+1):
        pos, vel = stepper(pos, vel, dt)
        plt_pos[i] = pos
        plt_vel[i] = vel