Step 2: Gravity

Welcome to step 2. This is the most important step — computing the gravitational acceleration. Turns out this is also the most expensive part in N-body simulation, so we will spend some time on optimization.

Newton's law of gravitation

I believe most of you are familiar with Newton's law of gravitation

\[ \mathbf{F}_{ij} = m_{i} \mathbf{a}_{ij} = \frac{G m_i m_j}{r_{ij}^2} \hat{r}_{ij}, \]

where \(\mathbf{F}_{ij}\) is the force on particle \(i\) due to particle \(j\), and \(\hat{r}_{ij}\) is the unit vector pointing from particle \(i\) to particle \(j\). That is,

\[ \mathbf{r}_{ij} = \mathbf{r}_j - \mathbf{r}_i, \quad \hat{r}_{ij} = \mathbf{r}_{ij} / r_{ij}. \]

In practice, we are only interested in the acceleration. To compute the acceleration of particle \(i \in \{1, \dots, N\}\), we have

\[ \mathbf{a}_{i} = \sum_{j \neq i} \frac{G m_j}{r_{ij}^3} \mathbf{r}_{ij}, \]

which can be easily done as this only involve simple vector operations.

Implementation 1

Below shows our first naive implementation of the acceleration function. The outer loop \(i\) iterates over all particles, and the inner loop \(j\) iterates over all particles again to compute the acceleration between all pairs of particles. However, this implementation is very slow.

