Search code examples
pythonsimulationnumerical-methodsdifferential-equations

Issue in N-body simulations (Python implementation)


import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

G = 6.67430e-11


def getAcc(pos ,mass, n): ## returns (n,3) acceleration
    
    ## pos :  (n,3)
    ## mass : (n,1) 

    ## getting x , y , z seperately
    x = pos[:, 0:1].reshape(n,1)
    y = pos[:, 1:2].reshape(n,1)
    z = pos[:, 2:3].reshape(n,1)

    ## relative coordiantes Xij
    x_rel = x.T - x
    y_rel = y.T - y
    z_rel = z.T - z

    ## r^3
    r3 = x_rel**2 + y_rel**2 + z_rel**2

    ## r^(-1.5) keeping in mind rii is 0
    r3[r3!=0] = r3[r3!=0]**(-1.5)
    

    ax = G * (r3 * x_rel) @ mass
    ay = G * (r3 * y_rel) @ mass
    az = G * (r3 * z_rel) @ mass
    return np.hstack((ax , ay , az))


def solve(mass , pos , vel, n):

    ## pos : (n,3) 
    ## vel : (n,3) 
    ## mass : (n,1)

    # Parameters
    t = 0
    t_final = 200
    dt = 0.01
    total_steps = np.int64(np.ceil(t_final/dt))

    # vel -= np.mean(vel*mass , 0) / np.mean(mass)

    # Initial acceleration
    acc = getAcc(pos , mass , n)

    # Position matrix to save all positions (Number of bodies , 3 , total_time_frames)
    pos_m = np.zeros((n,3,total_steps+1) , dtype=np.float64)
    pos_m[:, :, 0] = pos

    # Calculate positions

    for i in range(total_steps):

        pos += vel*dt + acc*dt*dt/2
        vel += 0.5 * acc * dt
        acc = getAcc(pos , mass , n)
        vel += 0.5 * acc * dt
        
        t+=dt

        pos_m[: , : , i+1] = pos

    
    return pos_m


def showplot(pos_m,n):
    
    # Create a 3D plot
    t = pos_m.shape[2]  # number of time steps

    # Create a 3D plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    

    # Plot the positions of each body over time with different colors
    colors = ['r', 'g', 'b']  # one color for each body
    for i in range(n):
        x = pos_m[i, 0, :]
        y = pos_m[i, 1, :]
        z = pos_m[i, 2, :]
        ax.plot(x, y, z, color=colors[i], label='Body {}'.format(i+1))

    # Add labels and legend
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.legend()

    # Show the plot
    plt.show()



def main():

    # Masses of the two bodies
    mass = np.array([5.97e24, 7.35e22], dtype=np.float64).reshape(2, 1)

    # Initial positions of the two bodies
    pos = np.array([[0, 0, 0], [384400000, 0, 0]], dtype=np.float64)

    # Initial velocities of the two bodies
    vel = np.array([[0, 0, 0], [0, 1022, 0]], dtype=np.float64)

    # Number of bodies
    n = 2

    # Solve for the positions of the two bodies over time
    pos_m = solve(mass, pos, vel, n)

    # Plot the positions of the two bodies over time
    showplot(pos_m,n)
    

main()

I tried to use parameters for an elliptical path, but the output looked something like this: 3D plot

It looks like it is only moving in a single direction of initial velocity, and has no acceleration in other directions.

I even tried earth-sun system, still had the same issue.

Is the numerical integration too simple, leading up to a lot of error in solving the ODE?
Or is there some other mathematical issue?


Solution

  • The parameters (mass and initial positions and velocities) look like the Earth-Moon system, which has an orbital period of about 27.3 days.

    However, your t = 200 means you're simulating for a period of 200 seconds. That's only 0.0085% of an orbital period, so the trajectory you plot is a very small arc, and its curvature is almost imperceptible. For all intents and purposes it looks like a straight line.

    If you increase your t to a longer period, the curvature becomes noticeable (for example t = 86400 * 7 and dt = 60 simulates a week of motion with a timestep of 1 minute, and as expected the trajectory looks like a quarter of a circle).

    While the numerical integration could be improved, your method gives results that don't look crazy when you model for a few days or months.