simulation.py

from typing import Any
import numpy as np
import astropy.constants as c
import time

# Adapted from Python for Astronomers: An Introduction to Scientific Computing
# by Imad Pasha & Christopher Agostino
# https://prappleizer.github.io/Tutorials/RK4/RK4_Tutorial.html

# Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License
# http://creativecommons.org/licenses/by-nc-sa/4.0/


class Body:
    def __init__(self, mass, x_vec, v_vec, name=None, has_units=True):
        """
        spawn instance of the Body class, which is used in Simulations.

        :param: mass | mass of particle. if has_units=True, an Astropy Quantity, otherwise a float
        :param: x_vec | a vector len(3) containing the x, y, z initial positions of the body.
                     the array can be unitless if has_units=False, or be of the form np.array([0,0,0])*u.km
        :param: v_vec | vector len(3) containing the v_x, v_y, v_z initial velocities of the body.
        :param: name | string containing a name, used for plotting later
        :param: has_units | defines how the code treats the problem, as unit-ed, or unitless.
        """
        self.name = name
        self.has_units = has_units
        if self.has_units:
            self.mass = mass.cgs
            self.x_vec = x_vec.cgs.value
            self.v_vec = v_vec.cgs.value
        else:
            self.mass = mass
            self.x_vec = x_vec
            self.v_vec = v_vec

    def return_vec(self):
        """
        Concatenates the x and v vector into 1 vector 'y' used in RK formalism.
        """
        return np.concatenate((self.x_vec, self.v_vec))

    def return_mass(self):
        """
        handler to strip the mass units if present (after converting to cgs) or return float
        """
        if self.has_units:
            return self.mass.cgs.value
        else:
            return self.mass

    def return_name(self):
        return self.name


class Simulation:
    def __init__(self, bodies, has_units=True):
        """
        Initializes instance of Simulation object.
        -------------------------------------------
        Params:
            bodies (list): a list of Body() objects
            has_units (bool): set whether bodies entered have units or not.
        """
        self.has_units = has_units
        self.bodies = bodies
        self.N_bodies = len(self.bodies)
        self.nDim = 6.0
        self.quant_vec = np.concatenate(np.array([i.return_vec() for i in self.bodies]))
        self.mass_vec = np.array([i.return_mass() for i in self.bodies])
        self.name_vec = [i.return_name() for i in self.bodies]

    def set_diff_eq(self, calc_diff_eqs, **kwargs):
        """
        Method which assigns an external solver function as the diff-eq solver for RK4.
        For N-body or gravitational setups, this is the function which calculates accelerations.
        ---------------------------------
        Params:
            calc_diff_eqs: A function which returns a [y] vector for RK4
            **kwargs: Any additional inputs/hyperparameters the external function requires
        """
        self.diff_eq_kwargs = kwargs
        self.calc_diff_eqs = calc_diff_eqs

    def rk4(self, t, dt):
        """
        RK4 integrator. Calculates the K values and returns a new y vector
        --------------------------------
        Params:
            t: a time. Only used if the diff eq depends on time (gravity doesn't).
            dt: timestep. Non adaptive in this case
        """
        k1 = dt * self.calc_diff_eqs(
            t, self.quant_vec, self.mass_vec, **self.diff_eq_kwargs
        )
        k2 = dt * self.calc_diff_eqs(
            t + 0.5 * dt,
            self.quant_vec + 0.5 * k1,
            self.mass_vec,
            **self.diff_eq_kwargs,
        )
        k3 = dt * self.calc_diff_eqs(
            t + 0.5 * dt,
            self.quant_vec + 0.5 * k2,
            self.mass_vec,
            **self.diff_eq_kwargs,
        )
        k4 = dt * self.calc_diff_eqs(
            t + dt, self.quant_vec + k2, self.mass_vec, **self.diff_eq_kwargs
        )

        y_new = self.quant_vec + ((k1 + 2 * k2 + 2 * k3 + k4) / 6.0)

        return y_new

    def run(self, T, dt, t0=0, progress=None):
        """
        Method which runs the simulation on a given set of bodies.
        ---------------------
        Params:
            T: total time (in simulation units) to run the simulation. Can have units or not, just set has_units appropriately.
            dt: timestep (in simulation units) to advance the simulation. Same as above
            t0 (optional): set a non-zero start time to the simulation.
            progress (optional): A shiny.ui.Progress object which will be used to send progress updates.

        Returns:
            None, but leaves an attribute history accessed via
            'simulation.history' which contains all y vectors for the simulation.
            These are of shape (Nstep,Nbodies * 6), so the x and y positions of particle 1 are
            simulation.history[:,0], simulation.history[:,1], while the same for particle 2 are
            simulation.history[:,6], simulation.history[:,7]. Velocities are also extractable.
        """
        if not hasattr(self, "calc_diff_eqs"):
            raise AttributeError("You must set a diff eq solver first.")
        if self.has_units:
            try:
                _ = t0.unit
            except:
                t0 = (t0 * T.unit).cgs.value
            T = T.cgs.value
            dt = dt.cgs.value

        self.history: Any = [self.quant_vec]
        clock_time = t0
        nsteps = int((T - t0) / dt)
        start_time = time.time()
        for step in range(nsteps):
            if progress is not None and step % 5 == 0:
                progress.set(
                    step,
                    message=f"Integrating step = {step} / {nsteps}",
                    detail=f"Elapsed time = {round(clock_time/1e6, 1)}",
                )
            y_new = self.rk4(0, dt)
            self.history.append(y_new)
            self.quant_vec = y_new
            clock_time += dt
        runtime = time.time() - start_time
        self.history = np.array(self.history)


def nbody_solve(t, y, masses):
    N_bodies = int(len(y) / 6)
    solved_vector = np.zeros(y.size)
    for i in range(N_bodies):
        ioffset = i * 6
        for j in range(N_bodies):
            joffset = j * 6
            solved_vector[ioffset] = y[ioffset + 3]
            solved_vector[ioffset + 1] = y[ioffset + 4]
            solved_vector[ioffset + 2] = y[ioffset + 5]
            if i != j:
                dx = y[ioffset] - y[joffset]
                dy = y[ioffset + 1] - y[joffset + 1]
                dz = y[ioffset + 2] - y[joffset + 2]
                r = (dx**2 + dy**2 + dz**2) ** 0.5
                ax = (-c.G.cgs * masses[j] / r**3) * dx
                ay = (-c.G.cgs * masses[j] / r**3) * dy
                az = (-c.G.cgs * masses[j] / r**3) * dz
                ax = ax.value
                ay = ay.value
                az = az.value
                solved_vector[ioffset + 3] += ax
                solved_vector[ioffset + 4] += ay
                solved_vector[ioffset + 5] += az
    return solved_vector


def spherical_to_cartesian(
    theta: float, phi: float, rho: float
) -> tuple[float, float, float]:
    x = rho * sind(phi) * cosd(theta)
    y = rho * sind(phi) * sind(theta)
    z = rho * cosd(phi)
    return (x, y, z)


def cosd(x):
    return np.cos(x / 180 * np.pi)


def sind(x):
    return np.sin(x / 180 * np.pi)