N-body problem with differents solvers (RK2, RK4, Euler symplectic, Stormer-Verlet) : planets drift to infinity

Question

I'm trying to write an integrator for the 2 and 3-body problem. I choose to start from a generalisation to N-body problem so I can just pass my bodies to the same integrator in the two cases.

I'm using Hamiltonian mechanincs, the Hamiltonian I get for the N-body problem is the following,

$$ \mathcal{H} = \sum_{i=1}^N \frac{\vec{p}_i^2}{2m_i} - \sum_{i<j}^{N} \frac{Gm_i m_j}{|\vec{r}_i - \vec{r}_j|} $$ where $|\vec{r}_i - \vec{r}_j| \equiv \vec{r}_{ij}$ is the relative distance between the masses $i$ and $j$.

Following that, the canonical equations for the mass $i$ are,

\begin{align} \dot{\vec{r}}_i = \begin{pmatrix} \dot{x}_i \\ \dot{y}_i \\ \dot{z}_i \end{pmatrix} = \begin{pmatrix} \frac{p_{x_i}}{m_i} \\ \frac{p_{y_i}}{m_i} \\ \frac{p_{z_i}}{m_i} \\ \end{pmatrix} \end{align}

\begin{align} \dot{\vec{p}}_i = \begin{pmatrix} \dot{p}_{x_i} \\ \dot{p}_{y_i} \\ \dot{p}_{z_i} \end{pmatrix} = \begin{pmatrix} - \sum_{j \neq i}^{N} \frac{Gm_im_j}{(\vec{r}_{ij})^3} (x_i - x_j) \\ - \sum_{j \neq i}^{N} \frac{Gm_im_j}{(\vec{r}_{ij})^3} (y_i - y_j) \\ - \sum_{j \neq i}^{N} \frac{Gm_im_j}{(\vec{r}_{ij})^3} (z_i - z_j) \end{pmatrix} \end{align}

I am doing the simulation for the four following methods : RK2, RK4, Euler Symplectic, Stormer-Verlet.

Simulating for 50 years with a step time of 30 days (and using astronomical units), in any cases I get straight lines in the 3D plots,

EDIT : Found the error. See the comment below. Post edited with correction (as well as in the plot)

Definitely something wrong. I've been checking the code multiple times so unless I'm wrong with the equations I don't know were is the error.

Here is the code,

# index.py
#!/usr/bin/env python3.9
import click
import matplotlib.pyplot as plt
import pyfiglet # ascii art

from src.two_body import TwoBody
from src.nbody import NBodySimulation
from src.body import Body

from consts import (sun_position0, sun_impulsion0, jupiter_position0, jupiter_impulsion0, saturn_position0, saturn_impulsion0, M_sun, M_jup, M_sat)
from consts import (t0, tN, dt)

plt.style.use("science")

# prepare bodies
Sun = Body("Sun", sun_position0, sun_impulsion0, M_sun)
Jupiter = Body("Jupiter", jupiter_position0, jupiter_impulsion0, M_jup)
Saturn = Body("Saturn", saturn_position0, saturn_impulsion0, M_sat)

@click.command()
@click.option("--stype", default="twobody", help="2-Body problem or 3-Body problem")
@click.option("--plot", default="static", help="static or animated plot")
def main(stype, plot):
  # clear terminal (even history)
  print('\033c', end=None)
  # ascii art - for fun.
  print(pyfiglet.print_figlet("CELESTIAL"))
  if stype == "twobody":
    twobody = NBodySimulation(bodies=[Sun, Jupiter], t0=t0, tN=tN, dt=dt)
    twobody.simulate()
    if plot == "static":
      twobody.plot()
    elif plot == "animated":
      twobody.animate()

  elif stype == "threebody":
    threeBody = NBodySimulation(bodies=[Sun, Jupiter, Saturn], t0=t0, tN=tN, dt=dt)
    threeBody.simulate()
    if plot == "static":
      threeBody.plot()
    elif plot == "animated":
      threeBody.animate()

