shiny-orbit-simulation / simulation.py
julien-c's picture
julien-c HF staff
initial import
6b6806c verified
raw
history blame
7.59 kB
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)