/*
Nyx, blazing fast astrodynamics
Copyright (C) 2021 Christopher Rabotin <christopher.rabotin@gmail.com>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published
by the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
use super::hyperdual::linalg::norm;
use super::hyperdual::{extract_jacobian_and_result, hyperspace_from_vector, Float, Hyperdual};
use super::{AccelModel, Dynamics, NyxError};
use crate::celestia::{Bodies, Cosm, Frame, LTCorr, Orbit};
use crate::dimensions::{
DimName, Matrix3, Matrix6, Vector3, Vector6, VectorN, U3, U36, U4, U42, U6, U7,
};
use crate::State;
use std::f64;
use std::sync::Arc;
pub use super::sph_harmonics::Harmonics;
/// `OrbitalDynamics` provides the equations of motion for any celestial dynamic, without state transition matrix computation.
#[derive(Clone)]
pub struct OrbitalDynamics<'a> {
pub accel_models: Vec<Arc<dyn AccelModel + Sync + 'a>>,
}
impl<'a> OrbitalDynamics<'a> {
/// Initialize point mass dynamics given the EXB IDs and a Cosm
pub fn point_masses(bodies: &[Bodies], cosm: Arc<Cosm>) -> Arc<Self> {
// Create the point masses
Self::new(vec![PointMasses::new(bodies, cosm)])
}
/// Initializes a OrbitalDynamics which does not simulate the gravity pull of other celestial objects but the primary one.
pub fn two_body() -> Arc<Self> {
Self::new(vec![])
}
/// Initialize orbital dynamics with a list of acceleration models
pub fn new(accel_models: Vec<Arc<dyn AccelModel + Sync + 'a>>) -> Arc<Self> {
Arc::new(Self::new_raw(accel_models))
}
/// Initialize orbital dynamics with a list of acceleration models, _without_ encapsulating it in an Arc
/// Use this only if you need to mutate the dynamics as you'll need to wrap it in an Arc before propagation.
pub fn new_raw(accel_models: Vec<Arc<dyn AccelModel + Sync + 'a>>) -> Self {
Self { accel_models }
}
/// Initialize new orbital mechanics with the provided model.
/// **Note:** Orbital dynamics _always_ include two body dynamics, these cannot be turned off.
pub fn from_model(accel_model: Arc<dyn AccelModel + Sync + 'a>) -> Arc<Self> {
Self::new(vec![accel_model])
}
/// Add a model to the currently defined orbital dynamics
pub fn add_model(&mut self, accel_model: Arc<dyn AccelModel + Sync + 'a>) {
self.accel_models.push(accel_model);
}
/// Clone these dynamics and add a model to the currently defined orbital dynamics
pub fn with_model(self, accel_model: Arc<dyn AccelModel + Sync + 'a>) -> Arc<Self> {
let mut me = self.clone();
me.add_model(accel_model);
Arc::new(me)
}
}
impl<'a> Dynamics for OrbitalDynamics<'a> {
type HyperdualSize = U7;
type StateType = Orbit;
fn eom(
&self,
delta_t_s: f64,
state: &VectorN<f64, U42>,
ctx: &Orbit,
) -> Result<VectorN<f64, U42>, NyxError> {
let (new_state, new_stm) = if ctx.stm.is_some() {
// Then call the dual_eom with the correct state size
let pos_vel = state.fixed_rows::<U6>(0).into_owned();
let (state, grad) = self.eom_grad(delta_t_s, &pos_vel, ctx)?;
let stm_dt = ctx.stm() * grad;
// Rebuild the STM as a vector.
let mut stm_as_vec = VectorN::<f64, U36>::zeros();
let mut stm_idx = 0;
for i in 0..U6::dim() {
for j in 0..U6::dim() {
stm_as_vec[(stm_idx, 0)] = stm_dt[(i, j)];
stm_idx += 1;
}
}
(state, stm_as_vec)
} else {
// Still return something of size 42, but the STM will be zeros.
let osc = ctx.ctor_from(delta_t_s, state);
// TODO: Speed check this with the PointMasses only, including the integration frame point mass
let body_acceleration = (-osc.frame.gm() / osc.rmag().powi(3)) * osc.radius();
let mut d_x = Vector6::from_iterator(
osc.velocity()
.iter()
.chain(body_acceleration.iter())
.cloned(),
);
// Apply the acceleration models
for model in &self.accel_models {
let model_acc = model.eom(&osc)?;
for i in 0..3 {
d_x[i + 3] += model_acc[i];
}
}
(d_x, VectorN::<f64, U36>::zeros())
};
Ok(VectorN::<f64, U42>::from_iterator(
new_state.iter().chain(new_stm.iter()).cloned(),
))
}
fn dual_eom(
&self,
_delta_t_s: f64,
state: &VectorN<Hyperdual<f64, U7>, U6>,
ctx: &Orbit,
) -> Result<(Vector6<f64>, Matrix6<f64>), NyxError> {
// Extract data from hyperspace
let radius = state.fixed_rows::<U3>(0).into_owned();
let velocity = state.fixed_rows::<U3>(3).into_owned();
// Code up math as usual
let rmag = norm(&radius);
let body_acceleration =
radius * (Hyperdual::<f64, U7>::from_real(-ctx.frame.gm()) / rmag.powi(3));
// Extract result into Vector6 and Matrix6
let mut fx = Vector6::zeros();
let mut grad = Matrix6::zeros();
for i in 0..U6::dim() {
fx[i] = if i < 3 {
velocity[i].real()
} else {
body_acceleration[i - 3].real()
};
for j in 1..U7::dim() {
grad[(i, j - 1)] = if i < 3 {
velocity[i][j]
} else {
body_acceleration[i - 3][j]
};
}
}
// Apply the acceleration models
for model in &self.accel_models {
let (model_acc, model_grad) = model.dual_eom(&radius, ctx)?;
for i in 0..U3::dim() {
fx[i + 3] += model_acc[i];
for j in 1..U4::dim() {
grad[(i + 3, j - 1)] += model_grad[(i, j - 1)];
}
}
}
Ok((fx, grad))
}
}
/// PointMasses model
pub struct PointMasses {
pub bodies: Vec<Frame>,
/// Optional point to a Cosm, needed if extra point masses are needed
pub cosm: Arc<Cosm>,
/// Light-time correction computation if extra point masses are needed
pub correction: LTCorr,
}
impl PointMasses {
/// Initializes the multibody point mass dynamics with the provided list of bodies
pub fn new(bodies: &[Bodies], cosm: Arc<Cosm>) -> Arc<Self> {
Arc::new(Self::with_correction(bodies, cosm, LTCorr::None))
}
/// Initializes the multibody point mass dynamics with the provided list of bodies, and accounting for some light time correction
pub fn with_correction(bodies: &[Bodies], cosm: Arc<Cosm>, correction: LTCorr) -> Self {
let mut refs = Vec::new();
// Check that these celestial bodies exist and build their references
for body in bodies {
refs.push(cosm.frame_from_ephem_path(&body.ephem_path()));
}
Self {
bodies: refs,
cosm,
correction,
}
}
/// Allows using bodies by name, defined in the non-default XB
pub fn specific(body_names: &[String], cosm: Arc<Cosm>, correction: LTCorr) -> Self {
let mut refs = Vec::new();
// Check that these celestial bodies exist and build their references
for body in body_names {
refs.push(cosm.frame(body));
}
Self {
bodies: refs,
cosm,
correction,
}
}
}
impl AccelModel for PointMasses {
fn eom(&self, osc: &Orbit) -> Result<Vector3<f64>, NyxError> {
let mut d_x = Vector3::zeros();
// Get all of the position vectors between the center body and the third bodies
for third_body in &self.bodies {
if third_body == &osc.frame {
// Ignore the contribution of the integration frame, that's handled by OrbitalDynamics
continue;
}
// Orbit of j-th body as seen from primary body
let st_ij = self.cosm.celestial_state(
&third_body.ephem_path(),
osc.dt,
osc.frame,
self.correction,
);
let r_ij = st_ij.radius();
let r_ij3 = st_ij.rmag().powi(3);
let r_j = osc.radius() - r_ij; // sc as seen from 3rd body
let r_j3 = r_j.norm().powi(3);
d_x += -third_body.gm() * (r_j / r_j3 + r_ij / r_ij3);
}
Ok(d_x)
}
fn dual_eom(
&self,
state: &VectorN<Hyperdual<f64, U7>, U3>,
osc: &Orbit,
) -> Result<(Vector3<f64>, Matrix3<f64>), NyxError> {
// Extract data from hyperspace
let radius = state.fixed_rows::<U3>(0).into_owned();
// Extract result into Vector6 and Matrix6
let mut fx = Vector3::zeros();
let mut grad = Matrix3::zeros();
// Get all of the position vectors between the center body and the third bodies
for third_body in &self.bodies {
if third_body == &osc.frame {
// Ignore the contribution of the integration frame, that's handled by OrbitalDynamics
continue;
}
let gm_d = Hyperdual::<f64, U7>::from_real(-third_body.gm());
// Orbit of j-th body as seen from primary body
let st_ij = self.cosm.celestial_state(
&third_body.ephem_path(),
osc.dt,
osc.frame,
self.correction,
);
let r_ij: Vector3<Hyperdual<f64, U7>> = hyperspace_from_vector(&st_ij.radius());
let r_ij3 = norm(&r_ij).powi(3) / gm_d; // Dividing the future divisor
// The difference leads to the dual parts nulling themselves out, so let's fix that.
let mut r_j = radius - r_ij; // sc as seen from 3rd body
r_j[0][1] = 1.0;
r_j[1][2] = 1.0;
r_j[2][3] = 1.0;
let r_j3 = norm(&r_j).powi(3) / gm_d; // Dividing the future divisor
let third_body_acc_d = r_j / r_j3 + r_ij / r_ij3;
let (fxp, gradp) = extract_jacobian_and_result::<_, U3, U3, _>(&third_body_acc_d);
fx += fxp;
grad += gradp;
}
Ok((fx, grad))
}
}