N-Body Gravitational Simulation of Point-Masses in Python

Question

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.

Answer 1

Below is a list of my suggestions. You can find my changes to your code here: https://gist.github.com/bjourne/6270cff58c0238bb863e189ea63d6254

Kill your darlings

There's a saying in movie production that one must kill one's darlings. It means that if something is not serving it's purpose it must be deleted even if the author is fond of it or has worked very hard with it. I can tell that you have worked hard with the UnitConversions class, but it is an example of such a darling. You can replace it with constants:

# Masses
SOLAR_MASS = 1.988e30
EARTH_MASS = 5.9722e24

# Distances
ASTRO_UNIT = 1.496e11
MILE = 1609

# Durations
HOUR = 3600
DAY = 24 * HOUR
YEAR = 365 * DAY

m_earth = 1 * EARTH_MASS
d_earth = 1 * ASTRO_UNIT
v_earth = (66_600 * MILE) / (1 * HOUR)

Put code in functions or classes

It looks like your main code is running at the module level. That can cause weird problems with shadowed variables. I suggest always using the following structure:

def main():
    ... main code here...
if __name__ == '__main__':
    main()

Follow PEP 8

You are violating some of the style guidelines in PEP 8. For example, you use initial uppercase letters in some variable names and many lines are longer than 79 characters.

Delete stuff you don't use

I found lots of ununsed variables and methods in the Body class; notes, __str__(), scalar_velocity(), scalar_acceleration(), vector_acceleration(), and vector_velocity().

Don't put type declarations in comments

Personally, I'm not a fan of Python's optional type declarations. But if you want to use them, you should put them in the parameter lists like this

def get_vector(self, parameter: str):
    attribute = 'vector_{}'.format(parameter)
    return getattr(self, attribute)

and not in the comments.

Prefer positional arguments over keyword arguments

What I mean is that instead of

Sun = Body(identifier='Sun', facecolor='yellow',
           mass=m_sun, position=[0, 0, 0])

you can just write

Sun = Body('Sun', 'yellow', [0, 0, 0], m_sun)

which makes the code less verbose. Keyword arguments are sometimes useful but I don't think they are useful in your code.

Avoid "magic"

Using features that depend on introspection makes code harder to understand. So instead of

def get_scalar(self, parameter):
    attribute = 'scalar_{}'.format(parameter)
    return getattr(self, attribute)
...
scalar = body.get_scalar('position')

simply write

scalar = body.scalar_position

Avoid @property

Methods declared as properties can be misleading to readers because the method call is "hidden." I would suggest avoiding them unless you have compelling reasons to use them (that they look nice is not a good reason :)).

Don't hide privates

Java recommends you to hide implementation details. But that is not the case in Python and it is ok for users of objects to access member attributes directly.

Consider using abbreviations

Shorter variable names are preferable because they are easier to read. On the other hand, short names doesn't carry as much information as long ones so one has to try and find a happy middle ground. In your code, I suggest that you use pos, vel and acc as abbreviations for position, velocity and acceleration respectively. For someone interested in N-body simulations, these abbreviations should be self-explanatory.