Search code examples
pythonphysics

Why does my nbody solver have the earth moving in a straight line instead of an orbit?


I'm making an Nbody solver for a physics project, and I've gotten to the point where I'm plotting the position of the Earth around the Sun. When I do this, however, the earth just travels away from the sun in a straight line, and does not perform an orbit. Is there something I'm missing with my acceleration calculation thats causing this? Thank you in advance.

import numpy as np 
import matplotlib.pyplot as plt 
import astropy.units as u 
import astropy.constants as c 
import sys 
import time
from mpl_toolkits.mplot3d import Axes3D

#making a class for Celestial Objects
class CelestialObjects():
    def __init__(self,mass,pos_vec,vel_vec,name=None, has_units=True):
        self.name=name
        self.has_units=has_units
        if self.has_units:
            self.mass=mass.cgs
            self.pos=pos_vec.cgs.value
            self.vel=vel_vec.cgs.value
        else:
            self.mass=mass 
            #3d array for position of body in 3d space in AU
            self.pos=pos_vec 
            #3d array for velocity of body in 3d space in km/s
            self.vel=vel_vec
        
    def return_vec(self):
        return np.concatenate((self.pos,self.vel))
    def return_name(self):
        return self.name
    def return_mass(self):
        if self.has_units:
            return self.mass.cgs.value
        else:
            return self.mass
  

#set up first instance of a celestial object, Earth
Earth=CelestialObjects(name='Earth',
                       pos_vec=np.array([0,1,0])*u.AU,
                       vel_vec=np.array([0,30,0])*u.km/u.s,
                       mass=1.0*c.M_earth)
#set up second instance of a celestial object, the Sun
Sun=CelestialObjects(name='Sun',
                     pos_vec=np.array([0,0,0])*u.AU,
                     vel_vec=np.array([0,0,0])*u.km/u.s,
                     mass=1*u.Msun)
bodies=[Earth,Sun]

#making a class for system
class Simulation():
    def __init__(self,bodies,has_units=True):
        self.has_units=has_units
        self.bodies=bodies
        self.Nbodies=len(self.bodies)
        self.Ndim=6
        self.quant_vec=np.concatenate(np.array([i.return_vec() for i in self.bodies]))
        self.mass_vec=np.array([i.return_mass() for i in self.bodies])
        self.name_vec=[i.return_name() for i in self.bodies]
        
    def set_diff_eqs(self,calc_diff_eqs,**kwargs):
        self.diff_eqs_kwargs=kwargs
        self.calc_diff_eqs=calc_diff_eqs
    
    def rk4(self,t,dt):
        k1= dt* self.calc_diff_eqs(t,self.quant_vec,self.mass_vec,**self.diff_eqs_kwargs)
        k2=dt*self.calc_diff_eqs(t+dt*0.5,self.quant_vec+0.5*k1,self.mass_vec,**self.diff_eqs_kwargs)
        k3=dt*self.calc_diff_eqs(t+dt*0.5,self.quant_vec+0.5*k2,self.mass_vec,**self.diff_eqs_kwargs)
        k4=dt*self.calc_diff_eqs(t+dt,self.quant_vec+k3,self.mass_vec,**self.diff_eqs_kwargs)
        
        y_new=self.quant_vec+((k1+2*k2+2*k3+k4)/6)
        return y_new
    
    def run(self,T,dt,t0=0):
        if not hasattr(self,'calc_diff_eqs'):
            raise AttributeError('You must set a diff eq solver first.')
        if self.has_units:
            try:
                _=t0.unit
            except:
                t0=(t0*T.unit).cgs.value
            T=T.cgs.value
            dt=dt.cgs.value
        
        self.history=[self.quant_vec]
        clock_time=t0
        nsteps=int((T-t0)/dt)
        start_time=time.time()
        for step in range(nsteps):
            sys.stdout.flush()
            sys.stdout.write('Integrating: step = {}/{}| Simulation Time = {}'.format(step,nsteps,round(clock_time,3))+'\r')
            y_new=self.rk4(0,dt)
            self.history.append(y_new)
            self.quant_vec=y_new
            clock_time+=dt
        runtime=time.time()-start_time
        print('\n')
        print('Simulation completed in {} seconds'.format(runtime))
        self.history=np.array(self.history)
    
def nbody_solver(t,y,masses):
    N_bodies=int(len(y)/6)
    solved_vector=np.zeros(y.size)
    distance=[]
    for i in range(N_bodies):
        ioffset=i * 6
        for j in range(N_bodies):
            joffset=j * 6
            solved_vector[ioffset]=y[ioffset+3]
            solved_vector[ioffset+1]=y[ioffset+4]
            solved_vector[ioffset+2]=y[ioffset+5]
            if i != j:
                dx= y[ioffset]-y[joffset]
                dy=y[ioffset+1]-y[joffset+1]
                dz=y[ioffset+2]-y[joffset+2]
                r=(dx**2+dy**2+dz**2)**0.5
                ax=(-c.G.cgs*masses[j]/r**3)*dx
                ay=(-c.G.cgs*masses[j]/r**3)*dy
                az=(-c.G.cgs*masses[j]/r**3)*dz
                ax=ax.value
                ay=ay.value
                az=az.value
                solved_vector[ioffset+3]+=ax
                solved_vector[ioffset+4]+=ay
                solved_vector[ioffset+5]+=az
    return solved_vector

simulation=Simulation(bodies)
simulation.set_diff_eqs(nbody_solver)
simulation.run(365*u.day,1*u.hr)

earth_position = simulation.history[:, :3]  # Extracting position for Earth
sun_position = simulation.history[:, 6:9]    # Extracting position for Sun
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Plot the trajectories
ax.plot(earth_position[:, 0], earth_position[:, 1], earth_position[:, 2], label='Earth')
ax.plot(sun_position[:, 0], sun_position[:, 1], sun_position[:, 2], label='Sun')

# Add labels and title
ax.set_xlabel('X (AU)')
ax.set_ylabel('Y (AU)')
ax.set_zlabel('Z (AU)')
ax.set_title('Trajectories of Earth and Sun')

# Add a legend
ax.legend()

# Show the plot
plt.show()

Solution

  • Looks to me like its working: You should probably just give your "Earth" an starting velocity which is perpinducular to your suns gravity. Else the earth will just fly straight into the sun! (swap vel_vec=np.array([0, 30, 0]) with vel_vec=np.array([30, 0, 0]) for example)

    # set up first instance of a celestial object, Earth
    Earth = CelestialObjects(name='Earth',
                             pos_vec=np.array([0, 1, 0]) * u.AU,
                             vel_vec=np.array([30, 0, 0]) * u.km / u.s,
                             mass=1.0 * c.M_earth)
    

    enter image description here