Some initial benchmarks and utilities, including keplerian elements. More typed coordinate stuff to come

benches/kepler.rs

@@ -0,0 +1,128 @@
+//! Element conversions and the two-body relations.
+//!
+//! These run per measurement in a filter and per step in some force models, so a
+//! slow conversion is not a rounding error in the budget. The point of tracking
+//! them is to notice a regression, not to hit a number.
+
+use std::hint::black_box;
+
+use astrokit::kepler::{self, Degeneracy, Keplerian};
+use criterion::{criterion_group, criterion_main, Criterion};
+use nalgebra::Vector3;
+
+const MU: f64 = 3.986004415e14;
+
+/// Inclined, near-circular LEO. Just outside the degeneracy tolerance, so it
+/// takes the general branch.
+fn leo() -> (Vector3<f64>, Vector3<f64>) {
+    (
+        Vector3::new(
+            4.263868426884883e6,
+            5.146189057155391e6,
+            1.1310208421331816e6,
+        ),
+        Vector3::new(-5923.454461876975, 4496.802639690076, 1870.3893008991558),
+    )
+}
+
+/// GEO: circular and equatorial, so it takes the doubly-degenerate branch.
+fn geo() -> (Vector3<f64>, Vector3<f64>) {
+    let radius = 4.2164e7;
+    (
+        Vector3::new(radius, 0.0, 0.0),
+        Vector3::new(0.0, kepler::circular_speed(radius, MU), 0.0),
+    )
+}
+
+fn bench_scalar_relations(c: &mut Criterion) {
+    let mut group = c.benchmark_group("kepler/relations");
+    let sma = 6.7781363e6;
+
+    group.bench_function("period", |b| {
+        b.iter(|| kepler::period(black_box(sma), black_box(MU)))
+    });
+    group.bench_function("semi_major_axis", |b| {
+        b.iter(|| kepler::semi_major_axis(black_box(5553.6), black_box(MU)))
+    });
+    group.bench_function("mean_motion", |b| {
+        b.iter(|| kepler::mean_motion(black_box(sma), black_box(MU)))
+    });
+    group.finish();
+}
+
+fn bench_anomalies(c: &mut Criterion) {
+    let mut group = c.benchmark_group("kepler/anomalies");
+
+    group.bench_function("eccentric_from_true", |b| {
+        b.iter(|| kepler::eccentric_from_true(black_box(1.2), black_box(0.3)))
+    });
+    group.bench_function("mean_from_eccentric", |b| {
+        b.iter(|| kepler::mean_from_eccentric(black_box(1.1), black_box(0.3)))
+    });
+
+    // The only iterative routine here, and the one worth watching. Near-circular
+    // converges in a step or two; high eccentricity is the expensive case.
+    for eccentricity in [0.001, 0.3, 0.9] {
+        group.bench_function(format!("eccentric_from_mean/e={eccentricity}"), |b| {
+            b.iter(|| kepler::eccentric_from_mean(black_box(1.0), black_box(eccentricity)))
+        });
+    }
+    group.finish();
+}
+
+fn bench_conversions(c: &mut Criterion) {
+    let mut group = c.benchmark_group("kepler/conversions");
+
+    let (r, v) = leo();
+    group.bench_function("from_cartesian/leo", |b| {
+        b.iter(|| Keplerian::from_cartesian(black_box(r), black_box(v), black_box(MU)))
+    });
+
+    let (r_geo, v_geo) = geo();
+    group.bench_function("from_cartesian/geo", |b| {
+        b.iter(|| Keplerian::from_cartesian(black_box(r_geo), black_box(v_geo), black_box(MU)))
+    });
+
+    let elements = Keplerian {
+        semi_major_axis: 2.0e7,
+        eccentricity: 0.3,
+        inclination: 0.8,
+        raan: 0.5,
+        aop: 1.0,
+        true_anomaly: 2.0,
+        degeneracy: Degeneracy::None,
+    };
+    group.bench_function("to_cartesian", |b| {
+        b.iter(|| black_box(elements).to_cartesian(black_box(MU)))
+    });
+
+    group.bench_function("round_trip", |b| {
+        b.iter(|| {
+            let elements = Keplerian::from_cartesian(black_box(r), black_box(v), MU);
+            black_box(elements.to_cartesian(MU))
+        })
+    });
+    group.finish();
+}
+
+fn bench_invariants(c: &mut Criterion) {
+    let mut group = c.benchmark_group("kepler/invariants");
+    let (r, v) = leo();
+
+    group.bench_function("specific_energy", |b| {
+        b.iter(|| kepler::specific_energy(black_box(r), black_box(v), black_box(MU)))
+    });
+    group.bench_function("specific_angular_momentum", |b| {
+        b.iter(|| kepler::specific_angular_momentum(black_box(r), black_box(v)))
+    });
+    group.finish();
+}
+
+criterion_group!(
+    benches,
+    bench_scalar_relations,
+    bench_anomalies,
+    bench_conversions,
+    bench_invariants
+);
+criterion_main!(benches);

benches/propagation.rs