def acceleration_1(
    a: np.ndarray,
    system: System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    num_particles = system.num_particles
    x = system.x
    m = system.m
    G = system.G

    # Calculations
    for i in range(num_particles):
        for j in range(num_particles):
            if i == j:
                continue

            R = x[j] - x[i]
            a[i] += G * m[j] * R / (np.linalg.norm(R) ** 3)

Where is the return statement?

Actually, there is no need to return the acceleration array a because we are modifying the memory in-place.

Implementation 2

To optimize the code, we utilize the fact that the distance between particles \(i\) and \(j\) is the same:

\[ \mathbf{r}_{ij} = - \mathbf{r}_{ji} \implies r_{ij} = r_{ji}. \]

Note

In our notation, the lowercase, non-bold \(r_{ij}\) is the vector norm, which is always positive.

This allows us to effectively reduce half of the distance calculations. (Calculating the distance is quite expensive as it involves the computation of sqrt.) The outer loop \(i\) still iterates over all particles, but the inner loop \(j\) only iterates from \(i + 1\) to \(N\). (Why? because all combinations of \(i\) and \(j \leq i\) has already been computed in the previous iterations. Therefore, we only need to care about \(j > i\).)

Tip

If you want to rewrite this in C, this is the implementation you should use.

def acceleration_2(
    a: np.ndarray,
    system: System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    num_particles = system.num_particles
    x = system.x
    m = system.m
    G = system.G

    # Calculations
    for i in range(num_particles):
        for j in range(i + 1, num_particles):
            R = x[j] - x[i]
            temp_value = G * R / (np.linalg.norm(R) ** 3)
            a[i] += temp_value * m[j]
            a[j] -= temp_value * m[i]

Implementation 3* (Advanced)

The above implementation is still quite slow because we are using Python loops to iterate over the particles. NumPy is (partly) implemented in C, which makes it much faster than Python operations. If we were able to avoid Python loops completely, we can achieve a significant speedup. This can be done by vectorizing the code. Note that this would be quite difficult for beginners, but learning this could help you understand a lot about NumPy arrays.

1. We first compute a displacement matrix \(\mathbf{R}\), where \(\mathbf{R}_{ij} = \mathbf{r}_j - \mathbf{r}_i\). Therefore, it is a 3D array of shape \((N, N, 3)\), and the diagonal elements are all zero. This is computed by broadcasting first \(\mathbf{r}\) along the axis 0 (row) and the second \(\mathbf{r}\) along the axis 1 (column): r_ij = x[:, np.newaxis, :] - x[np.newaxis, :, :]

   \[ \mathbf{R} = \begin{bmatrix} \mathbf{r}_1 - \mathbf{r}_1 & \mathbf{r}_2 - \mathbf{r}_1 & \cdots & \mathbf{r}_N - \mathbf{r}_1 \\ \mathbf{r}_1 - \mathbf{r}_2 & \mathbf{r}_2 - \mathbf{r}_2 & \cdots & \mathbf{r}_N - \mathbf{r}_2 \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{r}_1 - \mathbf{r}_N & \mathbf{r}_2 - \mathbf{r}_N & \cdots & \mathbf{r}_N - \mathbf{r}_N \end{bmatrix} \]

2. We compute the distance matrix \(\mathbf{R}_{\mathrm{norm}}\), which is a 2D array of shape \((N, N)\). It is computed by taking the norm along the last axis with length 3 (r_norm = np.linalg.norm(r_ij, axis=2)).

3. We compute \(1 / \mathbf{R}_\text{norm}^3\) (element-wise). Because the diagonal elements are all zero, the division will produce undefined values along the diagonal. Therefore, we want to silence the warnings from NumPy and set the diagonal elements to zero.

   # Compute 1 / r^3
   with np.errstate(divide='ignore', invalid='ignore'):
       inv_r_cubed = 1.0 / (r_norm * r_norm * r_norm)
   
   # Set diagonal elements to 0 to avoid self-interaction
   np.fill_diagonal(inv_r_cubed, 0.0)
   

4. We compute the acceleration by

\[ \mathbf{a}_{i} = \sum_{j \neq i} \frac{G m_j}{r_{ij}^3} \mathbf{r}_{ij}, \]

The last step can be done by using NumPy's broadcasting feature. The resulting acceleration array will be of shape \((N, 3)\).

a[:] = G * np.sum(
    r_ij * inv_r_cubed[:, :, np.newaxis] * m[:, np.newaxis, np.newaxis], axis=0
)

This line of code is a bit complicated. Let us break it down:

• G is a constant, so it can be factored out of the summation.
• Ignore the last dimension with length 3 for now. We have

  a[:, 0] = G * np.sum(
      r_ij[:, :, 0] * inv_r_cubed[:, :] * m[:, np.newaxis], axis=0
  )
  

  We are summing over the axis 0 (row), so the mass vector \(\mathbf{m}\) needs to be broadcasted along the axis 1 (column) to \(\mathbf{M}\) with the shape of \((N, N)\). We have

\[ \mathbf{M} = \begin{bmatrix} m_1 & m_2 & \cdots & m_N \\ m_1 & m_2 & \cdots & m_N \\ \vdots & \vdots & \ddots & \vdots \\ m_1 & m_2 & \cdots & m_N \end{bmatrix}. \]

The element-wise multiplication of \(\mathbf{M}\) with \(\mathbf{R}\) divided by \(\mathbf{R}_\text{norm}^3\) gives

\[ \begin{bmatrix} 0 & m_2 \mathbf{x}_{12} / x_{12}^3 & \cdots & m_N \mathbf{x}_{1N} / x_{1N}^3 \\ m_1 \mathbf{x}_{21} / x_{21}^3 & 0 & \cdots & m_N \mathbf{x}_{2N} / x_{2N}^3 \\ \vdots & \vdots & \ddots & \vdots \\ m_1 \mathbf{x}_{N1} / x_{N1}^3 & m_2 \mathbf{x}_{N2} / x_{N2}^3 & \cdots & 0 \end{bmatrix} \]

We are summing along the axis 0 (row). For particle \(i\), we have

\[ \mathbf{a}_{i, 0} = G \left[m_1 \frac{\mathbf{x}_{i1}}{x_{i1}^3} + m_2 \frac{\mathbf{x}_{i2}}{x_{i2}^3} + \cdots + 0 + \cdots + m_N \frac{\mathbf{x}_{iN}}{x_{iN}^3} \right] = \sum_{j \neq i} \frac{G m_j}{x_{ij}^3} \mathbf{x}_{ij} \]

This is exactly what we want! Now, to also include the last dimension with length 3, we simple add np.newaxis to the last axis. This gives us the vectorized version of the full acceleration function. Later, we will see in the benchmark that this is much faster than the previous implementations.

def acceleration_3(
    a: np.ndarray,
    system: System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    x = system.x
    m = system.m
    G = system.G

    # Compute the displacement vector
    r_ij = x[:, np.newaxis, :] - x[np.newaxis, :, :]

    # Compute the distance
    r_norm = np.linalg.norm(r_ij, axis=2)

    # Compute 1 / r^3
    with np.errstate(divide='ignore', invalid='ignore'):
        inv_r_cubed = 1.0 / (r_norm * r_norm * r_norm)

    # Set diagonal elements to 0 to avoid self-interaction
    np.fill_diagonal(inv_r_cubed, 0.0)

    # Compute the acceleration
    a[:] = G * np.sum(
        r_ij * inv_r_cubed[:, :, np.newaxis] * m[:, np.newaxis, np.newaxis], axis=0
    )

Implementation 4* (Advanced)

In the last implementation, we are using np.sum and broadcasting to compute the acceleration. In NumPy, there is a faster method np.einsum. Therefore, in this implementation, we will replace the np.sum with np.einsum.

In our original implementation, notice how the broadcasting is done:

a[:] = G * np.sum(
    r_ij * inv_r_cubed[:, :, np.newaxis] * m[:, np.newaxis, np.newaxis], axis=0
)

Denote the axis 0, 1, and 2 as \(i\), \(j\), and \(k\) respectively.

• r_ij is a 3D array multiplied without broadcasting \(\implies ijk\).
• inv_r_cubed is a 2D array multiplied with broadcasting along axis 2 \(\implies ij\).
• m is a 1D array multiplied with broadcasting along axis 1 and 2 \(\implies i\).

The final sum is done along axis 0 \(\implies ijk \to jk\).

Using np.einsum, we can specify the indices to be summed over. The following line of code is equivalent to the original implementation:

a[:] = G * np.einsum("ijk,ij,i->jk", r_ij, inv_r_cubed, m)

Full implementation:

def acceleration_4(
    a: np.ndarray,
    system: System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    x = system.x
    m = system.m
    G = system.G

    # Compute the displacement vector
    r_ij = x[:, np.newaxis, :] - x[np.newaxis, :, :]

    # Compute the distance
    r_norm = np.linalg.norm(r_ij, axis=2)

    # Compute 1 / r^3
    with np.errstate(divide='ignore', invalid='ignore'):
        inv_r_cubed = 1.0 / (r_norm * r_norm * r_norm)

    # Set diagonal elements to 0 to avoid self-interaction
    np.fill_diagonal(inv_r_cubed, 0.0)

    # Compute the acceleration
    a[:] = G * np.einsum("ijk,ij,i->jk", r_ij, inv_r_cubed, m)

Benchmark

To benchmark the performance, we will use the timeit module and repeat each function 10000 times. We will take the mean with standard error = \(\sigma / \sqrt{N_\text{repeats}}\).

step2.py

import math
import timeit

import numpy as np

import common

INITIAL_CONDITION = "solar_system"
NUM_REPEATS = 10000


def main() -> None:
    # Get initial conditions
    system, _, _, _ = common.get_initial_conditions(INITIAL_CONDITION)

    ### Benchmark ###
    print("Benchmarking with 10000 repetitions")
    print()

    # Allocate memory
    a = np.zeros((system.num_particles, 3))

    # Acceleration 1
    run_time_1 = np.zeros(NUM_REPEATS)
    for i in range(NUM_REPEATS):
        start = timeit.default_timer()
        acceleration_1(a, system)
        end = timeit.default_timer()
        run_time_1[i] = end - start
    print(
        f"acceleration_1: {run_time_1.mean():.6f} +- {run_time_1.std(ddof=1) / math.sqrt(NUM_REPEATS):.3g} seconds"
    )


    ... # (Repeat for acceleration_2, 3, and 4)

Finally, we do a error check by comparing the results from the first naive implementation.

def main() -> None:
    ...
    # Check for relative errors
    ### Error check ###
    acceleration_1(a, system)
    a_1 = a.copy()
    acceleration_2(a, system)
    a_2 = a.copy()
    acceleration_3(a, system)
    a_3 = a.copy()
    acceleration_4(a, system)
    a_4 = a.copy()

    rel_error_2 = np.sum(np.abs(a_1 - a_2)) / np.sum(a_1)
    rel_error_3 = np.sum(np.abs(a_1 - a_3)) / np.sum(a_1)
    rel_error_4 = np.sum(np.abs(a_1 - a_4)) / np.sum(a_1)

    print()
    print("Error check: (relative difference from acceleration_1)")
    print(f"acceleration_2: {rel_error_2:.3g}")
    print(f"acceleration_3: {rel_error_3:.3g}")
    print(f"acceleration_4: {rel_error_4:.3g}")

The results are as follows:

Benchmarking with 10000 repetitions

acceleration_1: 0.000203 +- 8.08e-08 seconds
acceleration_2: 0.000164 +- 1.25e-06 seconds
acceleration_3: 0.000013 +- 2.21e-08 seconds
acceleration_4: 0.000012 +- 1.38e-08 seconds

Error check: (relative difference from acceleration_1)
acceleration_2: 0
acceleration_3: 1.31e-15
acceleration_4: 1.31e-15

The vectorized implementation is about 10 - 20 times faster than the naive implementation! As for the error check, the small relative difference is likely due to rounding errors, which could be ignored. (For 64-bit floating point numbers, the machine epsilon is about \(10^{-16}\).) By the way, since acceleration_4 is the fastest, we put it into common.py.

Performance in C

By the way, if you are interested in the performance in C, below is a benchmark using our grav_sim package:

Test 0:    Method: Pairwise
    Number of times: 10000000
    Avg time: 2.06e-07 (+- 1.30e-10) s

This is about 58 times faster than the vectorized NumPy implementation. But beware that this may not be totally accurate as the run time for each run is too short.

Full scripts

The full scripts are available at 5_steps_to_n_body_simulation/python/, or https://github.com/alvinng4/grav_sim/blob/main/5_steps_to_n_body_simulation/python/

step2.py (Click to expand)

import math
import timeit

import numpy as np

import common

INITIAL_CONDITION = "solar_system"
NUM_REPEATS = 10000


def main() -> None:
    # Get initial conditions
    system, _, _, _ = common.get_initial_conditions(INITIAL_CONDITION)

    ### Benchmark ###
    print("Benchmarking with 10000 repetitions")
    print()

    # Allocate memory
    a = np.zeros((system.num_particles, 3))

    # Acceleration 1
    run_time_1 = np.zeros(NUM_REPEATS)
    for i in range(NUM_REPEATS):
        start = timeit.default_timer()
        acceleration_1(a, system)
        end = timeit.default_timer()
        run_time_1[i] = end - start
    print(
        f"acceleration_1: {run_time_1.mean():.6f} +- {run_time_1.std(ddof=1) / math.sqrt(NUM_REPEATS):.3g} seconds"
    )

    # Acceleration 2
    run_time_2 = np.zeros(NUM_REPEATS)
    for i in range(NUM_REPEATS):
        start = timeit.default_timer()
        acceleration_2(a, system)
        end = timeit.default_timer()
        run_time_2[i] = end - start
    print(
        f"acceleration_2: {run_time_2.mean():.6f} +- {run_time_2.std(ddof=1) / math.sqrt(NUM_REPEATS):.3g} seconds"
    )

    # Acceleration 3
    run_time_3 = np.zeros(NUM_REPEATS)
    for i in range(NUM_REPEATS):
        start = timeit.default_timer()
        acceleration_3(a, system)
        end = timeit.default_timer()
        run_time_3[i] = end - start
    print(
        f"acceleration_3: {run_time_3.mean():.6f} +- {run_time_3.std(ddof=1) / math.sqrt(NUM_REPEATS):.3g} seconds"
    )

    # Acceleration 4
    run_time_4 = np.zeros(NUM_REPEATS)
    for i in range(NUM_REPEATS):
        start = timeit.default_timer()
        acceleration_4(a, system)
        end = timeit.default_timer()
        run_time_4[i] = end - start
    print(
        f"acceleration_4: {run_time_4.mean():.6f} +- {run_time_4.std(ddof=1) / math.sqrt(NUM_REPEATS):.3g} seconds"
    )

    ### Error check ###
    acceleration_1(a, system)
    a_1 = a.copy()
    acceleration_2(a, system)
    a_2 = a.copy()
    acceleration_3(a, system)
    a_3 = a.copy()
    acceleration_4(a, system)
    a_4 = a.copy()

    rel_error_2 = np.sum(np.abs(a_1 - a_2)) / np.sum(a_1)
    rel_error_3 = np.sum(np.abs(a_1 - a_3)) / np.sum(a_1)
    rel_error_4 = np.sum(np.abs(a_1 - a_4)) / np.sum(a_1)

    print()
    print("Error check: (relative difference from acceleration_1)")
    print(f"acceleration_2: {rel_error_2:.3g}")
    print(f"acceleration_3: {rel_error_3:.3g}")
    print(f"acceleration_4: {rel_error_4:.3g}")


def acceleration_1(
    a: np.ndarray,
    system: common.System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    num_particles = system.num_particles
    x = system.x
    m = system.m
    G = system.G

    # Calculations
    for i in range(num_particles):
        for j in range(num_particles):
            if i == j:
                continue

            R = x[j] - x[i]
            a[i] += G * m[j] * R / (np.linalg.norm(R) ** 3)


def acceleration_2(
    a: np.ndarray,
    system: common.System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    num_particles = system.num_particles
    x = system.x
    m = system.m
    G = system.G

    # Calculations
    for i in range(num_particles):
        for j in range(i + 1, num_particles):
            R = x[j] - x[i]
            temp_value = G * R / (np.linalg.norm(R) ** 3)
            a[i] += temp_value * m[j]
            a[j] -= temp_value * m[i]


def acceleration_3(
    a: np.ndarray,
    system: common.System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    x = system.x
    m = system.m
    G = system.G

    # Compute the displacement vector
    r_ij = x[:, np.newaxis, :] - x[np.newaxis, :, :]

    # Compute the distance
    r_norm = np.linalg.norm(r_ij, axis=2)

    # Compute 1 / r^3
    with np.errstate(divide="ignore", invalid="ignore"):
        inv_r_cubed = 1.0 / (r_norm * r_norm * r_norm)

    # Set diagonal elements to 0 to avoid self-interaction
    np.fill_diagonal(inv_r_cubed, 0.0)

    # Compute the acceleration
    a[:] = G * np.sum(
        r_ij * inv_r_cubed[:, :, np.newaxis] * m[:, np.newaxis, np.newaxis], axis=0
    )


def acceleration_4(
    a: np.ndarray,
    system: common.System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    x = system.x
    m = system.m
    G = system.G

    # Compute the displacement vector
    r_ij = x[:, np.newaxis, :] - x[np.newaxis, :, :]

    # Compute the distance
    r_norm = np.linalg.norm(r_ij, axis=2)

    # Compute 1 / r^3
    with np.errstate(divide="ignore", invalid="ignore"):
        inv_r_cubed = 1.0 / (r_norm * r_norm * r_norm)

    # Set diagonal elements to 0 to avoid self-interaction
    np.fill_diagonal(inv_r_cubed, 0.0)

    # Compute the acceleration
    a[:] = G * np.einsum("ijk,ij,i->jk", r_ij, inv_r_cubed, m)


if __name__ == "__main__":
    main()

common.py (Click to expand)

from typing import Tuple, List, Optional

import numpy as np
import matplotlib.pyplot as plt


##### Step 1 #####
class System:
    def __init__(
        self, num_particles: int, x: np.ndarray, v: np.ndarray, m: np.ndarray, G: float
    ) -> None:
        self.num_particles = num_particles
        self.x = x
        self.v = v
        self.m = m
        self.G = G

    def center_of_mass_correction(self) -> None:
        """Set center of mass of position and velocity to zero"""
        M = np.sum(self.m)
        x_cm = np.einsum("i,ij->j", self.m, self.x) / M
        v_cm = np.einsum("i,ij->j", self.m, self.v) / M

        self.x -= x_cm
        self.v -= v_cm


def get_initial_conditions(
    initial_condition: str,
) -> Tuple[System, List[Optional[str]], List[Optional[str]], bool]:
    """
    Returns the initial conditions for solar system,
    with units AU, days, and M_sun.

    Parameters
    ----------
    initial_condition : str
        Name for the initial condition.

    Returns
    -------
    system: System
        System object with initial conditions.
    labels: list
        Labels for the particles.
    colors: list
        Colors for the particles.
    legend: bool
        Whether to show the legend.
    """
    # Conversion factor from km^3 s^-2 to AU^3 d^-2
    CONVERSION_FACTOR = (86400**2) / (149597870.7**3)

    # GM values (km^3 s^-2)
    # ref: https://ssd.jpl.nasa.gov/doc/Park.2021.AJ.DE440.pdf
    GM_KM_S = {
        "Sun": 132712440041.279419,
        "Mercury": 22031.868551,
        "Venus": 324858.592000,
        "Earth": 398600.435507,
        "Mars": 42828.375816,
        "Jupiter": 126712764.100000,
        "Saturn": 37940584.841800,
        "Uranus": 5794556.400000,
        "Neptune": 6836527.100580,
        "Moon": 4902.800118,
        "Pluto": 975.500000,
        "Ceres": 62.62890,
        "Vesta": 17.288245,
    }

    # GM values (AU^3 d^-2)
    GM_AU_DAY = {
        "Sun": 132712440041.279419 * CONVERSION_FACTOR,
        "Mercury": 22031.868551 * CONVERSION_FACTOR,
        "Venus": 324858.592000 * CONVERSION_FACTOR,
        "Earth": 398600.435507 * CONVERSION_FACTOR,
        "Mars": 42828.375816 * CONVERSION_FACTOR,
        "Jupiter": 126712764.100000 * CONVERSION_FACTOR,
        "Saturn": 37940584.841800 * CONVERSION_FACTOR,
        "Uranus": 5794556.400000 * CONVERSION_FACTOR,
        "Neptune": 6836527.100580 * CONVERSION_FACTOR,
        "Moon": 4902.800118 * CONVERSION_FACTOR,
        "Pluto": 975.500000 * CONVERSION_FACTOR,
        "Ceres": 62.62890 * CONVERSION_FACTOR,
        "Vesta": 17.288245 * CONVERSION_FACTOR,
    }

    # Solar system masses (M_sun^-1)
    SOLAR_SYSTEM_MASSES = {
        "Sun": 1.0,
        "Mercury": GM_KM_S["Mercury"] / GM_KM_S["Sun"],
        "Venus": GM_KM_S["Venus"] / GM_KM_S["Sun"],
        "Earth": GM_KM_S["Earth"] / GM_KM_S["Sun"],
        "Mars": GM_KM_S["Mars"] / GM_KM_S["Sun"],
        "Jupiter": GM_KM_S["Jupiter"] / GM_KM_S["Sun"],
        "Saturn": GM_KM_S["Saturn"] / GM_KM_S["Sun"],
        "Uranus": GM_KM_S["Uranus"] / GM_KM_S["Sun"],
        "Neptune": GM_KM_S["Neptune"] / GM_KM_S["Sun"],
        "Moon": GM_KM_S["Moon"] / GM_KM_S["Sun"],
        "Pluto": GM_KM_S["Pluto"] / GM_KM_S["Sun"],
        "Ceres": GM_KM_S["Ceres"] / GM_KM_S["Sun"],
        "Vesta": GM_KM_S["Vesta"] / GM_KM_S["Sun"],
    }

    G = GM_AU_DAY["Sun"]

    # Solar system position and velocities data
    # Units: AU-D
    # Coordinate center: Solar System Barycenter
    # Data dated on A.D. 2024-Jan-01 00:00:00.0000 TDB
    # Computational data generated by NASA JPL Horizons System https://ssd.jpl.nasa.gov/horizons/
    SOLAR_SYSTEM_POS = {
        "Sun": [-7.967955691533730e-03, -2.906227441573178e-03, 2.103054301547123e-04],
        "Mercury": [
            -2.825983269538632e-01,
            1.974559795958082e-01,
            4.177433558063677e-02,
        ],
        "Venus": [
            -7.232103701666379e-01,
            -7.948302026312400e-02,
            4.042871428174315e-02,
        ],
        "Earth": [-1.738192017257054e-01, 9.663245550235138e-01, 1.553901854897183e-04],
        "Mars": [-3.013262392582653e-01, -1.454029331393295e00, -2.300531433991428e-02],
        "Jupiter": [3.485202469657674e00, 3.552136904413157e00, -9.271035442798399e-02],
        "Saturn": [8.988104223143450e00, -3.719064854634689e00, -2.931937777323593e-01],
        "Uranus": [1.226302417897505e01, 1.529738792480545e01, -1.020549026883563e-01],
        "Neptune": [
            2.983501460984741e01,
            -1.793812957956852e00,
            -6.506401132254588e-01,
        ],
        "Moon": [-1.762788124769829e-01, 9.674377513177153e-01, 3.236901585768862e-04],
        "Pluto": [1.720200478843485e01, -3.034155683573043e01, -1.729127607100611e00],
        "Ceres": [-1.103880510367569e00, -2.533340440444230e00, 1.220283937721780e-01],
        "Vesta": [-8.092549658731499e-02, 2.558381434460076e00, -6.695836142398572e-02],
    }
    SOLAR_SYSTEM_VEL = {
        "Sun": [4.875094764261564e-06, -7.057133213976680e-06, -4.573453713094512e-08],
        "Mercury": [
            -2.232165900189702e-02,
            -2.157207103176252e-02,
            2.855193410495743e-04,
        ],
        "Venus": [
            2.034068201002341e-03,
            -2.020828626592994e-02,
            -3.945639843855159e-04,
        ],
        "Earth": [
            -1.723001232538228e-02,
            -2.967721342618870e-03,
            6.382125383116755e-07,
        ],
        "Mars": [1.424832259345280e-02, -1.579236181580905e-03, -3.823722796161561e-04],
        "Jupiter": [
            -5.470970658852281e-03,
            5.642487338479145e-03,
            9.896190602066252e-05,
        ],
        "Saturn": [
            1.822013845554067e-03,
            5.143470425888054e-03,
            -1.617235904887937e-04,
        ],
        "Uranus": [
            -3.097615358317413e-03,
            2.276781932345769e-03,
            4.860433222241686e-05,
        ],
        "Neptune": [
            1.676536611817232e-04,
            3.152098732861913e-03,
            -6.877501095688201e-05,
        ],
        "Moon": [
            -1.746667306153906e-02,
            -3.473438277358121e-03,
            -3.359028758606074e-05,
        ],
        "Pluto": [2.802810313667557e-03, 8.492056438614633e-04, -9.060790113327894e-04],
        "Ceres": [
            8.978653480111301e-03,
            -4.873256528198994e-03,
            -1.807162046049230e-03,
        ],
        "Vesta": [
            -1.017876585480054e-02,
            -5.452367109338154e-04,
            1.255870551153315e-03,
        ],
    }

    SOLAR_SYSTEM_COLORS = {
        "Sun": "orange",
        "Mercury": "slategrey",
        "Venus": "wheat",
        "Earth": "skyblue",
        "Mars": "red",
        "Jupiter": "darkgoldenrod",
        "Saturn": "gold",
        "Uranus": "paleturquoise",
        "Neptune": "blue",
    }

    SOLAR_SYSTEM_PLUS_COLORS = {
        "Sun": "orange",
        "Mercury": "slategrey",
        "Venus": "wheat",
        "Earth": "skyblue",
        "Mars": "red",
        "Jupiter": "darkgoldenrod",
        "Saturn": "gold",
        "Uranus": "paleturquoise",
        "Neptune": "blue",
        "Pluto": None,
        "Ceres": None,
        "Vesta": None,
    }

    if initial_condition == "pyth-3-body":
        # Pythagorean 3-body problem
        R1 = np.array([1.0, 3.0, 0.0])
        R2 = np.array([-2.0, -1.0, 0.0])
        R3 = np.array([1.0, -1.0, 0.0])
        V1 = np.array([0.0, 0.0, 0.0])
        V2 = np.array([0.0, 0.0, 0.0])
        V3 = np.array([0.0, 0.0, 0.0])

        x = np.array([R1, R2, R3])
        v = np.array([V1, V2, V3])
        m = np.array([3.0 / G, 4.0 / G, 5.0 / G])

        system = System(
            num_particles=len(m),
            x=x,
            v=v,
            m=m,
            G=G,
        )
        system.center_of_mass_correction()

        labels: List[Optional[str]] = [None, None, None]
        colors: List[Optional[str]] = [None, None, None]
        legend = False

        return system, labels, colors, legend

    elif initial_condition == "solar_system":
        m = np.array(
            [
                SOLAR_SYSTEM_MASSES["Sun"],
                SOLAR_SYSTEM_MASSES["Mercury"],
                SOLAR_SYSTEM_MASSES["Venus"],
                SOLAR_SYSTEM_MASSES["Earth"],
                SOLAR_SYSTEM_MASSES["Mars"],
                SOLAR_SYSTEM_MASSES["Jupiter"],
                SOLAR_SYSTEM_MASSES["Saturn"],
                SOLAR_SYSTEM_MASSES["Uranus"],
                SOLAR_SYSTEM_MASSES["Neptune"],
            ]
        )

        R1 = np.array(SOLAR_SYSTEM_POS["Sun"])
        R2 = np.array(SOLAR_SYSTEM_POS["Mercury"])
        R3 = np.array(SOLAR_SYSTEM_POS["Venus"])
        R4 = np.array(SOLAR_SYSTEM_POS["Earth"])
        R5 = np.array(SOLAR_SYSTEM_POS["Mars"])
        R6 = np.array(SOLAR_SYSTEM_POS["Jupiter"])
        R7 = np.array(SOLAR_SYSTEM_POS["Saturn"])
        R8 = np.array(SOLAR_SYSTEM_POS["Uranus"])
        R9 = np.array(SOLAR_SYSTEM_POS["Neptune"])

        V1 = np.array(SOLAR_SYSTEM_VEL["Sun"])
        V2 = np.array(SOLAR_SYSTEM_VEL["Mercury"])
        V3 = np.array(SOLAR_SYSTEM_VEL["Venus"])
        V4 = np.array(SOLAR_SYSTEM_VEL["Earth"])
        V5 = np.array(SOLAR_SYSTEM_VEL["Mars"])
        V6 = np.array(SOLAR_SYSTEM_VEL["Jupiter"])
        V7 = np.array(SOLAR_SYSTEM_VEL["Saturn"])
        V8 = np.array(SOLAR_SYSTEM_VEL["Uranus"])
        V9 = np.array(SOLAR_SYSTEM_VEL["Neptune"])

        x = np.array(
            [
                R1,
                R2,
                R3,
                R4,
                R5,
                R6,
                R7,
                R8,
                R9,
            ]
        )
        v = np.array(
            [
                V1,
                V2,
                V3,
                V4,
                V5,
                V6,
                V7,
                V8,
                V9,
            ]
        )

        system = System(
            num_particles=len(m),
            x=x,
            v=v,
            m=m,
            G=G,
        )
        system.center_of_mass_correction()

        labels = list(SOLAR_SYSTEM_COLORS.keys())
        colors = list(SOLAR_SYSTEM_COLORS.values())
        legend = True

        return system, labels, colors, legend

    elif initial_condition == "solar_system_plus":
        m = np.array(
            [
                SOLAR_SYSTEM_MASSES["Sun"],
                SOLAR_SYSTEM_MASSES["Mercury"],
                SOLAR_SYSTEM_MASSES["Venus"],
                SOLAR_SYSTEM_MASSES["Earth"],
                SOLAR_SYSTEM_MASSES["Mars"],
                SOLAR_SYSTEM_MASSES["Jupiter"],
                SOLAR_SYSTEM_MASSES["Saturn"],
                SOLAR_SYSTEM_MASSES["Uranus"],
                SOLAR_SYSTEM_MASSES["Neptune"],
                SOLAR_SYSTEM_MASSES["Pluto"],
                SOLAR_SYSTEM_MASSES["Ceres"],
                SOLAR_SYSTEM_MASSES["Vesta"],
            ]
        )

        R1 = np.array(SOLAR_SYSTEM_POS["Sun"])
        R2 = np.array(SOLAR_SYSTEM_POS["Mercury"])
        R3 = np.array(SOLAR_SYSTEM_POS["Venus"])
        R4 = np.array(SOLAR_SYSTEM_POS["Earth"])
        R5 = np.array(SOLAR_SYSTEM_POS["Mars"])
        R6 = np.array(SOLAR_SYSTEM_POS["Jupiter"])
        R7 = np.array(SOLAR_SYSTEM_POS["Saturn"])
        R8 = np.array(SOLAR_SYSTEM_POS["Uranus"])
        R9 = np.array(SOLAR_SYSTEM_POS["Neptune"])
        R10 = np.array(SOLAR_SYSTEM_POS["Pluto"])
        R11 = np.array(SOLAR_SYSTEM_POS["Ceres"])
        R12 = np.array(SOLAR_SYSTEM_POS["Vesta"])

        V1 = np.array(SOLAR_SYSTEM_VEL["Sun"])
        V2 = np.array(SOLAR_SYSTEM_VEL["Mercury"])
        V3 = np.array(SOLAR_SYSTEM_VEL["Venus"])
        V4 = np.array(SOLAR_SYSTEM_VEL["Earth"])
        V5 = np.array(SOLAR_SYSTEM_VEL["Mars"])
        V6 = np.array(SOLAR_SYSTEM_VEL["Jupiter"])
        V7 = np.array(SOLAR_SYSTEM_VEL["Saturn"])
        V8 = np.array(SOLAR_SYSTEM_VEL["Uranus"])
        V9 = np.array(SOLAR_SYSTEM_VEL["Neptune"])
        V10 = np.array(SOLAR_SYSTEM_VEL["Pluto"])
        V11 = np.array(SOLAR_SYSTEM_VEL["Ceres"])
        V12 = np.array(SOLAR_SYSTEM_VEL["Vesta"])

        x = np.array(
            [
                R1,
                R2,
                R3,
                R4,
                R5,
                R6,
                R7,
                R8,
                R9,
                R10,
                R11,
                R12,
            ]
        )
        v = np.array(
            [
                V1,
                V2,
                V3,
                V4,
                V5,
                V6,
                V7,
                V8,
                V9,
                V10,
                V11,
                V12,
            ]
        )

        system = System(
            num_particles=len(m),
            x=x,
            v=v,
            m=m,
            G=G,
        )
        system.center_of_mass_correction()

        labels = list(SOLAR_SYSTEM_PLUS_COLORS.keys())
        colors = list(SOLAR_SYSTEM_PLUS_COLORS.values())
        legend = True

        return system, labels, colors, legend

    else:
        raise ValueError(f"Initial condition not recognized: {initial_condition}.")


def plot_initial_conditions(
    system: System,
    labels: list,
    colors: list,
    legend: bool,
) -> None:
    """
    Plot the initial positions.

    Parameters
    ----------
    system : System
        System object.
    labels : list
        Labels for the particles.
    colors : list
        Colors for the particles.
    legend : bool
        Whether to show the legend.
    """
    fig, ax = plt.subplots()
    ax.set_xlabel("$x$ (AU)")
    ax.set_ylabel("$y$ (AU)")

    for i in range(system.num_particles):
        ax.scatter(
            system.x[i, 0], system.x[i, 1], marker="o", color=colors[i], label=labels[i]
        )

    if legend:
        ax.legend()

    plt.show()


##### Step 2 #####
def acceleration(
    a: np.ndarray,
    system: System,
) -> None:
    """
    Compute the gravitational acceleration

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array to be modified in-place,
        with shape (N, 3)
    system : System
        System object.
    """
    # Empty acceleration array
    a.fill(0.0)

    # Declare variables
    x = system.x
    m = system.m
    G = system.G

    # Compute the displacement vector
    r_ij = x[:, np.newaxis, :] - x[np.newaxis, :, :]

    # Compute the distance
    r_norm = np.linalg.norm(r_ij, axis=2)

    # Compute 1 / r^3
    with np.errstate(divide="ignore", invalid="ignore"):
        inv_r_cubed = 1.0 / (r_norm * r_norm * r_norm)

    # Set diagonal elements to 0 to avoid self-interaction
    np.fill_diagonal(inv_r_cubed, 0.0)

    # Compute the acceleration
    a[:] = G * np.einsum("ijk,ij,i->jk", r_ij, inv_r_cubed, m)


##### Step 3 #####
def euler(a: np.ndarray, system: System, dt: float) -> None:
    """
    Advance one step with the Euler's method.

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array with shape (N, 3).
    system : System
        System object.
    dt : float
        Time step.
    """
    acceleration(a, system)
    system.x += system.v * dt
    system.v += a * dt


def print_simulation_info_fixed_step_size(
    system: System,
    tf: float,
    dt: float,
    num_steps: int,
    output_interval: float,
    sol_size: int,
) -> None:
    print("----------------------------------------------------------")
    print("Simulation Info:")
    print(f"num_particles: {system.num_particles}")
    print(f"G: {system.G}")
    print(f"tf: {tf} days (Actual tf = dt * num_steps = {dt * num_steps} days)")
    print(f"dt: {dt} days")
    print(f"Num_steps: {num_steps}")
    print()
    print(f"Output interval: {output_interval} days")
    print(f"Estimated solution size: {sol_size}")
    print("----------------------------------------------------------")


def plot_trajectory(
    sol_x: np.ndarray,
    labels: list,
    colors: list,
    legend: bool,
) -> None:
    """
    Plot the 2D trajectory.

    Parameters
    ----------
    sol_x : np.ndarray
        Solution position array with shape (N_steps, num_particles, 3).
    labels : list
        List of labels for the particles.
    colors : list
        List of colors for the particles.
    legend : bool
        Whether to show the legend.
    """
    fig = plt.figure()
    ax = fig.add_subplot(111, aspect="equal")
    ax.set_xlabel("$x$ (AU)")
    ax.set_ylabel("$y$ (AU)")

    for i in range(sol_x.shape[1]):
        traj = ax.plot(
            sol_x[:, i, 0],
            sol_x[:, i, 1],
            color=colors[i],
        )
        # Plot the last position with marker
        ax.scatter(
            sol_x[-1, i, 0],
            sol_x[-1, i, 1],
            marker="o",
            color=traj[0].get_color(),
            label=labels[i],
        )

    if legend:
        fig.legend(loc="center right", borderaxespad=0.2)
        fig.tight_layout()

    plt.show()


##### Step 4 #####
def compute_rel_energy_error(
    sol_x: np.ndarray, sol_v: np.ndarray, system: System
) -> np.ndarray:
    """
    Compute the relative energy error of the simulation.

    Parameters
    ----------
    sol_x : np.ndarray
        Solution position array with shape (N_steps, num_particles, 3).
    sol_v : np.ndarray
        Solution velocity array with shape (N_steps, num_particles, 3).
    system : System
        System object.

    Returns
    -------
    energy_error : np.ndarray
        Relative energy error of the simulation, with shape (N_steps,).
    """
    # Allocate memory and initialize arrays
    n_steps = sol_x.shape[0]
    num_particles = system.num_particles
    m = system.m
    G = system.G
    rel_energy_error = np.zeros(n_steps)

    # Compute the total energy (KE + PE)
    for count in range(n_steps):
        x = sol_x[count]
        v = sol_v[count]
        for i in range(num_particles):
            # KE
            rel_energy_error[count] += 0.5 * m[i] * np.linalg.norm(v[i]) ** 2
            # PE
            for j in range(i + 1, num_particles):
                rel_energy_error[count] -= G * m[i] * m[j] / np.linalg.norm(x[i] - x[j])

    # Compute the relative energy error
    initial_energy = rel_energy_error[0]
    rel_energy_error = (rel_energy_error - initial_energy) / initial_energy
    rel_energy_error = np.abs(rel_energy_error)

    return rel_energy_error


def plot_rel_energy_error(rel_energy_error: np.ndarray, sol_t: np.ndarray) -> None:
    """
    Plot the relative energy error.

    Parameters
    ----------
    rel_energy_error : np.ndarray
        Relative energy error of the simulation, with shape (N_steps,).
    sol_t : np.ndarray
        Solution time array with shape (N_steps,).
    """
    plt.figure()
    plt.plot(sol_t, rel_energy_error)
    plt.yscale("log")
    plt.xlabel("Time step")
    plt.ylabel("Relative Energy Error")
    plt.title("Relative Energy Error vs Time Step")
    plt.show()


def euler_cromer(a: np.ndarray, system: System, dt: float) -> None:
    """
    Advance one step with the Euler-Cromer method.

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array with shape (N, 3).
    system : System
        System object.
    dt : float
        Time step.
    """
    acceleration(a, system)
    system.v += a * dt
    system.x += system.v * dt


def rk4(a: np.ndarray, system: System, dt: float) -> None:
    """
    Advance one step with the RK4 method.

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array with shape (N, 3).
    system : System
        System object.
    dt : float
        Time step.
    """
    num_stages = 4
    coeff = np.array([0.5, 0.5, 1.0])
    weights = np.array([1.0, 2.0, 2.0, 1.0]) / 6.0

    # Allocate memory and initialize arrays
    x0 = system.x.copy()
    v0 = system.v.copy()
    xk = np.zeros((num_stages, system.num_particles, 3))
    vk = np.zeros((num_stages, system.num_particles, 3))

    # Initial stage
    acceleration(a, system)
    xk[0] = v0
    vk[0] = a

    # Compute the stages
    for stage in range(1, num_stages):
        # Compute acceleration
        system.x = x0 + dt * coeff[stage - 1] * xk[stage - 1]
        acceleration(a, system)

        # Compute xk and vk
        xk[stage] = v0 + dt * coeff[stage - 1] * vk[stage - 1]
        vk[stage] = a

    # Advance step
    # dx = 0.0
    # dv = 0.0
    # for stage in range(num_stages):
    #     dx += weights[stage] * xk[stage]
    #     dv += weights[stage] * vk[stage]

    dx = np.einsum("i,ijk->jk", weights, xk)
    dv = np.einsum("i,ijk->jk", weights, vk)

    system.x = x0 + dt * dx
    system.v = v0 + dt * dv


def leapfrog(a: np.ndarray, system: System, dt: float) -> None:
    """
    Advance one step with the LeapFrog method.

    Parameters
    ----------
    a : np.ndarray
        Gravitational accelerations array with shape (N, 3).
    system : System
        System object.
    dt : float
        Time step.
    """
    # Velocity kick (v_1/2)
    acceleration(a, system)
    system.v += a * 0.5 * dt

    # Position drift (x_1)
    system.x += system.v * dt

    # Velocity kick (v_1)
    acceleration(a, system)
    system.v += a * 0.5 * dt


##### Step 5 #####
def print_simulation_info_adaptive_step_size(
    system: System,
    tf: float,
    tolerance: float,
    initial_dt: float,
    output_interval: float,
    sol_size: int,
) -> None:
    print("----------------------------------------------------------")
    print("Simulation Info:")
    print(f"num_particles: {system.num_particles}")
    print(f"G: {system.G}")
    print(f"tf: {tf} days")
    print(f"tolerance: {tolerance}")
    print(f"Initial dt: {initial_dt} days")
    print()
    print(f"Output interval: {output_interval} days")
    print(f"Estimated solution size: {sol_size}")
    print("----------------------------------------------------------")


def plot_dt(sol_dt: np.ndarray, sol_t: np.ndarray) -> None:
    """
    Plot the time step.

    Parameters
    ----------
    sol_dt : np.ndarray
        Time step array with shape (N_steps,).
    sol_t : np.ndarray
        Solution time array with shape (N_steps,).
    """
    plt.figure()
    plt.semilogy(sol_t, sol_dt)
    plt.xlabel("Time")
    plt.ylabel("dt")
    plt.show()


##### Extra #####
def set_3d_axes_equal(ax: plt.Axes) -> None:
    """
    Make axes of 3D plot have equal scale

    Parameters
    ----------
    ax : matplotlib axis
        The axis to set equal scale

    Reference
    ---------
    karlo, https://stackoverflow.com/questions/13685386/how-to-set-the-equal-aspect-ratio-for-all-axes-x-y-z
    """

    x_limits = ax.get_xlim3d()  # type: ignore
    y_limits = ax.get_ylim3d()  # type: ignore
    z_limits = ax.get_zlim3d()  # type: ignore

    x_range = abs(x_limits[1] - x_limits[0])
    x_middle = np.mean(x_limits)
    y_range = abs(y_limits[1] - y_limits[0])
    y_middle = np.mean(y_limits)
    z_range = abs(z_limits[1] - z_limits[0])
    z_middle = np.mean(z_limits)

    # The plot bounding box is a sphere in the sense of the infinity
    # norm, hence I call half the max range the plot radius.
    plot_radius = 0.5 * max([x_range, y_range, z_range])

    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])  # type: ignore
    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])  # type: ignore
    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])  # type: ignore


def plot_3d_trajectory(
    sol_x: np.ndarray,
    labels: list,
    colors: list,
    legend: bool,
) -> None:
    """
    Plot the 3D trajectory.

    Parameters
    ----------
    sol_x : np.ndarray
        Solution position array with shape (N_steps, num_particles, 3).
    labels : list
        List of labels for the particles.
    colors : list
        List of colors for the particles.
    legend : bool
        Whether to show the legend.
    """

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    ax.set_xlabel("$x$ (AU)")
    ax.set_ylabel("$y$ (AU)")
    ax.set_zlabel("$z$ (AU)")  # type: ignore

    for i in range(sol_x.shape[1]):
        traj = ax.plot(
            sol_x[:, i, 0],
            sol_x[:, i, 1],
            sol_x[:, i, 2],
            color=colors[i],
        )
        # Plot the last position with marker
        ax.scatter(
            sol_x[-1, i, 0],
            sol_x[-1, i, 1],
            sol_x[-1, i, 2],
            marker="o",
            color=traj[0].get_color(),
            label=labels[i],
        )

    set_3d_axes_equal(ax)

    if legend:
        ax.legend(loc="center right", bbox_to_anchor=(1.325, 0.5))
        fig.subplots_adjust(right=0.7)

    plt.show()

Comments