From 34813f5a8b6675d9faba446c0bc0cf233d497516 Mon Sep 17 00:00:00 2001 From: Connor Johnstone Date: Wed, 30 Sep 2026 12:44:05 -0400 Subject: [PATCH] Initial stab at propagator + two body dynamics --- Cargo.toml | 1 + examples/two_body.rs | 41 ++- src/dynamics/mod.rs | 278 +++++++++++++++++++ src/dynamics/two_body.rs | 74 ++++++ src/lib.rs | 13 +- src/propagation.rs | 561 +++++++++++++++++++++++++++++++++++++++ src/state.rs | 199 ++++++++++++++ 7 files changed, 1142 insertions(+), 25 deletions(-) create mode 100644 src/dynamics/mod.rs create mode 100644 src/dynamics/two_body.rs create mode 100644 src/propagation.rs diff --git a/Cargo.toml b/Cargo.toml index f3e2498..93bb562 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,5 +27,6 @@ thiserror = "2.0" [dev-dependencies] approx = "0.5" +serde_json = "1.0" criterion = "0.7" num-dual = "0.15" diff --git a/examples/two_body.rs b/examples/two_body.rs index b8273ef..ea9028b 100644 --- a/examples/two_body.rs +++ b/examples/two_body.rs @@ -2,9 +2,8 @@ //! //! Run with `cargo run --example two_body`. -use astrokit::state::{HasPosition, HasVelocity, Orbit}; -use nalgebra::{SVector, Vector6}; -use ordinary_diffeq::prelude::*; +use astrokit::prelude::*; +use nalgebra::Vector6; use std::f64::consts::PI; /// Earth's gravitational parameter, m^3/s^2. @@ -13,21 +12,12 @@ const MU: f64 = 3.986004415e14; /// Semi-major axis of the test orbit, m. const SMA: f64 = 6.7781363e6; -/// The two-body derivative. -fn two_body + HasVelocity>( - _t: T, - y: SVector, - &(mu,): &(T,), -) -> SVector { - let r = L::position(&y); - let mut dy = SVector::::zeros(); - L::set_position(&mut dy, L::velocity(&y)); - L::set_velocity(&mut dy, -(r * mu) / r.norm().powi(3)); - dy -} +/// Earth's equatorial radius, m. Comes from a frame once frames land. +const EARTH_RADIUS: f64 = 6.378137e6; fn main() { - let period = 2.0 * PI * (SMA.powi(3) / MU).sqrt(); + let epoch = Epoch::from_gregorian_utc_hms(2026, 1, 1, 0, 0, 0); + let period = Duration::from_seconds(2.0 * PI * (SMA.powi(3) / MU).sqrt()); let y0 = Vector6::new( 4.263868426884883e6, @@ -38,24 +28,27 @@ fn main() { 1870.3893008991558, ); - let ode = ODE::new(&two_body::, 0.0, period, y0, (MU,)); - let vern7 = Vern7::new().a_tol(1e-12).r_tol(1e-12); - let mut problem = Problem::new(ode, vern7, PIController::default()); - let solution = problem.solve(); + let forces = vec![DynamicsSpec::TwoBody { mu: MU }]; + + let mut propagation = Propagation::new(epoch, Orbit, y0, &forces) + .duration(period) + .accuracy(Accuracy::High); + let solution = propagation.run(); let end = solution.states.last().unwrap(); let position_error = (Orbit::position(end) - Orbit::position(&y0)).norm(); let velocity_error = (Orbit::velocity(end) - Orbit::velocity(&y0)).norm(); - println!("period: {period:.3} s"); + println!("epoch: {epoch}"); + println!("period: {period}"); println!("steps: {}", solution.states.len() - 1); println!("position error: {position_error:.3e} m"); println!("velocity error: {velocity_error:.3e} m/s"); - // Dense output: the state at any time in the arc, not just at a step. - let half = solution.interpolate(period / 2.0); + // Dense output: the state at any time in the arc, not only at a step. + let half = solution.interpolate(period.to_seconds() / 2.0); println!( "altitude at half a period: {:.3} km", - (Orbit::position(&half).norm() - 6.378137e6) / 1e3 + (Orbit::position(&half).norm() - EARTH_RADIUS) / 1e3 ); } diff --git a/src/dynamics/mod.rs b/src/dynamics/mod.rs new file mode 100644 index 0000000..c051725 --- /dev/null +++ b/src/dynamics/mod.rs @@ -0,0 +1,278 @@ +//! The models that contribute to a state's derivative. +//! +//! A [`Dynamics`] returns a [`Contribution`] with four channels — acceleration, +//! torque, mass rate and power — all defaulting to zero, so a model implements +//! only what it affects. A solar array contributes power and nothing else. +//! +//! # Specification and instantiation are separate +//! +//! A trait object is tied to one scalar: a `Vec>>` +//! cannot be propagated in a dual. So what a user composes is +//! [`DynamicsSpec`] — plain data, serializable, scalar-free — and +//! [`DynamicsSpec::build`] instantiates trait objects at whatever scalar is being +//! propagated: +//! +//! ``` +//! use astrokit::dynamics::{build, DynamicsSpec}; +//! use astrokit::state::Orbit; +//! +//! let spec = vec![DynamicsSpec::TwoBody { mu: 3.986004415e14 }]; +//! +//! // The same spec, at two scalars. +//! let plain = build::(&spec); +//! # #[cfg(feature = "autodiff")] +//! let dual = build::, 6, Orbit>(&spec); +//! ``` +//! +//! The spec is also what gets serialized into a run's provenance, so a +//! propagation's force composition is recoverable from the record rather than +//! hidden behind a mode flag. +//! +//! A side benefit: because a model is constructed at the propagated scalar, it +//! stores its constants already converted. `TwoBody` holds `mu` as a `T`, so +//! `T::lift` happens once at build time and never in the derivative. + +use hifitime::{Duration, Epoch}; +use nalgebra::{SVector, Vector3}; +use ordinary_diffeq::scalar::Scalar; +use serde::{Deserialize, Serialize}; +use std::ops::{Add, AddAssign}; + +use crate::state::{HasPosition, HasVelocity, Layout}; + +pub mod two_body; + +pub use two_body::TwoBody; + +/// What one [`Dynamics`] contributes at an instant. +/// +/// Every channel defaults to zero. Build one with [`Contribution::default`] and +/// struct update syntax, or [`Contribution::from_acceleration`] for the common +/// case. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Contribution { + /// Inertial acceleration, m/s^2. + pub acceleration: Vector3, + /// Body-frame torque, N m. + pub torque: Vector3, + /// Rate of change of mass, kg/s. Negative when propellant is spent. + pub mass_rate: T, + /// Net power, W. Positive into the power system. + pub power: T, +} + +impl Default for Contribution { + fn default() -> Self { + Self { + acceleration: Vector3::zeros(), + torque: Vector3::zeros(), + mass_rate: T::lift(0.0), + power: T::lift(0.0), + } + } +} + +impl Contribution { + /// A contribution that is only an acceleration. + #[inline] + pub fn from_acceleration(acceleration: Vector3) -> Self { + Self { + acceleration, + ..Default::default() + } + } +} + +impl AddAssign for Contribution { + #[inline] + fn add_assign(&mut self, rhs: Self) { + self.acceleration += rhs.acceleration; + self.torque += rhs.torque; + self.mass_rate += rhs.mass_rate; + self.power += rhs.power; + } +} + +impl Add for Contribution { + type Output = Self; + + #[inline] + fn add(mut self, rhs: Self) -> Self { + self += rhs; + self + } +} + +/// What a [`Dynamics`] needs that is not in the state vector. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Context { + /// The epoch the propagation started from. The `t` a dynamics receives is + /// seconds after this. + pub epoch: Epoch, +} + +impl Context { + /// A context anchored at `epoch`. + pub fn new(epoch: Epoch) -> Self { + Self { epoch } + } + + /// The epoch `t` seconds after the start of the propagation. + /// + /// `t` is projected to its real part: an epoch is a real quantity, and there + /// is no such thing as a dual date. A propagation differentiated with respect + /// to its initial state is unaffected by this, since the time grid does not + /// depend on the state. It does mean sensitivities with respect to epoch + /// itself are not available this way. + pub fn epoch_at(&self, t: T) -> Epoch { + self.epoch + Duration::from_seconds(t.re()) + } +} + +/// A model contributing to the state derivative. +/// +/// Generic at the trait level rather than the method level, so the methods stay +/// non-generic and the trait stays object-safe: `Box>` +/// works, which is what lets a propagation hold a heterogeneous force list. +/// +/// Implement it for the capabilities the model actually needs. `TwoBody` requires +/// only `L: HasPosition`, so it works for every layout with a position. +pub trait Dynamics> { + /// This model's contribution `t` seconds after the propagation epoch. + fn contribution(&self, t: T, y: &SVector, context: &Context) -> Contribution; + + /// A short name, for provenance records and diagnostics. + fn name(&self) -> &'static str; +} + +/// A dynamics as plain, serializable, scalar-free data. +/// +/// This is what a user composes and what a run's provenance records. +/// [`build`] turns a list of these into trait objects at the scalar being +/// propagated. +#[non_exhaustive] +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +pub enum DynamicsSpec { + /// Point-mass gravity of a central body, `mu` in m^3/s^2. + TwoBody { + /// Gravitational parameter, m^3/s^2. + mu: f64, + }, +} + +impl DynamicsSpec { + /// Instantiates this model at the scalar `T` and layout `L`. + pub fn build(&self) -> Box> + where + T: Scalar + 'static, + L: HasPosition + HasVelocity + 'static, + { + match *self { + Self::TwoBody { mu } => Box::new(TwoBody::new(mu)), + } + } +} + +/// Instantiates a whole force list at the scalar `T` and layout `L`. +pub fn build(specs: &[DynamicsSpec]) -> Vec>> +where + T: Scalar + 'static, + L: HasPosition + HasVelocity + 'static, +{ + specs.iter().map(DynamicsSpec::build).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::Orbit; + use nalgebra::Vector6; + + #[test] + fn test_contribution_defaults_to_zero() { + let c = Contribution::::default(); + assert_eq!(c.acceleration, Vector3::zeros()); + assert_eq!(c.torque, Vector3::zeros()); + assert_eq!(c.mass_rate, 0.0); + assert_eq!(c.power, 0.0); + } + + #[test] + fn test_contributions_accumulate_per_channel() { + let a = Contribution { + acceleration: Vector3::new(1.0, 0.0, 0.0), + power: 50.0, + ..Default::default() + }; + let b = Contribution { + acceleration: Vector3::new(0.0, 2.0, 0.0), + mass_rate: -0.01, + ..Default::default() + }; + + let total = a + b; + assert_eq!(total.acceleration, Vector3::new(1.0, 2.0, 0.0)); + assert_eq!(total.mass_rate, -0.01); + assert_eq!(total.power, 50.0); + assert_eq!(total.torque, Vector3::zeros()); + } + + #[test] + fn test_spec_round_trips_through_json() { + let spec = vec![DynamicsSpec::TwoBody { mu: 3.986004415e14 }]; + let json = serde_json::to_string(&spec).unwrap(); + let back: Vec = serde_json::from_str(&json).unwrap(); + assert_eq!(spec, back); + } + + /// The property the spec/instantiation split exists for: one spec value, + /// built at two different scalars. + #[test] + fn test_one_spec_builds_at_two_scalars() { + let spec = vec![DynamicsSpec::TwoBody { mu: 3.986004415e14 }]; + + let plain = build::(&spec); + assert_eq!(plain.len(), 1); + assert_eq!(plain[0].name(), "two_body"); + + #[cfg(feature = "autodiff")] + { + let dual = build::, 6, Orbit>(&spec); + assert_eq!(dual.len(), 1); + } + } + + #[test] + fn test_context_advances_an_epoch() { + let epoch = Epoch::from_gregorian_utc_hms(2026, 1, 1, 0, 0, 0); + let context = Context::new(epoch); + assert_eq!(context.epoch_at(0.0), epoch); + assert_eq!( + context.epoch_at(3600.0), + Epoch::from_gregorian_utc_hms(2026, 1, 1, 1, 0, 0) + ); + } + + /// A dynamics bounded on a capability works for any layout that has it. + #[test] + fn test_two_body_acceleration_is_layout_agnostic() { + let context = Context::new(Epoch::from_gregorian_utc_hms(2026, 1, 1, 0, 0, 0)); + let mu = 3.986004415e14; + let r = 7.0e6; + + let orbit = Vector6::new(r, 0.0, 0.0, 0.0, 7546.0, 0.0); + let model = TwoBody::::new(mu); + let a_orbit = + as Dynamics>::contribution(&model, 0.0, &orbit, &context); + + let mut full = SVector::::zeros(); + crate::state::Full::set_position(&mut full, Vector3::new(r, 0.0, 0.0)); + let a_full = as Dynamics>::contribution( + &model, 0.0, &full, &context, + ); + + assert_eq!(a_orbit.acceleration, a_full.acceleration); + // a = -mu/r^2 along x + approx::assert_relative_eq!(a_orbit.acceleration[0], -mu / (r * r), max_relative = 1e-15); + } +} diff --git a/src/dynamics/two_body.rs b/src/dynamics/two_body.rs new file mode 100644 index 0000000..1777a4f --- /dev/null +++ b/src/dynamics/two_body.rs @@ -0,0 +1,74 @@ +//! Point-mass gravity. + +use nalgebra::SVector; +use ordinary_diffeq::scalar::Scalar; + +use super::{Context, Contribution, Dynamics}; +use crate::state::HasPosition; + +/// Point-mass gravity of a single central body, at the origin of the state's +/// frame. +/// +/// `mu` is stored in the propagated scalar, converted once at construction, so +/// the derivative does no lifting. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct TwoBody { + mu: T, +} + +impl TwoBody { + /// Point-mass gravity with gravitational parameter `mu`, in m^3/s^2. + pub fn new(mu: f64) -> Self { + Self { mu: T::lift(mu) } + } + + /// The gravitational parameter, m^3/s^2. + pub fn mu(&self) -> T { + self.mu + } +} + +impl> Dynamics for TwoBody { + #[inline] + fn contribution(&self, _t: T, y: &SVector, _context: &Context) -> Contribution { + let r = L::position(y); + Contribution::from_acceleration(-(r * self.mu) / r.norm().powi(3)) + } + + fn name(&self) -> &'static str { + "two_body" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::Orbit; + use approx::assert_relative_eq; + use hifitime::Epoch; + use nalgebra::Vector6; + + const MU: f64 = 3.986004415e14; + + #[test] + fn test_acceleration_points_at_the_origin() { + let context = Context::new(Epoch::from_gregorian_utc_hms(2026, 1, 1, 0, 0, 0)); + let model = TwoBody::::new(MU); + let y = Vector6::new(3.0e6, 4.0e6, 0.0, 0.0, 0.0, 0.0); + + let c = as Dynamics>::contribution(&model, 0.0, &y, &context); + + // Antiparallel to the position, magnitude mu / r^2 at r = 5e6. + let r = 5.0e6; + assert_relative_eq!(c.acceleration.norm(), MU / (r * r), max_relative = 1e-14); + let along = c + .acceleration + .normalize() + .dot(&Orbit::position(&y).normalize()); + assert_relative_eq!(along, -1.0, max_relative = 1e-14); + + // And nothing else is contributed. + assert_eq!(c.mass_rate, 0.0); + assert_eq!(c.power, 0.0); + } +} diff --git a/src/lib.rs b/src/lib.rs index 801314c..194e6bd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,7 +9,18 @@ #![cfg_attr(feature = "autodiff", doc = include_str!("../readme.md"))] #![warn(missing_docs)] +pub mod dynamics; +pub mod propagation; pub mod state; /// The types you want in scope to set up and run a propagation. -pub mod prelude {} +pub mod prelude { + pub use crate::dynamics::{Context, Contribution, Dynamics, DynamicsSpec}; + pub use crate::propagation::{Accuracy, Propagation}; + pub use crate::state::{Full, Orbit, OrbitMass}; + pub use crate::state::{ + HasAngularRate, HasAttitude, HasEnergy, HasMass, HasPosition, HasVelocity, Layout, + }; + pub use hifitime::{Duration, Epoch, Unit}; + pub use ordinary_diffeq::scalar::Scalar; +} diff --git a/src/propagation.rs b/src/propagation.rs new file mode 100644 index 0000000..3dc6830 --- /dev/null +++ b/src/propagation.rs @@ -0,0 +1,561 @@ +//! Assembling a state, a force list and an integrator into a propagation. + +use hifitime::{Duration, Epoch}; +use nalgebra::SVector; +use ordinary_diffeq::integrator::bs3::BS3; +use ordinary_diffeq::prelude::*; +use serde::{Deserialize, Serialize}; + +use crate::dynamics::{self, Context, Contribution, Dynamics, DynamicsSpec}; +use crate::state::{HasPosition, HasVelocity, Layout}; + +/// How precisely the equations are solved. +/// +/// Measured on `ordinary-diffeq` 0.4, two-body, Dormand-Prince, `a_tol = r_tol`: +/// +/// | Preset | tolerance | steps/day | position error after one period | +/// |---|---|---|---| +/// | [`Accuracy::Low`] | 1e-8 | 1,573 | 4.4e-2 m | +/// | [`Accuracy::Medium`] | 1e-10 | 3,803 | 1.8e-3 m | +/// | [`Accuracy::High`] | 1e-12 | 9,391 | 2.4e-5 m | +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum Accuracy { + /// 1e-8. + Low, + /// 1e-10. + Medium, + /// 1e-12. + High, +} + +impl Accuracy { + /// The absolute and relative tolerance this preset sets. + pub fn tolerance(self) -> f64 { + match self { + Accuracy::Low => 1e-8, + Accuracy::Medium => 1e-10, + Accuracy::High => 1e-12, + } + } +} + +/// An integrator whose error tolerances can be set generically. +/// +/// `ordinary-diffeq` exposes `a_tol` and `r_tol` as inherent builder methods on +/// each integrator rather than through a trait, so this is the adapter that lets +/// [`Accuracy`] and [`Propagation::a_tol`] work over whichever one is in use. +pub trait Tolerances: Integrator { + /// Sets the absolute tolerance for every component. + fn with_a_tol(self, a_tol: f64) -> Self; + + /// Sets the absolute tolerance per component. + fn with_a_tol_per_component(self, a_tol: SVector) -> Self; + + /// Sets the relative tolerance. + fn with_r_tol(self, r_tol: f64) -> Self; +} + +macro_rules! impl_tolerances { + ($ty:ident) => { + impl Tolerances for $ty + where + $ty: Integrator, + { + fn with_a_tol(self, a_tol: f64) -> Self { + self.a_tol(a_tol) + } + fn with_a_tol_per_component(self, a_tol: SVector) -> Self { + self.a_tol_full(a_tol) + } + fn with_r_tol(self, r_tol: f64) -> Self { + self.r_tol(r_tol) + } + } + }; +} + +impl_tolerances!(Vern7); +impl_tolerances!(DormandPrince45); +impl_tolerances!(BS3); + +/// The built force list and its context: what the right-hand side reads. +/// +/// Built fresh per propagation, at the scalar being propagated, from a +/// [`DynamicsSpec`] list. +pub struct Assembly> { + models: Vec>>, + context: Context, +} + +impl> Assembly { + /// The names of the models in the list, in order. + pub fn names(&self) -> Vec<&'static str> { + self.models.iter().map(|m| m.name()).collect() + } +} + +/// The state derivative: kinematics, plus every model's contribution routed into +/// the components the layout has. +/// +/// Position takes the velocity, velocity takes the summed acceleration, and mass +/// and stored energy take the summed mass rate and power when the layout carries +/// them. A channel with nowhere to go is a programmer error and trips a +/// `debug_assert`. +pub fn state_derivative( + t: T, + y: SVector, + assembly: &Assembly, +) -> SVector +where + T: Scalar, + L: HasPosition + HasVelocity, +{ + let mut total = Contribution::default(); + for model in &assembly.models { + total += model.contribution(t, &y, &assembly.context); + } + + let mut dy = SVector::::zeros(); + L::set_position(&mut dy, L::velocity(&y)); + L::set_velocity(&mut dy, total.acceleration); + let mass_routed = L::try_set_mass(&mut dy, total.mass_rate); + let power_routed = L::try_set_energy(&mut dy, total.power); + + debug_assert!( + mass_routed || total.mass_rate.re() == 0.0, + "a dynamics contributed a mass rate but layout {} has no mass component", + L::NAME + ); + debug_assert!( + power_routed || total.power.re() == 0.0, + "a dynamics contributed power but layout {} has no energy component", + L::NAME + ); + debug_assert!( + total.torque.norm().re() == 0.0, + "a dynamics contributed a torque, which needs an inertia tensor to route" + ); + + dy +} + +/// The `ordinary-diffeq` problem a [`Propagation`] assembles and holds. +/// +/// Held rather than built in a local because a `Solution` borrows its ODE. +type BuiltProblem = + Problem<'static, T, D, S, Assembly, PIController>; + +/// A propagation, configured then run. +/// +/// The layout is passed as a value so that it can be inferred: +/// `Propagation::new(epoch, Orbit, y0, &forces)` rather than a turbofish. +/// +/// The integrator is a type parameter, defaulting to [`Vern7`], because +/// `ordinary-diffeq`'s `Integrator` trait keeps stage counts in associated +/// constants — so it cannot be chosen at runtime, and [`Accuracy`] sets +/// tolerances only. Change it with [`Propagation::method`]. +pub struct Propagation< + T: Scalar + 'static, + const D: usize, + L: Layout + 'static, + S: Integrator = Vern7, +> { + epoch: Epoch, + initial_state: SVector, + duration: Option, + specs: Vec, + integrator: S, + controller: PIController, + problem: Option>, +} + +impl Propagation> +where + T: Scalar + 'static, + L: HasPosition + HasVelocity + 'static, + Vern7: Integrator, +{ + /// A propagation starting at `epoch` from `initial_state`, under `dynamics`. + /// + /// Set a duration with [`Propagation::duration`] or an end with + /// [`Propagation::until`] before running. + pub fn new( + epoch: Epoch, + _layout: L, + initial_state: SVector, + dynamics: &[DynamicsSpec], + ) -> Self { + Self { + epoch, + initial_state, + duration: None, + specs: dynamics.to_vec(), + integrator: Vern7::new().a_tol(1e-10).r_tol(1e-10), + controller: PIController::default(), + problem: None, + } + } +} + +impl Propagation +where + T: Scalar + 'static, + L: HasPosition + HasVelocity + 'static, + S: Tolerances + Copy, +{ + /// How long to propagate for. + pub fn duration(mut self, duration: Duration) -> Self { + self.duration = Some(duration); + self + } + + /// Propagate until this epoch. + pub fn until(mut self, epoch: Epoch) -> Self { + self.duration = Some(epoch - self.epoch); + self + } + + /// Sets both tolerances from a preset. + /// + /// Apply it before [`Propagation::a_tol`] or [`Propagation::r_tol`] if you + /// want to override one of them. + pub fn accuracy(mut self, accuracy: Accuracy) -> Self { + let tol = accuracy.tolerance(); + self.integrator = self.integrator.with_a_tol(tol).with_r_tol(tol); + self + } + + /// Sets the absolute tolerance for every component. + pub fn a_tol(mut self, a_tol: f64) -> Self { + self.integrator = self.integrator.with_a_tol(a_tol); + self + } + + /// Sets the absolute tolerance per component. + pub fn a_tol_per_component(mut self, a_tol: SVector) -> Self { + self.integrator = self.integrator.with_a_tol_per_component(a_tol); + self + } + + /// Sets the relative tolerance. + pub fn r_tol(mut self, r_tol: f64) -> Self { + self.integrator = self.integrator.with_r_tol(r_tol); + self + } + + /// Replaces the step size controller. + pub fn controller(mut self, controller: PIController) -> Self { + self.controller = controller; + self + } + + /// Replaces the integrator, keeping everything else. + /// + /// Tolerances travel with the integrator in `ordinary-diffeq`, so set them + /// after this, not before. + pub fn method(self, integrator: S2) -> Propagation + where + S2: Tolerances + Copy, + { + Propagation { + epoch: self.epoch, + initial_state: self.initial_state, + duration: self.duration, + specs: self.specs, + integrator, + controller: self.controller, + problem: None, + } + } + + /// The force list, as the data it was specified by. + pub fn dynamics(&self) -> &[DynamicsSpec] { + &self.specs + } + + /// Runs the propagation. + /// + /// The returned solution borrows this `Propagation`, because + /// `ordinary-diffeq`'s `Solution` borrows the ODE it came from — Vern7 + /// computes six of its sixteen stages only when something interpolates inside + /// a step, so it needs the right-hand side after the solve has finished. + /// + /// # Panics + /// + /// If neither [`Propagation::duration`] nor [`Propagation::until`] was called. + pub fn run(&mut self) -> Solution<'_, T, S, D, Assembly> { + let duration = self + .duration + .expect("set a duration or an end epoch before running a propagation"); + + let assembly = Assembly { + models: dynamics::build::(&self.specs), + context: Context::new(self.epoch), + }; + + let ode = ODE::new( + &state_derivative::, + T::lift(0.0), + T::lift(duration.to_seconds()), + self.initial_state, + assembly, + ); + + self.problem = Some(Problem::new(ode, self.integrator, self.controller)); + self.problem.as_mut().expect("just assigned").solve() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::{Full, HasAttitude, HasMass, Orbit, OrbitMass}; + use approx::assert_relative_eq; + use nalgebra::{Vector3, Vector6}; + use std::f64::consts::PI; + + const MU: f64 = 3.986004415e14; + const SMA: f64 = 6.7781363e6; + + fn epoch() -> Epoch { + Epoch::from_gregorian_utc_hms(2026, 1, 1, 0, 0, 0) + } + + fn two_body() -> Vec { + vec![DynamicsSpec::TwoBody { mu: MU }] + } + + fn leo() -> Vector6 { + Vector6::new( + 4.263868426884883e6, + 5.146189057155391e6, + 1.1310208421331816e6, + -5923.454461876975, + 4496.802639690076, + 1870.3893008991558, + ) + } + + fn period() -> f64 { + 2.0 * PI * (SMA.powi(3) / MU).sqrt() + } + + #[test] + fn test_two_body_orbit_closes_over_one_period() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(Duration::from_seconds(period())) + .accuracy(Accuracy::High); + let solution = propagation.run(); + + let end = solution.states.last().unwrap(); + let position_error = (Orbit::position(end) - Orbit::position(&leo())).norm(); + let velocity_error = (Orbit::velocity(end) - Orbit::velocity(&leo())).norm(); + + assert!(position_error < 1e-5, "position error {position_error:e} m"); + assert!( + velocity_error < 1e-8, + "velocity error {velocity_error:e} m/s" + ); + } + + /// Two-body motion conserves specific energy and angular momentum. Both are + /// independent of the integrator's own error estimate, so they are a real + /// check rather than a restatement of the tolerance. + #[test] + fn test_two_body_conserves_energy_and_angular_momentum() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(Duration::from_seconds(5.0 * period())) + .accuracy(Accuracy::High); + let solution = propagation.run(); + + let energy = |y: &Vector6| { + let (r, v) = (Orbit::position(y), Orbit::velocity(y)); + v.norm_squared() * 0.5 - MU / r.norm() + }; + let momentum = |y: &Vector6| Orbit::position(y).cross(&Orbit::velocity(y)); + + let e0 = energy(&leo()); + let h0 = momentum(&leo()); + for state in &solution.states { + assert_relative_eq!(energy(state), e0, max_relative = 1e-12); + assert_relative_eq!(momentum(state).norm(), h0.norm(), max_relative = 1e-12); + } + } + + /// The same spec and the same dynamics, propagated on three layouts of three + /// different dimensions, must give the same trajectory. + #[test] + fn test_translational_result_is_independent_of_layout() { + let duration = Duration::from_seconds(period()); + + let mut on_orbit = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(duration) + .accuracy(Accuracy::High); + let reference = *on_orbit.run().states.last().unwrap(); + + let mut y7 = SVector::::zeros(); + OrbitMass::set_position(&mut y7, Orbit::position(&leo())); + OrbitMass::set_velocity(&mut y7, Orbit::velocity(&leo())); + OrbitMass::set_mass(&mut y7, 500.0); + let mut on_orbit_mass = Propagation::new(epoch(), OrbitMass, y7, &two_body()) + .duration(duration) + .accuracy(Accuracy::High); + let end7 = *on_orbit_mass.run().states.last().unwrap(); + + let mut y14 = SVector::::zeros(); + Full::set_position(&mut y14, Orbit::position(&leo())); + Full::set_velocity(&mut y14, Orbit::velocity(&leo())); + Full::set_mass(&mut y14, 500.0); + let mut on_full = Propagation::new(epoch(), Full, y14, &two_body()) + .duration(duration) + .accuracy(Accuracy::High); + let end14 = *on_full.run().states.last().unwrap(); + + assert_relative_eq!( + OrbitMass::position(&end7), + Orbit::position(&reference), + max_relative = 1e-9 + ); + assert_relative_eq!( + Full::position(&end14), + Orbit::position(&reference), + max_relative = 1e-9 + ); + + // Mass is untouched with no model contributing a mass rate. + assert_eq!(OrbitMass::mass(&end7), 500.0); + // And attitude is inert, as `state_derivative` documents. + assert_eq!(Full::attitude(&end14), Vector3::zeros()); + } + + #[test] + fn test_accuracy_presets_trade_steps_for_error() { + let duration = Duration::from_seconds(period()); + let mut results = Vec::new(); + for accuracy in [Accuracy::Low, Accuracy::Medium, Accuracy::High] { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(duration) + .accuracy(accuracy); + let solution = propagation.run(); + let steps = solution.states.len() - 1; + let error = + (Orbit::position(solution.states.last().unwrap()) - Orbit::position(&leo())).norm(); + results.push((steps, error)); + } + + for pair in results.windows(2) { + assert!( + pair[1].0 > pair[0].0, + "a tighter tolerance should take more steps: {results:?}" + ); + assert!( + pair[1].1 < pair[0].1, + "a tighter tolerance should end up closer: {results:?}" + ); + } + } + + #[test] + fn test_accuracy_preset_can_be_overridden() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(Duration::from_seconds(period())) + .accuracy(Accuracy::High) + .r_tol(1e-6); + // Runs, and takes fewer steps than High would on its own. + let loosened = propagation.run().states.len(); + + let mut tight = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(Duration::from_seconds(period())) + .accuracy(Accuracy::High); + assert!(loosened < tight.run().states.len()); + } + + #[test] + fn test_method_can_be_replaced() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .method(DormandPrince45::new()) + .duration(Duration::from_seconds(period())) + .accuracy(Accuracy::High); + let solution = propagation.run(); + let error = + (Orbit::position(solution.states.last().unwrap()) - Orbit::position(&leo())).norm(); + assert!(error < 1e-4, "position error {error:e} m"); + } + + #[test] + fn test_until_matches_duration() { + let duration = Duration::from_seconds(period()); + let mut by_duration = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(duration) + .accuracy(Accuracy::Medium); + let a = *by_duration.run().states.last().unwrap(); + + let mut by_epoch = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .until(epoch() + duration) + .accuracy(Accuracy::Medium); + let b = *by_epoch.run().states.last().unwrap(); + + assert_eq!(a, b); + } + + #[test] + fn test_dense_output_is_available_after_the_solve() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(Duration::from_seconds(period())) + .accuracy(Accuracy::High); + let solution = propagation.run(); + + // Half a period from a circular-ish LEO puts us on the far side. + let half = solution.interpolate(period() / 2.0); + let r0 = Orbit::position(&leo()); + let r_half = Orbit::position(&half); + assert!( + r0.dot(&r_half) < 0.0, + "half a period should be on the opposite side of the primary" + ); + } + + #[test] + #[should_panic(expected = "set a duration or an end epoch")] + fn test_running_without_a_duration_panics() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()); + let _ = propagation.run(); + } + + #[test] + fn test_assembly_reports_model_names() { + let mut propagation = Propagation::new(epoch(), Orbit, leo(), &two_body()) + .duration(Duration::from_seconds(60.0)); + let solution = propagation.run(); + assert_eq!(solution.ode.params.names(), ["two_body"]); + } + + /// One propagation in a vector dual gives the whole 6x6 state transition + /// matrix. Two-body motion is Hamiltonian, so its flow is volume-preserving + /// and `det(Phi) = 1` exactly, for any interval. + #[cfg(feature = "autodiff")] + #[test] + fn test_state_transition_matrix_from_one_propagation() { + use nalgebra::Matrix6; + use num_dual::DualSVec64; + + let specs = two_body(); + let (end, stm): (Vector6, Matrix6) = num_dual::jacobian( + |y: Vector6>| { + type T = DualSVec64<6>; + let mut propagation = Propagation::new(epoch(), Orbit, y, &specs) + .duration(Duration::from_seconds(period())) + .accuracy(Accuracy::High); + *propagation.run().states.last().unwrap() as SVector + }, + &leo(), + ); + + assert_relative_eq!( + (Orbit::position(&end) - Orbit::position(&leo())).norm(), + 0.0, + epsilon = 1e-4 + ); + assert_relative_eq!(stm.determinant(), 1.0, max_relative = 1e-9); + } +} diff --git a/src/state.rs b/src/state.rs index 76d09ca..5b27f23 100644 --- a/src/state.rs +++ b/src/state.rs @@ -36,16 +36,176 @@ pub struct Component { pub width: usize, } +/// Compares two strings in a const context. +const fn str_eq(a: &str, b: &str) -> bool { + let (a, b) = (a.as_bytes(), b.as_bytes()); + if a.len() != b.len() { + return false; + } + let mut i = 0; + while i < a.len() { + if a[i] != b[i] { + return false; + } + i += 1; + } + true +} + +/// The index at which `name` starts, if `components` has it. +/// +/// A `const fn` so that [`Layout`]'s optional index constants can default to it, +/// which is what makes them correct for every layout without the macro emitting +/// anything per component. +pub const fn offset_of(components: &[Component], name: &str) -> Option { + let mut offset = 0; + let mut i = 0; + while i < components.len() { + if str_eq(components[i].name, name) { + return Some(offset); + } + offset += components[i].width; + i += 1; + } + None +} + /// A state layout of dimension `D`. /// /// Layouts are zero-sized marker types, so the accessors monomorphize to a /// direct offset. Define one with [`define_state!`](crate::define_state) rather than by hand. +/// +/// # Required and optional access +/// +/// A model that cannot work without a quantity should bound on the capability +/// trait — `L: HasPosition` — and use [`HasPosition::position`], which is a +/// plain offset with no branch. +/// +/// A model that can *use* a quantity but does not need it should bound on +/// `Layout` alone and use the `try_` accessors here. The `_AT` constants are +/// known at compile time, so the branch folds away at monomorphization and the +/// absent case costs nothing. pub trait Layout: Copy + Clone + Debug + Default { /// The layout's name, as written at its definition site const NAME: &'static str; /// The components, in index order. const COMPONENTS: &'static [Component]; + + /// Index of the position, if this layout has one. + const POSITION_AT: Option = offset_of(Self::COMPONENTS, "position"); + /// Index of the velocity, if this layout has one. + const VELOCITY_AT: Option = offset_of(Self::COMPONENTS, "velocity"); + /// Index of the attitude, if this layout has one. + const ATTITUDE_AT: Option = offset_of(Self::COMPONENTS, "attitude"); + /// Index of the angular rate, if this layout has one. + const ANGULAR_RATE_AT: Option = offset_of(Self::COMPONENTS, "angular_rate"); + /// Index of the mass, if this layout has one. + const MASS_AT: Option = offset_of(Self::COMPONENTS, "mass"); + /// Index of the stored energy, if this layout has one. + const ENERGY_AT: Option = offset_of(Self::COMPONENTS, "energy"); + + /// The position in meters, if this layout has one. + #[inline] + fn try_position(y: &SVector) -> Option> { + Self::POSITION_AT.map(|at| y.fixed_rows::<3>(at).into_owned()) + } + + /// The velocity in meters per second, if this layout has one. + #[inline] + fn try_velocity(y: &SVector) -> Option> { + Self::VELOCITY_AT.map(|at| y.fixed_rows::<3>(at).into_owned()) + } + + /// The attitude as Modified Rodrigues Parameters, if this layout has one. + #[inline] + fn try_attitude(y: &SVector) -> Option> { + Self::ATTITUDE_AT.map(|at| y.fixed_rows::<3>(at).into_owned()) + } + + /// The body angular rate in radians per second, if this layout has one. + #[inline] + fn try_angular_rate(y: &SVector) -> Option> { + Self::ANGULAR_RATE_AT.map(|at| y.fixed_rows::<3>(at).into_owned()) + } + + /// The mass in kilograms, if this layout has one. + #[inline] + fn try_mass(y: &SVector) -> Option { + Self::MASS_AT.map(|at| y[at]) + } + + /// The stored energy in joules, if this layout has one. + #[inline] + fn try_energy(y: &SVector) -> Option { + Self::ENERGY_AT.map(|at| y[at]) + } + + /// Writes the position, reporting whether this layout has one. + #[inline] + fn try_set_position(y: &mut SVector, value: Vector3) -> bool { + set_vector(y, Self::POSITION_AT, value) + } + + /// Writes the velocity, reporting whether this layout has one. + #[inline] + fn try_set_velocity(y: &mut SVector, value: Vector3) -> bool { + set_vector(y, Self::VELOCITY_AT, value) + } + + /// Writes the attitude, reporting whether this layout has one. + #[inline] + fn try_set_attitude(y: &mut SVector, value: Vector3) -> bool { + set_vector(y, Self::ATTITUDE_AT, value) + } + + /// Writes the angular rate, reporting whether this layout has one. + #[inline] + fn try_set_angular_rate(y: &mut SVector, value: Vector3) -> bool { + set_vector(y, Self::ANGULAR_RATE_AT, value) + } + + /// Writes the mass, reporting whether this layout has one. + #[inline] + fn try_set_mass(y: &mut SVector, value: T) -> bool { + set_scalar(y, Self::MASS_AT, value) + } + + /// Writes the stored energy, reporting whether this layout has one. + #[inline] + fn try_set_energy(y: &mut SVector, value: T) -> bool { + set_scalar(y, Self::ENERGY_AT, value) + } +} + +#[inline] +fn set_vector( + y: &mut SVector, + at: Option, + value: Vector3, +) -> bool { + match at { + Some(at) => { + y.fixed_rows_mut::<3>(at).copy_from(&value); + true + } + None => false, + } +} + +#[inline] +fn set_scalar( + y: &mut SVector, + at: Option, + value: T, +) -> bool { + match at { + Some(at) => { + y[at] = value; + true + } + None => false, + } } /// Defines a capability trait for a three-element component. @@ -471,6 +631,45 @@ mod tests { assert_eq!(MassOnly::mass(&y), 42.0); } + #[test] + fn test_optional_indices_agree_with_required_ones() { + assert_eq!( + >::POSITION_AT, + Some(>::POSITION) + ); + assert_eq!( + >::ENERGY_AT, + Some(>::ENERGY) + ); + + // And report absence for what a layout does not carry. + assert_eq!(>::MASS_AT, None); + assert_eq!(>::ATTITUDE_AT, None); + assert_eq!(>::MASS_AT, Some(6)); + assert_eq!(>::ENERGY_AT, None); + } + + /// The optional indices are const, so they are usable in a const context and + /// the absent branch folds away rather than costing a runtime check. + #[test] + fn test_optional_indices_are_const() { + const MASS: Option = >::MASS_AT; + assert_eq!(MASS, Some(12)); + } + + #[test] + fn test_try_accessors_report_absence() { + let mut orbit = Vector6::::zeros(); + assert_eq!(Orbit::try_mass(&orbit), None); + assert!(!Orbit::try_set_mass(&mut orbit, 500.0)); + assert_eq!(orbit, Vector6::zeros()); + + let mut full = SVector::::zeros(); + assert!(Full::try_set_mass(&mut full, 500.0)); + assert_eq!(Full::try_mass(&full), Some(500.0)); + assert_eq!(full[12], 500.0); + } + /// Accessors have to work at a dual scalar, or the layout is unusable in the /// differentiated path. #[cfg(feature = "autodiff")]