Quick and dirty solar system dynamics

sim.py

import numpy as np
from numpy.linalg import norm as l2
import matplotlib.pyplot as plt


class PlanetarySim():

    # Masses e24 scale
    DATA = (
        ('SUN', 1.98847e6),
        ('MERCURY', 0.330),
        ('VENUS', 4.87),
        ('EARTH', 5.97),
        ('MOON', 0.073),
        ('MARS', 0.642),
        ('JUPYTER', 1898),
        ('SATURN', 568),
        ('URANUS', 86.8),
        ('NEPTUNE', 102),
        ('PLUTO', 0.014)
    )

    G = 39

    def __init__(self):
        # Convert to solar masses
        solar_mass = self.DATA[0][1]
        self.nplanets = len(self.DATA)

        self.m = np.empty(self.nplanets)
        for i, (name, m) in enumerate(self.DATA):
            self.m[i] = m / solar_mass


    def step(self, s, v, f, dt):
        nplanets = self.nplanets

        # Estimate the forces
        f_new = np.zeros((nplanets, 3))
        for i in range(nplanets):
            for j in range(nplanets):
                if i == j:
                    continue
                r = s[j] - s[i]
                nr = l2(r)
                f_new[i, :] += self.G * self.m[j] * r / nr * nr * nr

        f = f_new if f is None else f

        # Integration step
        s += v * dt + f * dt * dt / 2
        v += (f + f_new) * dt / 2
        return s, v, f_new

    def sim(self, ntot, dt=0.001):
        nplanets = self.nplanets

        s = np.zeros((nplanets, 3))
        v = np.zeros((nplanets, 3))

        s[:, 0] = np.linspace(0, -1, nplanets, endpoint=True)
        v[:, 1] = -np.linspace(0, 5, nplanets, endpoint=True)

        orbit = np.empty((nplanets, ntot, 3))
        f = None
        for n in range(ntot):
            s, v, f = self.step(s, v, f, dt)
            orbit[:, n, :] = s

        return orbit


def main():
    p = PlanetarySim()
    orbit = p.sim(1000)
    print(orbit.shape)
    for pi in orbit:
        plt.plot(pi[:, 0], pi[:, 1])

    plt.ylim(-2, 2)
    plt.xlim(-2, 2)
    plt.show()



if __name__ == "__main__":
    main()