if __name__ == "__main__":
  main()

# consts.py
# https://nssdc.gsfc.nasa.gov/planetary/factsheet/saturnfact.html
# [M]: kg
#M_sat = 5.6834e26
#M_jup = 1.89812e27
#M_sun = 1.98841e30

# masses relative to the sun
M_sun = 1.00000597682
M_jup = 0.000954786104043
M_sat = 0.000285583733151

# 1 m = 6.6845... au
meter_to_au = 6.6845871222684e-12

# 1 day = 24 * 60 * 60 s
day_to_second = 24 * 60 * 60

# [G]: m^3 / kg s^2 => using a.u and days [G]: a.u / kg days^2
# G = (6.67430e-11 * (meter_to_au ** 3) * (day_to_second ** 2))
G = 2.95912208286e-4

# initial conditions
# positions are in a.u (astronomical units)
# speeds are in a.u / day => speed * mass = impulsion
# masses are in kg
# http://vo.imcce.fr/webservices/miriade/?forms

sun_position0 = [0., 0., 0.]
sun_impulsion0 = [0., 0., 0.]

jupiter_position0 = [3.4707903364632, -3.3666298150704, -1.5275164476857]
jupiter_impulsion0 = [0.0054161298800 * M_jup, 0.0051281500076 * M_jup, 0.0020662323714 * M_jup]

saturn_position0 = [5.8139930169916, -7.3998049325634, -3.3068297277913]
saturn_impulsion0 = [0.0042287765130 * M_sat, 0.0030656447687 * M_sat, 0.0010842701680 * M_sat]

# time in days
t0 = 0
tN = 5000 * 365.25 # integration over 5000 years
dt = 30 # time_step : 30 days

# body.py
class Body():
  """
  Body (Sun, Jupiter,...) in cartesian coordinates
  """
  def __init__(self, name, initial_positions, initial_impulsions, mass):
    self.name = name
    self.initial_positions = initial_positions
    self.initial_impulsions = initial_impulsions
    self.mass = mass

# nbody.py
import random
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from mpl_toolkits.mplot3d import Axes3D

from .edo import (n_body_dqdt, n_body_dpdt)
from .solvers import (heun, rk4, euler_symp, stormer_verlet)

class NBodySimulation():
  def __init__(self, bodies, t0, tN, dt):
    self.bodies = bodies
    self.t0 = t0
    self.tN = tN
    self.dt = dt
    self.legends = ["Heun (RK2)", "RK4", "Euler Symplectique", "Stormer-Verlet"]

    self.solvers = [
      {"call": heun, "name": "Heun (RK2)"},
      {"call": rk4, "name": "RK4"},
      {"call": euler_symp, "name": "Euler Symplectique"},
      {"call": stormer_verlet, "name": "Stormer Verlet"}
    ]

    self.results = []

  def solve(self, solver):
    # number of time step
    nt = int((self.tN - self.t0) / self.dt)

    # positions and impulsions state vectors
    # each body is in \R^3 (3D space x, y, z)
    q = np.zeros((nt, len(self.bodies) * 3))
    p = np.zeros((nt, len(self.bodies) * 3))

    # set initial conditions
    q[0] = np.concatenate(np.array([body.initial_positions for body in self.bodies]))
    p[0] = np.concatenate(np.array([body.initial_impulsions for body in self.bodies]))

    return solver(dqdt=n_body_dqdt, dpdt=n_body_dpdt, q=q, p=p, dt=self.dt, nt=nt, bodies=self.bodies)

  def simulate(self):
    self.results = []
    for solver in self.solvers:
      q, p = self.solve(solver["call"])
      self.results.append({"solver": solver["name"], "q": q, "p": p})

  def plot(self):
    self.fig = plt.figure(figsize=(8,8))
    colors = ['r', 'b', 'g', 'y', 'm', 'c']

    # loop for each result (corresponding to a specific solving method)
    for (index, result) in enumerate(self.results):
      # create a 3D plot
      ax = self.fig.add_subplot(2, 2, index + 1, projection="3d")
      # set plot parameters
      ax.set_title(result["solver"])

      max_range = 0

      # plotting each body
      for (ind, body) in enumerate(self.bodies):
        x_index = (ind * 3)
        y_index = (ind * 3) + 1
        z_index = (ind * 3) + 2

        x, y, z = result["q"][:,x_index], result["q"][:,y_index], result["q"][:,z_index]

        max_dim = max(max(x), max(y), max(z))
        if max_dim > max_range:
          max_range = max_dim

        ax.plot(xs=x, ys=y, zs=z, c=random.choice(colors), label=body.name)

      # limiting plot
      ax.set_xlim([-max_range,max_range])
      ax.set_ylim([-max_range,max_range])
      ax.set_zlim([-max_range,max_range])

      #ax.legend()

    plt.show()

  def animate(self):
    pass

