SpacesExamples/shiny-with-python

1from typing import Any2import numpy as np3import astropy.constants as c4import time5 6# Adapted from Python for Astronomers: An Introduction to Scientific Computing7# by Imad Pasha & Christopher Agostino8# https://prappleizer.github.io/Tutorials/RK4/RK4_Tutorial.html9 10# Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License11# http://creativecommons.org/licenses/by-nc-sa/4.0/12 13 14class Body:15    def __init__(self, mass, x_vec, v_vec, name=None, has_units=True):16        """17        spawn instance of the Body class, which is used in Simulations.18 19        :param: mass | mass of particle. if has_units=True, an Astropy Quantity, otherwise a float20        :param: x_vec | a vector len(3) containing the x, y, z initial positions of the body.21                     the array can be unitless if has_units=False, or be of the form np.array([0,0,0])*u.km22        :param: v_vec | vector len(3) containing the v_x, v_y, v_z initial velocities of the body.23        :param: name | string containing a name, used for plotting later24        :param: has_units | defines how the code treats the problem, as unit-ed, or unitless.25        """26        self.name = name27        self.has_units = has_units28        if self.has_units:29            self.mass = mass.cgs30            self.x_vec = x_vec.cgs.value31            self.v_vec = v_vec.cgs.value32        else:33            self.mass = mass34            self.x_vec = x_vec35            self.v_vec = v_vec36 37    def return_vec(self):38        """39        Concatenates the x and v vector into 1 vector 'y' used in RK formalism.40        """41        return np.concatenate((self.x_vec, self.v_vec))42 43    def return_mass(self):44        """45        handler to strip the mass units if present (after converting to cgs) or return float46        """47        if self.has_units:48            return self.mass.cgs.value49        else:50            return self.mass51 52    def return_name(self):53        return self.name54 55 56class Simulation:57    def __init__(self, bodies, has_units=True):58        """59        Initializes instance of Simulation object.60        -------------------------------------------61        Params:62            bodies (list): a list of Body() objects63            has_units (bool): set whether bodies entered have units or not.64        """65        self.has_units = has_units66        self.bodies = bodies67        self.N_bodies = len(self.bodies)68        self.nDim = 6.069        self.quant_vec = np.concatenate(np.array([i.return_vec() for i in self.bodies]))70        self.mass_vec = np.array([i.return_mass() for i in self.bodies])71        self.name_vec = [i.return_name() for i in self.bodies]72 73    def set_diff_eq(self, calc_diff_eqs, **kwargs):74        """75        Method which assigns an external solver function as the diff-eq solver for RK4.76        For N-body or gravitational setups, this is the function which calculates accelerations.77        ---------------------------------78        Params:79            calc_diff_eqs: A function which returns a [y] vector for RK480            **kwargs: Any additional inputs/hyperparameters the external function requires81        """82        self.diff_eq_kwargs = kwargs83        self.calc_diff_eqs = calc_diff_eqs84 85    def rk4(self, t, dt):86        """87        RK4 integrator. Calculates the K values and returns a new y vector88        --------------------------------89        Params:90            t: a time. Only used if the diff eq depends on time (gravity doesn't).91            dt: timestep. Non adaptive in this case92        """93        k1 = dt * self.calc_diff_eqs(94            t, self.quant_vec, self.mass_vec, **self.diff_eq_kwargs95        )96        k2 = dt * self.calc_diff_eqs(97            t + 0.5 * dt,98            self.quant_vec + 0.5 * k1,99            self.mass_vec,100            **self.diff_eq_kwargs,101        )102        k3 = dt * self.calc_diff_eqs(103            t + 0.5 * dt,104            self.quant_vec + 0.5 * k2,105            self.mass_vec,106            **self.diff_eq_kwargs,107        )108        k4 = dt * self.calc_diff_eqs(109            t + dt, self.quant_vec + k2, self.mass_vec, **self.diff_eq_kwargs110        )111 112        y_new = self.quant_vec + ((k1 + 2 * k2 + 2 * k3 + k4) / 6.0)113 114        return y_new115 116    def run(self, T, dt, t0=0, progress=None):117        """118        Method which runs the simulation on a given set of bodies.119        ---------------------120        Params:121            T: total time (in simulation units) to run the simulation. Can have units or not, just set has_units appropriately.122            dt: timestep (in simulation units) to advance the simulation. Same as above123            t0 (optional): set a non-zero start time to the simulation.124            progress (optional): A shiny.ui.Progress object which will be used to send progress updates.125 126        Returns:127            None, but leaves an attribute history accessed via128            'simulation.history' which contains all y vectors for the simulation.129            These are of shape (Nstep,Nbodies * 6), so the x and y positions of particle 1 are130            simulation.history[:,0], simulation.history[:,1], while the same for particle 2 are131            simulation.history[:,6], simulation.history[:,7]. Velocities are also extractable.132        """133        if not hasattr(self, "calc_diff_eqs"):134            raise AttributeError("You must set a diff eq solver first.")135        if self.has_units:136            try:137                _ = t0.unit138            except:139                t0 = (t0 * T.unit).cgs.value140            T = T.cgs.value141            dt = dt.cgs.value142 143        self.history: Any = [self.quant_vec]144        clock_time = t0145        nsteps = int((T - t0) / dt)146        start_time = time.time()147        for step in range(nsteps):148            if progress is not None and step % 5 == 0:149                progress.set(150                    step,151                    message=f"Integrating step = {step} / {nsteps}",152                    detail=f"Elapsed time = {round(clock_time/1e6, 1)}",153                )154            y_new = self.rk4(0, dt)155            self.history.append(y_new)156            self.quant_vec = y_new157            clock_time += dt158        runtime = time.time() - start_time159        self.history = np.array(self.history)160 161 162def nbody_solve(t, y, masses):163    N_bodies = int(len(y) / 6)164    solved_vector = np.zeros(y.size)165    for i in range(N_bodies):166        ioffset = i * 6167        for j in range(N_bodies):168            joffset = j * 6169            solved_vector[ioffset] = y[ioffset + 3]170            solved_vector[ioffset + 1] = y[ioffset + 4]171            solved_vector[ioffset + 2] = y[ioffset + 5]172            if i != j:173                dx = y[ioffset] - y[joffset]174                dy = y[ioffset + 1] - y[joffset + 1]175                dz = y[ioffset + 2] - y[joffset + 2]176                r = (dx**2 + dy**2 + dz**2) ** 0.5177                ax = (-c.G.cgs * masses[j] / r**3) * dx178                ay = (-c.G.cgs * masses[j] / r**3) * dy179                az = (-c.G.cgs * masses[j] / r**3) * dz180                ax = ax.value181                ay = ay.value182                az = az.value183                solved_vector[ioffset + 3] += ax184                solved_vector[ioffset + 4] += ay185                solved_vector[ioffset + 5] += az186    return solved_vector187 188 189def spherical_to_cartesian(190    theta: float, phi: float, rho: float191) -> tuple[float, float, float]:192    x = rho * sind(phi) * cosd(theta)193    y = rho * sind(phi) * sind(theta)194    z = rho * cosd(phi)195    return (x, y, z)196 197 198def cosd(x):199    return np.cos(x / 180 * np.pi)200 201 202def sind(x):203    return np.sin(x / 180 * np.pi)204