julien-c's picture
julien-c HF staff
initial import
6b6806c verified
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])