Spaces:
Build error
Build error
| 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 | |
| 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]) | |