0
$\begingroup$

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)

3D 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.

$\endgroup$
5
  • 1
    $\begingroup$ I'm not able to understand properly your code, but It could probably be a problem of units. I worked on a problem similar to yours recently and usually I got that kind of plots because of bad choice of units. I was using your same units actually and I got better results by using solar masses instead of kilograms. $\endgroup$ – Zebx Mar 29 at 15:30
  • $\begingroup$ It may be interesting to study the values of your time derivatives (or the forces directly) at the initial time. You may spot some issues, such as @Zebx suggested. $\endgroup$ – Laurent90 Mar 29 at 21:07
  • $\begingroup$ I found the error. It was a minus sign instead of a plus sign between the two terms in the Hamiltonian and there was an error in the position derivative ($dq/dt$). Post edited. $\endgroup$ – Mathieu Rousseau Mar 30 at 12:05
  • $\begingroup$ I'm doubtful in your use of day_to_seconds, as that should be the fraction of a day in a second and thus the reciprocal of the current value. That is, if you want to use it in the same way as meter_to_au. /// It does not seem to be an error in the code, but it is inconsistent in the naming. /// Is the plot picture current for the corrected code? $\endgroup$ – Lutz Lehmann Mar 31 at 15:01
  • $\begingroup$ I edited the post. Now the plot is for the corrected code. $\endgroup$ – Mathieu Rousseau Mar 31 at 16:20
2
$\begingroup$

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

enter image description here

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
$\endgroup$
2
  • $\begingroup$ What do you mean by "Removing the average velocity from the initial impulses" ? From the initial impulsions $p(t=0)$ ? $\endgroup$ – Mathieu Rousseau Apr 1 at 13:41
  • $\begingroup$ You compute $v_{mean}=\sum p_k/\sum m_k$ and then correct $p_k := p_k-m_k·v_{mean}$. Yes, that all relates to the initial values. $\endgroup$ – Lutz Lehmann Apr 1 at 14:02

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.