SpacesExamples/shiny-with-python
0
1from pathlib import Path2from simulation import Body, Simulation, nbody_solve, spherical_to_cartesian3import matplotlib.pyplot as plt4import astropy.units as u5import numpy as np6 7from shiny import App, reactive, render, ui8 9# This application adapted from RK4 Orbit Integrator tutorial in Python for Astronomers10# https://prappleizer.github.io/11 12 13def panel_box(*args, **kwargs):14 return ui.div(15 ui.div(*args, class_="card-body"),16 **kwargs,17 class_="card mb-3",18 )19 20 21app_ui = ui.page_fluid(22 {"class": "p-4"},23 ui.row(24 ui.column(25 4,26 panel_box(27 ui.input_slider("days", "Simulation duration (days)", 0, 200, value=60),28 ui.input_slider(29 "step_size",30 "Simulation time step (hours)",31 0,32 24,33 value=4,34 step=0.5,35 ),36 ui.input_action_button(37 "run", "Run simulation", class_="btn-primary w-100"38 ),39 ),40 ui.navset_tab_card(41 ui.nav(42 "Earth",43 ui.input_checkbox("earth", "Enable", True),44 ui.panel_conditional(45 "input.earth",46 ui.input_numeric(47 "earth_mass",48 "Mass (10^22 kg)",49 597.216,50 ),51 ui.input_slider(52 "earth_speed",53 "Speed (km/s)",54 0,55 1,56 value=0.0126,57 step=0.001,58 ),59 ui.input_slider("earth_theta", "Angle (5)", 0, 360, value=270),60 ui.input_slider("earth_phi", "5", 0, 180, value=90),61 ),62 ),63 ui.nav(64 "Moon",65 ui.input_checkbox("moon", "Enable", True),66 ui.panel_conditional(67 "input.moon",68 ui.input_numeric("moon_mass", "Mass (10^22 kg)", 7.347),69 ui.input_slider(70 "moon_speed", "Speed (km/s)", 0, 2, value=1.022, step=0.00171 ),72 ui.input_slider("moon_theta", "Angle (5)", 0, 360, value=90),73 ui.input_slider("moon_phi", "5", 0, 180, value=90),74 ),75 ),76 ui.nav(77 "Planet X",78 ui.input_checkbox("planetx", "Enable", False),79 ui.output_ui("planetx_controls"),80 ui.panel_conditional(81 "input.planetx",82 ui.input_numeric("planetx_mass", "Mass (10^22 kg)", 7.347),83 ui.input_slider(84 "planetx_speed",85 "Speed (km/s)",86 0,87 2,88 value=1.022,89 step=0.001,90 ),91 ui.input_slider("planetx_theta", "Angle (5)", 0, 360, 270),92 ui.input_slider("planetx_phi", "5", 0, 180, 90),93 ),94 ),95 ),96 ),97 ui.column(98 8,99 ui.output_plot("orbits", width="500px", height="500px"),100 ui.img(src="coords.png", style="width: 100%; max-width: 250px;"),101 ),102 ),103)104 105 106def server(input, output, session):107 def earth_body():108 v = spherical_to_cartesian(109 input.earth_theta(), input.earth_phi(), input.earth_speed()110 )111 112 return Body(113 mass=input.earth_mass() * 10e21 * u.kg,114 x_vec=np.array([0, 0, 0]) * u.km,115 v_vec=np.array(v) * u.km / u.s,116 name="Earth",117 )118 119 def moon_body():120 v = spherical_to_cartesian(121 input.moon_theta(), input.moon_phi(), input.moon_speed()122 )123 124 return Body(125 mass=input.moon_mass() * 10e21 * u.kg,126 x_vec=np.array([3.84e5, 0, 0]) * u.km,127 v_vec=np.array(v) * u.km / u.s,128 name="Moon",129 )130 131 def planetx_body():132 v = spherical_to_cartesian(133 input.planetx_theta(), input.planetx_phi(), input.planetx_speed()134 )135 136 return Body(137 mass=input.planetx_mass() * 10e21 * u.kg,138 x_vec=np.array([-3.84e5, 0, 0]) * u.km,139 v_vec=np.array(v) * u.km / u.s,140 name="Planet X",141 )142 143 def simulation():144 bodies = []145 if input.earth():146 bodies.append(earth_body())147 if input.moon():148 bodies.append(moon_body())149 if input.planetx():150 bodies.append(planetx_body())151 152 simulation_ = Simulation(bodies)153 simulation_.set_diff_eq(nbody_solve)154 155 return simulation_156 157 has_run = False158 159 @output160 @render.plot161 @reactive.event(input.run, ignore_none=False)162 def orbits():163 return make_orbit_plot()164 165 def make_orbit_plot():166 sim = simulation()167 n_steps = input.days() * 24 / input.step_size()168 with ui.Progress(min=1, max=n_steps) as p:169 sim.run(input.days() * u.day, input.step_size() * u.hr, progress=p)170 171 sim_hist = sim.history172 end_idx = len(sim_hist) - 1173 174 fig = plt.figure()175 176 ax = plt.axes(projection="3d")177 178 n_bodies = int(sim_hist.shape[1] / 6)179 for i in range(0, n_bodies):180 ax.scatter3D(181 sim_hist[end_idx, i * 6],182 sim_hist[end_idx, i * 6 + 1],183 sim_hist[end_idx, i * 6 + 2],184 s=50,185 )186 ax.plot3D(187 sim_hist[:, i * 6],188 sim_hist[:, i * 6 + 1],189 sim_hist[:, i * 6 + 2],190 )191 192 ax.view_init(30, 20)193 set_axes_equal(ax)194 195 return fig196 197 198www_dir = Path(__file__).parent / "www"199app = App(app_ui, server, static_assets=www_dir)200 201 202# https://stackoverflow.com/a/31364297/412655203def set_axes_equal(ax):204 """Make axes of 3D plot have equal scale so that spheres appear as spheres,205 cubes as cubes, etc.. This is one possible solution to Matplotlib's206 ax.set_aspect('equal') and ax.axis('equal') not working for 3D.207 208 Input209 ax: a matplotlib axis, e.g., as output from plt.gca().210 """211 212 x_limits = ax.get_xlim3d()213 y_limits = ax.get_ylim3d()214 z_limits = ax.get_zlim3d()215 216 x_range = abs(x_limits[1] - x_limits[0])217 x_middle = np.mean(x_limits)218 y_range = abs(y_limits[1] - y_limits[0])219 y_middle = np.mean(y_limits)220 z_range = abs(z_limits[1] - z_limits[0])221 z_middle = np.mean(z_limits)222 223 # The plot bounding box is a sphere in the sense of the infinity224 # norm, hence I call half the max range the plot radius.225 plot_radius = 0.5 * max([x_range, y_range, z_range])226 227 ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])228 ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])229 ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])230