# solvers.py
import numpy as np
import tqdm

def rk2_derivatives(edo, qk, pk, dt, bodies):
  k1 = dt * edo(qk, pk, bodies)
  k2 = dt * edo(qk + (dt * k1), pk + (dt * k1), bodies)
  return (k1 + k2) / 2.

def heun(dqdt, dpdt, q, p, dt, nt, bodies):
  for k in range(0, nt - 1):
    qk, pk = q[k], p[k]
    q[k + 1] = qk + rk2_derivatives(dqdt, qk, pk, dt, bodies)
    p[k + 1] = pk + rk2_derivatives(dpdt, qk, pk, dt, bodies)

  return q, p


def rk4_derivatives(edo, qk, pk, dt, bodies):
  k1 = dt * edo(qk, pk, bodies)
  k2 = dt * edo(qk + ((dt / 2) * k1), pk + ((dt / 2) * k1), bodies)
  k3 = dt * edo(qk + ((dt / 2) * k2), pk + ((dt / 2) * k2), bodies)
  k4 = dt * edo(qk + (dt * k3), pk + (dt * k3), bodies)
  return (k1 + 2*k2 + 2*k3 + k4) / 6.

def rk4(dqdt, dpdt, q, p, dt, nt, bodies):
  for k in range(0, nt - 1):
    qk, pk = q[k], p[k]
    q[k + 1] = qk + rk4_derivatives(dqdt, qk, pk, dt, bodies)
    p[k + 1] = pk + rk4_derivatives(dpdt, qk, pk, dt, bodies)

  return q, p


def euler_symp(dqdt, dpdt, q, p, dt, nt, bodies):
  for k in range(0, nt - 1):
    qk, pk = q[k], p[k]
    p[k + 1] = pk + dt * dpdt(qk, pk, bodies)
    q[k + 1] = qk + dt * dqdt(qk, p[k + 1], bodies)

  return q, p


def stormer_verlet(dqdt, dpdt, q, p, dt, nt, bodies):
  for k in range(0, nt - 1):
    qk, pk = q[k], p[k]
    p_half = pk + (dt / 2) * dpdt(qk, pk, bodies)
    q[k + 1] = qk + dt * dqdt(qk, p_half, bodies)
    p[k + 1] = p_half + (dt / 2) * dpdt(q[k + 1], p_half, bodies)

  return q, p

# edo.py
import numpy as np

from consts import (M_sun as m1, M_jup as m2, M_sat as m3, G)


