Python implementation of n-body problem issue

Question

I am currently trying to implement the N-body problem using Euler's method for solving differential equations. However, the graphical outputs do not seem correct, and I'm not sure where the issue in my code is after a while of testing. I'm currently using approximate values for Alpha Centauri A and B to test. This is my code:

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
        )

And this is the resulting plot. I would expect their orbits to be less variant and more circular, sort of in a Venn diagram-esque form.

Answer 1

There are 3 steps to a correctly looking plot.

First and most importantly, get the implementation of the physical model right. r3 is supposed to contain the third powers of the distances, thus the third power of the square root has exponent 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))

This then gives the cleaned up plot

Note the differences in the scales in the axes, one would have to horizontally compress the image by a factor of 40 to get the same scale in both directions.

Second, this means that the initial velocity is too large, the stars flee from each other. These velocities might be right in the position where the stars are closest together. As a quick fix, divide the velocities by 10. This gives the plot

Better initial values could be obtained by evaluating and transforming the supposedly more realistic data from 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

Third, the mass center is not at the origin and has a non-zero velocity. Normalize the initial values for that

    # 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

(or apply suitable broadcasting magic instead of the last two lines) to get the plot that you were probably expecting.

That the orbits spiral outwards is typical for the Euler method, as the individual steps move along the tangents to the convex ellipses of the exact solution.

The same only using RK4 with 5-day time steps gives prefect looking ellipses

For the RK4 implementation the most important step is to package the non-trivial derivatives computation into a separate sub-procedure

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)

Then one can take the Euler step out of the time loop

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

and similarly implement the RK4 step

    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 the method like

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

and then the time loop takes the generic form

    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