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)