def n_body_dqdt(qk, pk, bodies):
  """
  :param qk: n-body position state vector => qk: [x1, y1, z1, x2, y2, z2,..., xN, yN, zN]
  :param pk: n-body impulsion state vector => pk: [px1, py1, pz1,..., pxN, pyN, pzN]
  :param bodies: set of bodies (Sun, Jupiter,...)
  :type q_state: ndarray
  :type p_state: ndarray
  :type bodies: ndarray

  :return: \dot{q} = dh/dp
  :rtype: ndarray
  """
  dqdt = np.zeros(len(qk))
  for i in range(len(bodies)):
    offset = i * 3
    # px_i / m_i
    dqdt[offset] = pk[offset] / bodies[i].mass
    # py_i / m_i
    dqdt[offset + 1] = pk[offset + 1] / bodies[i].mass
    # pz_i / m_i
    dqdt[offset + 2] = pk[offset + 2] / bodies[i].mass

  return dqdt

def n_body_dpdt(qk, pk, bodies):
  """
  :param qk: n-body position state vector at k * dt time => qk: [x1, y1, z1, x2, y2, z2,..., xN, yN, zN]
  :param pk: n-body impulsion state vector at k * dt time => pk: [px1, py1, pz1,..., pxN, pyN, pzN]
  :param bodies: set of bodies (Sun, Jupiter,...)
  :type qk: ndarray
  :type pk: ndarray
  :type bodies: ndarray

  :return: \dot{p} = - dh/dr => \dot{p} = [\dot{px1}, \dot{py2}, \dot{pz1},..., \dot{pxN}, \dot{pyN}, \dot{pzN}]
  :rtype: ndarray
  """
  dpdt = np.zeros(len(qk))
  # loop each body
  for i in range(len(bodies)):
    i_offset = i * 3
    for j in range(len(bodies)):
      if j != i:
        j_offset = j * 3
        dpdt[i_offset] -= ((G * bodies[i].mass * bodies[j].mass) / r_dist(ri=[qk[i_offset], qk[i_offset + 1], qk[i_offset + 2]], rj=[qk[j_offset], qk[j_offset + 1], qk[j_offset + 2]]) ** 3) * (qk[i_offset] - qk[j_offset])

        dpdt[i_offset + 1] -= ((G * bodies[i].mass * bodies[j].mass) / r_dist(ri=[qk[i_offset], qk[i_offset + 1], qk[i_offset + 2]], rj=[qk[j_offset], qk[j_offset + 1], qk[j_offset + 2]]) ** 3) * (qk[i_offset + 1] - qk[j_offset + 1])

        dpdt[i_offset + 2] -= ((G * bodies[i].mass * bodies[j].mass) / r_dist(ri=[qk[i_offset], qk[i_offset + 1], qk[i_offset + 2]], rj=[qk[j_offset], qk[j_offset + 1], qk[j_offset + 2]]) ** 3) * (qk[i_offset + 2] - qk[j_offset + 2])

  # divide by two to avoid double count
  return dpdt

def r_dist(ri, rj):
  """
  Calculate distance between two points in \R^3
  :param ri: point of \R^3: [xi, yi, zi]
  :param rj: point of \R^3: [xj, yj, zj]

  :return: distance betwen the two points
  :rtype: float
  """
  return np.sqrt((ri[0] - rj[0]) ** 2 + (ri[1] - rj[1]) ** 2 + (ri[2] - rj[2]) ** 2)

Eventually, if you want an example to to reproduce the problem: https://github.com/Mathieu-R/celestial-mechanics

Thanks in advance.

Answer 1

All-over this is a nicely structured code. The main problems are related to the Runge-Kutta solvers, where the first-order system was not uniformly applied to the computation.

What is obviously wrong can be found in the first lines of the solver file

def rk2_derivatives(edo, qk, pk, dt, bodies):
  k1 = dt * edo(qk, pk, bodies)
  k2 = dt * edo(qk + (dt * k1), pk + (dt * k1), bodies)
  return (k1 + k2) / 2.

The update for qk and pk is rarely given by the same vector k1, each needs its own k1 from the associated derivative function. Also, the factor dt is applied twice where it only should be applied in one place. All-in-all this makes for an order 1 method, as the derivatives are still computed close to the point (qk,pk).

