astrokit/benches/kepler.rs

//! 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::body::Body;
use astrokit::cartesian::Cartesian;
use astrokit::equinoctial::Equinoctial;
use astrokit::frames::Frame;
use astrokit::kepler::{self, Degeneracy, Keplerian};
use criterion::{criterion_group, criterion_main, Criterion};
use nalgebra::Vector3;

const MU: f64 = Body::EARTH.mu;

/// Inclined, near-circular LEO. Just outside the degeneracy tolerance, so it
/// takes the general branch.
fn leo() -> Cartesian<f64> {
    Cartesian::new(
        Vector3::new(
            4.263868426884883e6,
            5.146189057155391e6,
            1.1310208421331816e6,
        ),
        Vector3::new(-5923.454461876975, 4496.802639690076, 1870.3893008991558),
        Frame::EARTH_J2000,
    )
}

/// GEO: circular and equatorial, so it takes the doubly-degenerate branch.
fn geo() -> Cartesian<f64> {
    let radius = 4.2164e7;
    Cartesian::new(
        Vector3::new(radius, 0.0, 0.0),
        Vector3::new(0.0, kepler::circular_speed(radius, MU), 0.0),
        Frame::EARTH_J2000,
    )
}

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 leo = leo();
    group.bench_function("from_cartesian/leo", |b| {
        b.iter(|| Keplerian::from_cartesian(&black_box(leo)))
    });

    let geo = geo();
    group.bench_function("from_cartesian/geo", |b| {
        b.iter(|| Keplerian::from_cartesian(&black_box(geo)))
    });

    let elements = Keplerian {
        semi_major_axis: 2.0e7,
        eccentricity: 0.3,
        inclination: 0.8,
        raan: 0.5,
        aop: 1.0,
        true_anomaly: 2.0,
        frame: Frame::EARTH_J2000,
        degeneracy: Degeneracy::None,
    };
    group.bench_function("to_cartesian", |b| {
        b.iter(|| black_box(elements).to_cartesian())
    });

    group.bench_function("round_trip", |b| {
        b.iter(|| Keplerian::from_cartesian(&black_box(leo)).to_cartesian())
    });
    group.finish();
}

fn bench_invariants(c: &mut Criterion) {
    let mut group = c.benchmark_group("kepler/invariants");
    let leo = leo();

    group.bench_function("specific_energy", |b| {
        b.iter(|| black_box(leo).specific_energy())
    });
    group.bench_function("specific_angular_momentum", |b| {
        b.iter(|| black_box(leo).specific_angular_momentum())
    });
    group.finish();
}

/// Equinoctial against Keplerian on the same states.
///
/// The interesting comparison is GEO: Keplerian takes its doubly-degenerate
/// branch there and equinoctial takes no branch at all.
fn bench_equinoctial(c: &mut Criterion) {
    let mut group = c.benchmark_group("equinoctial");

    let leo = leo();
    let geo = geo();

    group.bench_function("from_cartesian/leo", |b| {
        b.iter(|| Equinoctial::from_cartesian(&black_box(leo)))
    });
    group.bench_function("from_cartesian/geo", |b| {
        b.iter(|| Equinoctial::from_cartesian(&black_box(geo)))
    });

    let elements = Equinoctial::from_cartesian(&leo);
    group.bench_function("to_cartesian", |b| {
        b.iter(|| black_box(elements).to_cartesian())
    });
    group.bench_function("round_trip", |b| {
        b.iter(|| Equinoctial::from_cartesian(&black_box(leo)).to_cartesian())
    });

    let keplerian = leo.keplerian();
    group.bench_function("from_keplerian", |b| {
        b.iter(|| Equinoctial::from_keplerian(&black_box(keplerian)))
    });
    group.finish();
}

criterion_group!(
    benches,
    bench_scalar_relations,
    bench_anomalies,
    bench_conversions,
    bench_invariants,
    bench_equinoctial
);
criterion_main!(benches);