numeric_integration.py

#!/usr/bin/env python3

import numpy as np
from scipy.integrate import odeint
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import animation


# define constants/units

year = 0.0109
r_moon_earth = 2.5e-3
m_sun = 3.33e5
m_moon = 0.012
m_jupiter = 318


# define system and its initial state

masses = []
positions = []
velocities = []


def initial_speed(r, m=m_sun):
    return np.sqrt(m/r)


def rotate(v, phi):
    return [np.cos(phi) * v[0] - np.sin(phi) * v[1],
            np.sin(phi) * v[0] + np.cos(phi) * v[1]]


def add_celestial_body(r, m, mother=None, label="celestial body"):
    """Add a some celestial object, e.g. a planet or a moon.

    If mother is not None, the radius r is interpreted with respect to it;
    mother can be the returned tuple of a prior call of this function. Do not
    try to add a moon to a moon."""

    # determine radius and angle
    if mother is not None:
        last_radius, last_speed, last_mass, last_angle = mother
        r_to_sun = last_radius + r
        rotation_angle = last_angle
    else:
        r_to_sun = r
        rotation_angle = np.random.rand() * 2*np.pi
    last_radius = r
    last_angle = rotation_angle

    # store Cartesian position and velocity
    positions.extend(rotate([r_to_sun, 0], rotation_angle))
    if not r_to_sun == 0:
        if mother is not None:
            speed = initial_speed(r, last_mass) + last_speed
        else:
            speed = initial_speed(r_to_sun)
        velocities.extend(rotate([0, speed], rotation_angle))
        last_speed = speed
    else:
        velocities.extend([0, 0])
        last_speed = 0

    # store mass
    masses.append(m)
    last_mass = m

    return last_radius, last_speed, last_mass, last_angle


add_celestial_body(0, 3.33e5, label="Sun")
add_celestial_body(0.414, 0.00553, label="Mercury")
add_celestial_body(0.719, 0.815, label="Venus")
earth = add_celestial_body(1, 1, label="Earth")
add_celestial_body(r_moon_earth, 0.0123, mother=earth, label="Moon")
add_celestial_body(1.66, 0.107, label="Mars")
add_celestial_body(5.41, 318, label="Jupiter")

# use colors corresponding to planets
colors = ['yellow', 'grey', 'orange', 'cornflowerblue', 'silver', 'red',
          'darkorange']
mpl.rcParams['axes.color_cycle'] = colors

# represent initial state as phase space point and use numpy arrays
# NOTE: we used lists before as we dealt with variable lengths

nobj = len(masses)
y0 = np.array(positions + velocities)
positions = y0[:2 * nobj]
velocities = y0[2 * nobj:]


# transform to COM system

com_position = np.zeros(2)
for mass, position in zip(masses, np.split(positions, nobj)):
    com_position += mass * position
com_position /= sum(masses)
for i in range(nobj):
    positions[i * 2:(i + 1) * 2] -= com_position
com_velocity = np.zeros(2)
for mass, velocity in zip(masses, np.split(velocities, nobj)):
    com_velocity += mass * velocity
com_velocity /= sum(masses)
for i in range(nobj):
    velocities[i * 2:(i + 1) * 2] -= com_velocity


# integrate the system

# a Mercury year is roughly a fifth of an Earth year, we want 250 integration
# during that period
dt = year/5/250
tmax = 5*year
ts = np.arange(0.0, tmax, dt)


def f(y, t):
    """Determine derivatives for the phase space point y."""

    positions = y[:2 * nobj]
    velocities = y[2 * nobj:]

    f = np.zeros(len(y))
    f[:2 * nobj] = velocities

    accelerations = f[2 * nobj:]
    for i in range(nobj - 1):
        position_i = positions[i * 2:(i + 1) * 2]
        for j in range(i + 1, nobj):
            position_j = positions[j * 2:(j + 1) * 2]
            dr = position_i - position_j
            distance = np.linalg.norm(position_i - position_j)
            dr_by_distance3 = dr / distance**3
            accelerations[i*2:(i+1)*2] += - masses[j] * dr_by_distance3
            accelerations[j*2:(j+1)*2] += masses[i] * dr_by_distance3

    return f


result = odeint(f, y0, ts)
positions = result[:, :result.shape[1]/2]


# setup plotting
fig = plt.figure()
ax = fig.add_subplot(111,
                     autoscale_on=False, aspect='equal')
ax.grid()

lines = []
objects = []
for i in range(nobj):
    lines.append(ax.plot([], [], color=colors[i])[0])
    objects.append(ax.plot([], [], 'o', color=colors[i])[0])
time_template = 't = %.1fa'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)


# animate

def init():
    for line, celestial_object in zip(lines, objects):
        line.set_data([], [])
        celestial_object.set_data([], [])
    ax.set_xlim(-6, 6)
    ax.set_ylim(-6, 6)
    time_text.set_text('')
    return lines, objects, ax, time_text


def timing_function(t):
    return (t**2) * (3 - 2 * t)


def animate(i):
    i *= 10
    for j, line in enumerate(lines):
        position_j = positions[:i, j * 2:(j + 1) * 2]
        line.set_data(*position_j.transpose())
        objects[j].set_data(*position_j[-1:].transpose())
    timing = timing_function(i/tmax*dt)
    scaling_factor = 1 - timing * 0.98
    ax.set_xlim(-6*scaling_factor, 6*scaling_factor)
    ax.set_ylim(-6*scaling_factor, 6*scaling_factor)

    time_text.set_text(time_template % (i*dt/year))
    return lines, objects, ax, time_text


ani = animation.FuncAnimation(fig, animate, frames=result.shape[0]//10,
                              blit=False, init_func=init, repeat=True)
ani.save('solar_system.mp4', fps=30, extra_args=['-vcodec', 'libx264'])


# others plots

plt.subplot(aspect='equal')
plt.plot(*positions.transpose())
plt.savefig('ss_fs.pdf')

plt.subplots()
plt.plot(ts, positions[:, ::2])
plt.savefig('ss_positions.pdf')

plt.subplots()
earth_radii = np.sqrt(positions[:, 6]**2 + positions[:, 7]**2)
plt.plot(ts, earth_radii, color=colors[3])
plt.ylim(0.992, 1.008)
plt.savefig('ss_earth_radii.pdf')