use super::orbital::OrbitalDynamics;
use super::thrustctrl::ThrustControl;
use super::{Dynamics, ForceModel};
use crate::celestia::Spacecraft;
use crate::dimensions::{DimName, MatrixN, Vector1, Vector3, VectorN, U3, U4, U42, U43, U6, U7};
use crate::dynamics::Hyperdual;
use crate::errors::NyxError;
use crate::state::State;
use crate::time::TimeUnit;
use crate::TimeTagged;
use std::sync::Arc;
pub use super::solarpressure::SolarPressure;
const NORM_ERR: f64 = 1e-12;
const STD_GRAVITY: f64 = 9.80665;
#[derive(Clone)]
pub struct SpacecraftDynamics<'a> {
pub orbital_dyn: Arc<OrbitalDynamics<'a>>,
pub force_models: Vec<Arc<dyn ForceModel + 'a>>,
pub ctrl: Option<Arc<dyn ThrustControl + 'a>>,
pub decrement_mass: bool,
}
impl<'a> SpacecraftDynamics<'a> {
pub fn with_ctrl(
orbital_dyn: Arc<OrbitalDynamics<'a>>,
ctrl: Arc<dyn ThrustControl + 'a>,
) -> Arc<Self> {
Arc::new(Self {
orbital_dyn,
ctrl: Some(ctrl),
force_models: Vec::new(),
decrement_mass: true,
})
}
pub fn with_ctrl_no_decr(
orbital_dyn: Arc<OrbitalDynamics<'a>>,
ctrl: Arc<dyn ThrustControl + 'a>,
) -> Arc<Self> {
Arc::new(Self {
orbital_dyn,
ctrl: Some(ctrl),
force_models: Vec::new(),
decrement_mass: false,
})
}
pub fn new(orbital_dyn: Arc<OrbitalDynamics<'a>>) -> Arc<Self> {
Arc::new(Self::new_raw(orbital_dyn))
}
pub fn new_raw(orbital_dyn: Arc<OrbitalDynamics<'a>>) -> Self {
Self {
orbital_dyn,
ctrl: None,
force_models: Vec::new(),
decrement_mass: true,
}
}
pub fn from_model(
orbital_dyn: Arc<OrbitalDynamics<'a>>,
force_model: Arc<dyn ForceModel + 'a>,
) -> Arc<Self> {
let mut me = Self::new_raw(orbital_dyn);
me.add_model(force_model);
Arc::new(me)
}
pub fn from_models(
orbital_dyn: Arc<OrbitalDynamics<'a>>,
force_models: Vec<Arc<dyn ForceModel + 'a>>,
) -> Arc<Self> {
let mut me = Self::new_raw(orbital_dyn);
me.force_models = force_models;
Arc::new(me)
}
pub fn add_model(&mut self, force_model: Arc<dyn ForceModel + 'a>) {
self.force_models.push(force_model);
}
pub fn with_model(self, force_model: Arc<dyn ForceModel + 'a>) -> Arc<Self> {
let mut me = self.clone();
me.add_model(force_model);
Arc::new(me)
}
pub fn ctrl_achieved(&self, state: &Spacecraft) -> Result<bool, NyxError> {
match &self.ctrl {
Some(ctrl) => ctrl.achieved(state),
None => Err(NyxError::NoObjectiveDefined),
}
}
}
impl<'a> Dynamics for SpacecraftDynamics<'a> {
type HyperdualSize = U7; type StateType = Spacecraft;
fn finally(&self, next_state: Self::StateType) -> Result<Self::StateType, NyxError> {
if next_state.fuel_mass_kg < 0.0 {
error!("negative fuel mass at {}", next_state.epoch());
return Err(NyxError::FuelExhausted);
}
if let Some(ctrl) = &self.ctrl {
let mut state = next_state;
state.mode = ctrl.next(&state);
Ok(state)
} else {
Ok(next_state)
}
}
fn eom(
&self,
delta_t: f64,
state: &VectorN<f64, U43>,
ctx: &Spacecraft,
) -> Result<VectorN<f64, U43>, NyxError> {
let orbital_dyn_vec = state.fixed_rows::<U42>(0).into_owned();
let d_x_orbital_dyn = self
.orbital_dyn
.eom(delta_t, &orbital_dyn_vec, &ctx.orbit)?;
let mut d_x = VectorN::<f64, U43>::from_iterator(
d_x_orbital_dyn
.iter()
.chain(Vector1::new(0.0).iter())
.cloned(),
);
let mut total_mass = ctx.dry_mass_kg;
let osc_sc = ctx.ctor_from(delta_t, state);
if let Some(ctrl) = &self.ctrl {
let (thrust_force, fuel_rate) = {
if osc_sc.thruster.is_none() {
return Err(NyxError::CtrlExistsButNoThrusterAvail);
}
let thruster = osc_sc.thruster.unwrap();
let thrust_power = ctrl.throttle(&osc_sc);
if !(0.0..=1.0).contains(&thrust_power) {
return Err(NyxError::CtrlThrottleRangeErr(thrust_power));
} else if thrust_power > 0.0 {
let thrust_inertial = ctrl.direction(&osc_sc);
if (thrust_inertial.norm() - 1.0).abs() > NORM_ERR {
return Err(NyxError::CtrlNotAUnitVector(thrust_inertial.norm()));
}
let total_thrust = (thrust_power * thruster.thrust) * 1e-3; (
thrust_inertial * total_thrust,
if self.decrement_mass {
let fuel_usage =
thrust_power * thruster.thrust / (thruster.isp * STD_GRAVITY);
-fuel_usage
} else {
0.0
},
)
} else {
(Vector3::zeros(), 0.0)
}
};
total_mass += ctx.fuel_mass_kg + fuel_rate;
for i in 0..3 {
d_x[i + 3] += thrust_force[i] / (ctx.dry_mass_kg + state[U43::dim() - 1]);
}
d_x[U43::dim() - 1] += fuel_rate;
}
for model in &self.force_models {
let model_frc = model.eom(&osc_sc)? / total_mass;
for i in 0..3 {
d_x[i + 3] += model_frc[i];
}
}
Ok(d_x)
}
fn dual_eom(
&self,
delta_t_s: f64,
state_vec: &VectorN<Hyperdual<f64, Self::HyperdualSize>, U7>,
ctx: &Self::StateType,
) -> Result<(VectorN<f64, U7>, MatrixN<f64, U7>), NyxError> {
let pos_vel = state_vec.fixed_rows::<U6>(0).into_owned();
let (orb_state, orb_grad) = self.orbital_dyn.dual_eom(delta_t_s, &pos_vel, &ctx.orbit)?;
let mut d_x = VectorN::<f64, U7>::from_iterator(
orb_state.iter().chain(Vector1::new(0.0).iter()).cloned(),
);
let mut grad = MatrixN::<f64, U7>::zeros();
for i in 0..U6::dim() {
for j in 0..U6::dim() {
grad[(i, j)] = orb_grad[(i, j)];
}
}
if self.ctrl.is_some() {
return Err(NyxError::PartialsUndefined);
}
let total_mass = ctx.dry_mass_kg;
let radius = state_vec.fixed_rows::<U3>(0).into_owned();
let mut osc_sc = *ctx;
osc_sc.set_epoch(ctx.epoch() + delta_t_s * TimeUnit::Second);
osc_sc.orbit.x = orb_state[0];
osc_sc.orbit.y = orb_state[1];
osc_sc.orbit.z = orb_state[2];
osc_sc.orbit.vx = orb_state[3];
osc_sc.orbit.vy = orb_state[4];
osc_sc.orbit.vz = orb_state[5];
for model in &self.force_models {
let (model_frc, model_grad) = model.dual_eom(&radius, &osc_sc)?;
for i in 0..U3::dim() {
d_x[i + 3] += model_frc[i] / total_mass;
for j in 1..U4::dim() {
grad[(i + 3, j - 1)] += model_grad[(i, j - 1)] / total_mass;
}
}
}
Ok((d_x, grad))
}
}