@@ -0,0 +1,187 @@
+//! Two-body propagation, and the derivative evaluation underneath it.
+//!
+//! The derivative benchmarks are the ones that matter most: a propagation
+//! evaluates them thousands of times per day of arc, so anything wasteful in
+//! assembly or in a model shows up here before it shows up in a wall-clock
+//! figure. They also isolate the cost of a layout, which should be nil — the
+//! `Orbit`, `OrbitMass` and `Full` numbers ought to differ only by the work of
+//! moving more elements around.
+
+use std::hint::black_box;
+
+use astrokit::dynamics::{self, Context};
+use astrokit::kepler;
+use astrokit::prelude::*;
+use astrokit::propagation::{state_derivative, Assembly};
+use criterion::{criterion_group, criterion_main, Criterion};
+use nalgebra::Vector3;
+
+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 position() -> Vector3<f64> {
+    Vector3::new(
+        4.263868426884883e6,
+        5.146189057155391e6,
+        1.1310208421331816e6,
+    )
+}
+
+fn velocity() -> Vector3<f64> {
+    Vector3::new(-5923.454461876975, 4496.802639690076, 1870.3893008991558)
+}
+
+fn two_body() -> Vec<Dynamics> {
+    vec![TwoBody { mu: MU }]
+}
+
+/// One derivative evaluation, per layout. The layout should cost nothing.
+fn bench_derivative(c: &mut Criterion) {
+    let mut group = c.benchmark_group("derivative");
+
+    let assembly = Assembly::new(
+        dynamics::build::<f64, 6, Orbit>(&two_body()),
+        Context::new(epoch()),
+    );
+    let y = Orbit.with_position(position()).with_velocity(velocity());
+    group.bench_function("two_body/Orbit", |b| {
+        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &assembly))
+    });
+
+    let assembly = Assembly::new(
+        dynamics::build::<f64, 7, OrbitMass>(&two_body()),
+        Context::new(epoch()),
+    );
+    let y = OrbitMass
+        .with_position(position())
+        .with_velocity(velocity())
+        .with_mass(500.0);
+    group.bench_function("two_body/OrbitMass", |b| {
+        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &assembly))
+    });
+
+    let assembly = Assembly::new(
+        dynamics::build::<f64, 14, Full>(&two_body()),
+        Context::new(epoch()),
+    );
+    let y = Full
+        .with_position(position())
+        .with_velocity(velocity())
+        .with_mass(500.0);
+    group.bench_function("two_body/Full", |b| {
+        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &assembly))
+    });
+
+    group.finish();
+}
+
+/// A full propagation at each accuracy preset, over one day.
+fn bench_propagation(c: &mut Criterion) {
+    let mut group = c.benchmark_group("propagation/one_day");
+    group.sample_size(20);
+
+    let forces = two_body();
+    for accuracy in [Accuracy::Low, Accuracy::Medium, Accuracy::High] {
+        group.bench_function(format!("{accuracy:?}"), |b| {
+            b.iter(|| {
+                let y0 = Orbit.with_position(position()).with_velocity(velocity());
+                let mut propagation = Propagation::new(epoch(), y0, &forces)
+                    .duration(Duration::from_seconds(86400.0))
+                    .accuracy(accuracy);
+                black_box(propagation.run().states.len())
+            })
+        });
+    }
+    group.finish();
+}
+
+/// One orbital period, which is the unit most of the tests work in.
+fn bench_one_period(c: &mut Criterion) {
+    let forces = two_body();
+    let period = Duration::from_seconds(kepler::period(SMA, MU));
+
+    c.bench_function("propagation/one_period", |b| {
+        b.iter(|| {
+            let y0 = Orbit.with_position(position()).with_velocity(velocity());
+            let mut propagation = Propagation::new(epoch(), y0, &forces)
+                .duration(period)
+                .accuracy(Accuracy::Medium);
+            black_box(propagation.run().states.len())
+        })
+    });
+}
+
+/// The cost of a state transition matrix relative to a plain propagation.
+///
+/// Stage 1 measured 10.6x for a 6x6 on a two-body arc. Central differences would
+/// need thirteen propagations and lose about half the digits, so anything near
+/// that ratio is a win; this tracks it as the force models get heavier.
+#[cfg(feature = "autodiff")]
+fn bench_state_transition_matrix(c: &mut Criterion) {
+    use nalgebra::{Matrix6, Vector6};
+    use num_dual::DualSVec64;
+
+    let mut group = c.benchmark_group("propagation/stm");
+    group.sample_size(20);
+
+    let forces = two_body();
+    let period = Duration::from_seconds(kepler::period(SMA, MU));
+    let state = Vector6::new(
+        position()[0],
+        position()[1],
+        position()[2],
+        velocity()[0],
+        velocity()[1],
+        velocity()[2],
+    );
+
+    group.bench_function("plain", |b| {
+        b.iter(|| {
+            let y0 = Orbit.state(state);
+            let mut propagation = Propagation::new(epoch(), y0, &forces)
+                .duration(period)
+                .accuracy(Accuracy::Medium);
+            black_box(*propagation.run().states.last().unwrap())
+        })
+    });
+
+    group.bench_function("with_stm", |b| {
+        b.iter(|| {
+            let (_, stm): (Vector6<f64>, Matrix6<f64>) = num_dual::jacobian(
+                |y: Vector6<DualSVec64<6>>| {
+                    let mut propagation = Propagation::new(epoch(), Orbit.state(y), &forces)
+                        .duration(period)
+                        .accuracy(Accuracy::Medium);
+                    *propagation.run().states.last().unwrap()
+                },
+                &state,
+            );
+            black_box(stm)
+        })
+    });
+
+    group.finish();
+}
+
+#[cfg(feature = "autodiff")]
+criterion_group!(
+    benches,
+    bench_derivative,
+    bench_propagation,
+    bench_one_period,
+    bench_state_transition_matrix
+);
+
+#[cfg(not(feature = "autodiff"))]
+criterion_group!(
+    benches,
+    bench_derivative,
+    bench_propagation,
+    bench_one_period
+);
+
+criterion_main!(benches);