I made an n-body gravitational simulation in python. The algorithm does produce an approximate solution, which is shown at the bottom of the post. Additional methods to produce animations (among other things) can be found here. I would like to see if efficiency, readability, and/or accuracy can be improved.
The main idea of the algorithm is this: given two bodies of masses m_i
and m_j
separated by a distance r_ij
apart (for i ≠ j
), one can calculate the acceleration due to gravity of m_i
. This is repeated for all bodies at each iteration of time to get the net acceleration a_i
of m_i
. The velocity is obtained by integrating the acceleration (since dv = a dt); similarly, displacement is given by integrating velocity. The process repeats for each iteration of time.
For convenience, the class Body
contains array attributes such as position, velocity, and acceleration, which are updated at each iteration of time.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
class Body():
"""
This class contains adjustable parameters as attributes. These
parameters include current and previous positions, velocities, and
accelerations.
"""
def __init__(self, identifier, facecolor, position, mass=1, velocity=None, acceleration=None, notes=None):
"""
identifier:
type <str>
facecolor:
type <str>
position:
type <tuple / list / array>
mass:
type <int / float>
velocity:
type <tuple / list / array> or None
acceleration:
type <tuple / list / array> or None
notes:
type <str> or None
"""
self.identifier = identifier
self.facecolor = facecolor
self.position = np.array(position, dtype=float)
self.mass = mass
self.velocity = self.autocorrect_parameter(velocity)
self.acceleration = self.autocorrect_parameter(acceleration)
self._vector_position = [self.position]
self._vector_velocity = [self.velocity]
self._vector_acceleration = [self.acceleration]
self.notes = notes
def __str__(self):
res = ''
for key in ('identifier', 'vector_position'):
value = getattr(self, key)
text = '\n .. {}:\n\t{}\n'.format(key, value)
res += '{}'.format(text)
return res
def autocorrect_parameter(self, args):
"""
args:
type <tuple / list / array> or None
"""
if args is None:
return np.zeros(self.position.shape, dtype=float)
else:
return np.array(args, dtype=float)
@property
def vector_position(self):
return np.array(self._vector_position)
@property
def vector_velocity(self):
return np.array(self._vector_velocity)
@property
def vector_acceleration(self):
return np.array(self._vector_acceleration)
@property
def scalar_position(self):
return np.sqrt(np.sum(np.square(self.vector_position), axis=1))
@property
def scalar_velocity(self):
return np.sqrt(np.sum(np.square(self.vector_velocity), axis=1))
@property
def scalar_acceleration(self):
return np.sqrt(np.sum(np.square(self.vector_acceleration), axis=1))
def get_vector(self, parameter):
"""
parameter:
type <str>
"""
attribute = 'vector_{}'.format(parameter)
return getattr(self, attribute)
def get_scalar(self, parameter):
"""
parameter:
type <str>
"""
attribute = 'scalar_{}'.format(parameter)
return getattr(self, attribute)
Before creating instances of the Body
class, it's important for units to be scaled properly. The class UnitConversions
contains convenience methods for this purpose.
class UnitConversions():
"""
This class contains methods to convert measurements from one unit to
another. Measurable quantities include distance, time, velocity, mass,
and acceleration. All values are scaled to their corresponding units
such that each base quantity measured in SI units has a value of 1.
"""
def __init__(self):
self.distance_units = ('meter', 'kilometer', 'mile', 'Astronomical Unit', 'Light Year', 'parsec')
self.distance_values = (1, 1000, 1609, 1.496e11, 9.461e15, 3.086e16)
self.time_units = ('second', 'hour', 'day', 'year')
self.time_values = (1, 3600, 3600*24, 3600*24*365)
self.mass_units = ('gram', 'kilogram', 'lunar mass', 'earth mass', 'solar mass')
self.mass_values = (1/1000, 1, 7.3459e22, 5.9722e24, 1.988e30)
self.distance_conversion_factors = self.get_conversion_factors(self.distance_units, self.distance_values, base_unit='meter')
self.time_conversion_factors = self.get_conversion_factors(self.time_units, self.time_values, base_unit='second')
self.mass_conversion_factors = self.get_conversion_factors(self.mass_units, self.mass_values, base_unit='kilogram')
@staticmethod
def get_conversion_factors(units, values, base_unit):
"""
units:
type <tuple / list / array>
values:
type <tuple / list / array>
base_unit:
type <str>
"""
res = {}
for outer_key, outer_value in zip(units, values):
if outer_key == base_unit:
res[outer_key] = dict(zip(units, values))
else:
res[outer_key] = {}
for inner_key, inner_value in zip(units, values):
res[outer_key][inner_key] = inner_value / outer_value
return res
def convert_distance(self, distance, original_unit, prime_unit='meter'):
"""
distance:
type <int / float / array>
original_unit:
type <str>
prime_unit:
type <str>
"""
conversion_factor = self.distance_conversion_factors[prime_unit][original_unit]
return distance * conversion_factor
def convert_time(self, time, original_unit, prime_unit='second'):
"""
time:
type <int / float / array>
original_unit:
type <str>
prime_unit:
type <str>
"""
conversion_factor = self.time_conversion_factors[prime_unit][original_unit]
return time * conversion_factor
def convert_mass(self, mass, original_unit, prime_unit='kilogram'):
"""
mass:
type <int / float / array>
original_unit:
type <str>
prime_unit:
type <str>
"""
conversion_factor = self.mass_conversion_factors[prime_unit][original_unit]
return mass * conversion_factor
def convert_velocity(self, distance, time, original_distance_unit, original_time_unit, prime_distance_unit='meter', prime_time_unit='second'):
"""
velocity:
type <int / float / array>
original_distance_unit:
type <str>
original_time_unit:
type <str>
prime_distance_unit:
type <str>
prime_time_unit:
type <str>
"""
return self.convert_distance(distance, original_distance_unit, prime_distance_unit) / self.convert_time(time, original_time_unit, prime_time_unit)
def convert_acceleration(self, distance, time, original_distance_unit, original_time_unit, prime_distance_unit='meter', prime_time_unit='second'):
"""
velocity:
type <int / float / array>
original_distance_unit:
type <str>
original_time_unit:
type <str>
prime_distance_unit:
type <str>
prime_time_unit:
type <str>
"""
return self.convert_distance(distance, original_distance_unit, prime_distance_unit) / self.convert_time(time, original_time_unit, prime_time_unit)**2
The class GravitationalDynamics
is used to store and update bodies and their associated parameters.
class GravitationalDynamics():
"""
This class contains methods to run a simulation of N bodies that interact
gravitationally over some time. Each body in the simulation keeps a record
of parameters (position, velocity, and acceleration) as time progresses.
"""
def __init__(self, bodies, t=0, gravitational_constant=6.67e-11):
"""
bodies:
type <tuple / list / array>
t:
type <int / float>
gravitational_constant:
type <float>
"""
self._bodies = bodies
self.nbodies = len(bodies)
self.ndim = len(bodies[0].position)
self.t = t
self._moments = [t]
self.gravitational_constant = gravitational_constant
# self._collision_times = []
@property
def bodies(self):
return self._bodies
@property
def moments(self):
return np.array(self._moments)
def get_acceleration(self, ibody, jbody):
"""
ibody:
type <custom class: Body>
jbody:
type <custom class: Body>
"""
dpos = ibody.position - jbody.position
r = np.sum(dpos**2)
acc = -1 * self.gravitational_constant * jbody.mass * dpos / np.sqrt(r**3)
return acc
def update_accelerations(self):
for ith_body, ibody in enumerate(self.bodies):
ibody.acceleration *= 0
for jth_body, jbody in enumerate(self.bodies):
if ith_body != jth_body:
acc = self.get_acceleration(ibody, jbody)
ibody.acceleration += acc
ibody._vector_acceleration.append(np.copy(ibody.acceleration))
def update_velocities_and_positions(self, dt):
"""
dt:
type <int / float>
"""
for ith_body, ibody in enumerate(self.bodies):
ibody.velocity += ibody.acceleration * dt
ibody.position += ibody.velocity * dt
ibody._vector_velocity.append(np.copy(ibody.velocity))
ibody._vector_position.append(np.copy(ibody.position))
def simulate(self, dt, duration):
"""
dt:
type <int / float>
duration:
type <int / float>
"""
nsteps = int(duration / dt)
for ith_step in range(nsteps):
self.update_accelerations()
self.update_velocities_and_positions(dt)
self.t += dt
self._moments.append(self.t)
## check collisions
One can use the algorithm above using the commands below. The parameter values of each planet are taken from wolfram alpha.
UC = UnitConversions()
m_sun = UC.convert_mass(1, original_unit='solar mass', prime_unit='kilogram')
Sun = Body(identifier='Sun', facecolor='yellow', mass=m_sun, position=[0, 0, 0])
m_mercury = UC.convert_mass(0.05227, original_unit='earth mass', prime_unit='kilogram')
d_mercury = UC.convert_distance(0.4392, original_unit='Astronomical Unit', prime_unit='meter')
v_mercury = UC.convert_distance(106000, original_unit='mile', prime_unit='meter') / UC.convert_time(1, original_unit='hour', prime_unit='second')
Mercury = Body(identifier='Mercury', facecolor='gray', mass=m_mercury, position=[d_mercury, 0, 0], velocity=[0, v_mercury, 0])
m_earth = UC.convert_mass(1, original_unit='earth mass', prime_unit='kilogram')
d_earth = UC.convert_distance(1, original_unit='Astronomical Unit', prime_unit='meter')
v_earth = UC.convert_distance(66600, original_unit='mile', prime_unit='meter') / UC.convert_time(1, original_unit='hour', prime_unit='second')
Earth = Body(identifier='Earth', facecolor='blue', mass=m_earth, position=[d_earth, 0, 0], velocity=[0, v_earth, 0])
m_mars = UC.convert_mass(0.1704, original_unit='earth mass', prime_unit='kilogram')
d_mars = UC.convert_distance(1.653, original_unit='Astronomical Unit', prime_unit='meter')
v_mars = UC.convert_distance(53900, original_unit='mile', prime_unit='meter') / UC.convert_time(1, original_unit='hour', prime_unit='second')
Mars = Body(identifier='Mars', facecolor='darkred', mass=m_mars, position=[0, d_mars, 0], velocity=[v_mars, 0, 0])
bodies = [Sun, Mercury, Earth, Mars]
dt = UC.convert_time(2, original_unit='day', prime_unit='second')
duration = UC.convert_time(10, original_unit='year', prime_unit='second')
GD = GravitationalDynamics(bodies)
GD.simulate(dt, duration)
To check that the simulation ran, one can check various parameters of each body in GD.bodies
. Alternatively, we can visualize an image using the commands below (animation is provided in github link above). The image on the left shows a sequence of scattered points (one (x, y, z)
point with scaling in AU for each time t), and the image on the right shows that the orbits are elliptical (since the closest and farthest points appear periodic over time (AU against years).
fig = plt.figure(figsize=(11, 7))
ax_left = fig.add_subplot(1, 2, 1, projection='3d')
ax_left.set_title('3-D Position')
ax_right = fig.add_subplot(1, 2, 2)
ax_right.set_title('Displacement vs Time')
ts = UC.convert_time(GD.moments, original_unit='second', prime_unit='year')
xticks = np.arange(max(ts)+1).astype(int)
for body in GD.bodies:
vector = body.get_vector('position')
vector_coordinates = UC.convert_distance(vector, original_unit='meter', prime_unit='Astronomical Unit')
scalar = body.get_scalar('position')
scalar_coordinates = UC.convert_distance(scalar, original_unit='meter', prime_unit='Astronomical Unit')
ax_left.scatter(*vector_coordinates.T, marker='.', c=body.facecolor, label=body.identifier)
ax_right.scatter(ts, scalar_coordinates, marker='.', c=body.facecolor, label=body.identifier)
ax_right.set_xticks(xticks)
ax_right.grid(color='k', linestyle=':', alpha=0.3)
fig.subplots_adjust(bottom=0.3)
fig.legend(loc='lower center', mode='expand', ncol=len(GD.bodies))
plt.show()
plt.close(fig)
I should note that the simulation is only an approximation. Using both autocorrelation and power spectral density in a simulation containing all 8 planets of our solar system, some of the orbital periods appear exact (such Earth using velocity or acceleration), while others are off by some non-negligible amount (such as Mercury). I don't know whether or not to attribute this to float precision, the method of integration, the use of a discrete time-step to describe continuous motion, or the lack of relativity. That said, the method works for the lower-order case; the figure above shows that Earth's orbital period is one year, with fluctuations about 1 AU.
How can this algorithm be improved? Is anything implemented above something to be frowned upon? Trying to "self-learn" as much as possible, so (please) don't be nice.