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