Interleaving the derivatives computations of the coupled system correctly, it should be

def rk2_derivatives(dqdt, dpdt, qk, pk, dt, bodies):
  k1q, k1p = (dt * edo(qk, pk, bodies) for edo in (dqdt, dpdt))
  qt, pt = qk+k1q, pk+k1p # temporary or intermediary variables
  k2q, k2p = (dt * edo(qt, pt, bodies) for edo in (dqdt, dpdt))
  return (k1q + k2q) / 2., (k1p + k2p) / 2.

and then used in the RK2-Heun method as

def heun(dqdt, dpdt, q, p, dt, nt, bodies):
  for k in range(0, nt - 1):
    qk, pk = q[k], p[k]
    dqk, dpk = rk2_derivatives(dqdt, dpdt, qk, pk, dt, bodies)
    q[k+1] = qk + dqk
    p[k+1] = pk + dpk
  return q, p

The same goes for the step update function of RK4.

---

Now with only these corrections one gets to see an additional problem. As the data is taken from a larger ensemble, the location and velocity of the center-of-mass of the 3-body system are not zero (even if they might be for the full solar system). The resulting drift masks all differences between the methods in the trajectory plots.

Removing the average velocity from the initial impulses results in a very stable picture over 500 years with no visible differences, and over 5000 years the differences are still minimal, but well visible

The RK2-Heun method is visibly different from the others due to its order 2. The symplectic methods are practically the same, the difference is a shift in the initial impulse by half a time step. The energy is preserved to 2nd order, due to the time symmetry other errors accumulate in the 4th order, which is the same error order as RK4. In a more detailed analysis of the solutions it should turn out that the symplectic methods do have a second order deviation in the orbit period. So the orbits follow the same curve as a more exact solution, but at a different time table. RK4 should have a 4th order error also there.

---

additional code variations

edo.py The computation can be arranged to be much more compact and thus more easily maintainable. There might be a slight penalty for the data view re-organization, but with the below version a lot of duplicate computation is avoided.

def n_body_dqdt(qk, pk, bodies):
  # reorganize data as list of 3D vectors
  pk = pk.reshape([-1,3])
  dqdt = np.zeros(pk.shape)
  for i,pi in enumerate(pk):
    # dq_i/dt = p_i / m_i
    dqdt[i] = pi / bodies[i].mass
  # return an unstructured array
  return dqdt.flatten()

def n_body_dpdt(qk, pk, bodies):
  # reorganize data as list of 3D vectors
  qk = qk.reshape([-1,3])
  dpdt = np.zeros(qk.shape)
  # loop each body
  for i in range(len(bodies)):
    for j in range(len(bodies)):
      if j != i:
        r_ij = r_dist(qk[i],qk[j])
        dpdt[i] -= ((G * bodies[i].mass * bodies[j].mass) / r_ij ** 3) * (qk[i] - qk[j])
  # return an unstructured array
  return dpdt.flatten()

nbody.py In the initialization shift the coordinate frame so that the barycenter is at rest (and at zero)

class NBodySimulation():
  def __init__(self, bodies, t0, tN, dt):
    self.bodies = bodies
    total_mass = sum(b.mass for b in bodies)
    mean_pos = sum(b.initial_positions * b.mass for b in bodies)/total_mass
    mean_vel = sum(b.initial_impulsions for b in bodies)/total_mass
    for b in self.bodies:
      b.initial_positions -= mean_pos
      b.initial_impulsions -= b.mass * mean_vel
    ...

body.py For that to work, make the fields of the body class into numpy arrays

from numpy import array

class Body():
  """
  Body (Sun, Jupiter,...) in cartesian coordinates
  """
  def __init__(self, name, initial_positions, initial_impulsions, mass):
    self.name = name
    self.initial_positions = array(initial_positions)
    self.initial_impulsions = array(initial_impulsions)
    self.mass = mass