Duplicate from elonmuskceo/shiny-orbit-simulation

Co-authored-by: Elon Musk <elonmuskceo@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,22 @@
+FROM ubuntu:kinetic
+
+# Doesn't usually have an "upgrade"
+RUN apt-get update \
+ && DEBIAN_FRONTEND=noninteractive \
+    apt-get install --no-install-recommends --assume-yes \
+	  build-essential \
+      python3 \
+	  python3-dev \
+      python3-pip
+
+COPY requirements.txt .
+
+RUN pip install -r requirements.txt
+
+COPY . .
+
+ENTRYPOINT ["/bin/sh", "-c"]
+
+EXPOSE 7860
+
+CMD ["shiny run --port 7860 --host 0.0.0.0 app.py"]

README.md

@@ -0,0 +1,11 @@
+---
+title: Shiny Orbit Simulation
+emoji: 🪐
+colorFrom: green
+colorTo: green
+sdk: docker
+pinned: false
+duplicated_from: elonmuskceo/shiny-orbit-simulation
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,229 @@
+from pathlib import Path
+from simulation import Body, Simulation, nbody_solve, spherical_to_cartesian
+import matplotlib.pyplot as plt
+import astropy.units as u
+import numpy as np
+
+from shiny import App, reactive, render, ui
+
+# This application adapted from RK4 Orbit Integrator tutorial in Python for Astronomers
+# https://prappleizer.github.io/
+
+
+def panel_box(*args, **kwargs):
+    return ui.div(
+        ui.div(*args, class_="card-body"),
+        **kwargs,
+        class_="card mb-3",
+    )
+
+
+app_ui = ui.page_fluid(
+    {"class": "p-4"},
+    ui.row(
+        ui.column(
+            4,
+            panel_box(
+                ui.input_slider("days", "Simulation duration (days)", 0, 200, value=60),
+                ui.input_slider(
+                    "step_size",
+                    "Simulation time step (hours)",
+                    0,
+                    24,
+                    value=4,
+                    step=0.5,
+                ),
+                ui.input_action_button(
+                    "run", "Run simulation", class_="btn-primary w-100"
+                ),
+            ),
+            ui.navset_tab_card(
+                ui.nav(
+                    "Earth",
+                    ui.input_checkbox("earth", "Enable", True),
+                    ui.panel_conditional(
+                        "input.earth",
+                        ui.input_numeric(
+                            "earth_mass",
+                            "Mass (10^22 kg)",
+                            597.216,
+                        ),
+                        ui.input_slider(
+                            "earth_speed",
+                            "Speed (km/s)",
+                            0,
+                            1,
+                            value=0.0126,
+                            step=0.001,
+                        ),
+                        ui.input_slider("earth_theta", "Angle (5)", 0, 360, value=270),
+                        ui.input_slider("earth_phi", "5", 0, 180, value=90),
+                    ),
+                ),
+                ui.nav(
+                    "Moon",
+                    ui.input_checkbox("moon", "Enable", True),
+                    ui.panel_conditional(
+                        "input.moon",
+                        ui.input_numeric("moon_mass", "Mass (10^22 kg)", 7.347),
+                        ui.input_slider(
+                            "moon_speed", "Speed (km/s)", 0, 2, value=1.022, step=0.001
+                        ),
+                        ui.input_slider("moon_theta", "Angle (5)", 0, 360, value=90),
+                        ui.input_slider("moon_phi", "5", 0, 180, value=90),
+                    ),
+                ),
+                ui.nav(
+                    "Planet X",
+                    ui.input_checkbox("planetx", "Enable", False),
+                    ui.output_ui("planetx_controls"),
+                    ui.panel_conditional(
+                        "input.planetx",
+                        ui.input_numeric("planetx_mass", "Mass (10^22 kg)", 7.347),
+                        ui.input_slider(
+                            "planetx_speed",
+                            "Speed (km/s)",
+                            0,
+                            2,
+                            value=1.022,
+                            step=0.001,
+                        ),
+                        ui.input_slider("planetx_theta", "Angle (5)", 0, 360, 270),
+                        ui.input_slider("planetx_phi", "5", 0, 180, 90),
+                    ),
+                ),
+            ),
+        ),
+        ui.column(
+            8,
+            ui.output_plot("orbits", width="500px", height="500px"),
+            ui.img(src="coords.png", style="width: 100%; max-width: 250px;"),
+        ),
+    ),
+)
+
+
+def server(input, output, session):
+    def earth_body():
+        v = spherical_to_cartesian(
+            input.earth_theta(), input.earth_phi(), input.earth_speed()
+        )
+
+        return Body(
+            mass=input.earth_mass() * 10e21 * u.kg,
+            x_vec=np.array([0, 0, 0]) * u.km,
+            v_vec=np.array(v) * u.km / u.s,
+            name="Earth",
+        )
+
+    def moon_body():
+        v = spherical_to_cartesian(
+            input.moon_theta(), input.moon_phi(), input.moon_speed()
+        )
+
+        return Body(
+            mass=input.moon_mass() * 10e21 * u.kg,
+            x_vec=np.array([3.84e5, 0, 0]) * u.km,
+            v_vec=np.array(v) * u.km / u.s,
+            name="Moon",
+        )
+
+    def planetx_body():
+        v = spherical_to_cartesian(
+            input.planetx_theta(), input.planetx_phi(), input.planetx_speed()
+        )
+
+        return Body(
+            mass=input.planetx_mass() * 10e21 * u.kg,
+            x_vec=np.array([-3.84e5, 0, 0]) * u.km,
+            v_vec=np.array(v) * u.km / u.s,
+            name="Planet X",
+        )
+
+    def simulation():
+        bodies = []
+        if input.earth():
+            bodies.append(earth_body())
+        if input.moon():
+            bodies.append(moon_body())
+        if input.planetx():
+            bodies.append(planetx_body())
+
+        simulation_ = Simulation(bodies)
+        simulation_.set_diff_eq(nbody_solve)
+
+        return simulation_
+
+    has_run = False
+
+    @output
+    @render.plot
+    @reactive.event(input.run, ignore_none=False)
+    def orbits():
+        return make_orbit_plot()
+
+    def make_orbit_plot():
+        sim = simulation()
+        n_steps = input.days() * 24 / input.step_size()
+        with ui.Progress(min=1, max=n_steps) as p:
+            sim.run(input.days() * u.day, input.step_size() * u.hr, progress=p)
+
+        sim_hist = sim.history
+        end_idx = len(sim_hist) - 1
+
+        fig = plt.figure()
+
+        ax = plt.axes(projection="3d")
+
+        n_bodies = int(sim_hist.shape[1] / 6)
+        for i in range(0, n_bodies):
+            ax.scatter3D(
+                sim_hist[end_idx, i * 6],
+                sim_hist[end_idx, i * 6 + 1],
+                sim_hist[end_idx, i * 6 + 2],
+                s=50,
+            )
+            ax.plot3D(
+                sim_hist[:, i * 6],
+                sim_hist[:, i * 6 + 1],
+                sim_hist[:, i * 6 + 2],
+            )
+
+        ax.view_init(30, 20)
+        set_axes_equal(ax)
+
+        return fig
+
+
+www_dir = Path(__file__).parent / "www"
+app = App(app_ui, server, static_assets=www_dir)
+
+
+# https://stackoverflow.com/a/31364297/412655
+def set_axes_equal(ax):
+    """Make axes of 3D plot have equal scale so that spheres appear as spheres,
+    cubes as cubes, etc..  This is one possible solution to Matplotlib's
+    ax.set_aspect('equal') and ax.axis('equal') not working for 3D.
+
+    Input
+      ax: a matplotlib axis, e.g., as output from plt.gca().
+    """
+
+    x_limits = ax.get_xlim3d()
+    y_limits = ax.get_ylim3d()
+    z_limits = ax.get_zlim3d()
+
+    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])
+    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
+    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])

requirements.txt

@@ -0,0 +1,6 @@
+astropy==5.1
+Jinja2
+matplotlib
+pandas
+numpy
+shiny

simulation.py

@@ -0,0 +1,203 @@
+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)