deep_causality_physics 0.8.3

/*
 * SPDX-License-Identifier: MIT
 * Copyright (c) 2023 - 2026. The DeepCausality Authors and Contributors. All Rights Reserved.
 */

//! Exact 3-D two-body propagation by **Kustaanheimo–Stiefel (KS) regularisation** — the
//! singularity-free, perturbation-ready generalisation of the planar
//! [`TwoBodyPropagator`](crate::TwoBodyPropagator) (FS-1). This is the B1 conformal core of the
//! perturbed-conformal trajectory split (Gap-3 Resolution 1); the plasma-blackout reentry arc is
//! highly eccentric / near-radial, exactly where an orbital-elements propagator is ill-conditioned and
//! KS is not.
//!
//! # Method
//! KS lifts the 3-D Kepler position `r ∈ ℝ³` to a 4-vector `u ∈ ℝ⁴` (`|r| = |u|²`) and reparametrises
//! time by the fictitious time `s` with `dt = |u|² ds`. In `s` the **bound** motion is the constant
//! 4-D harmonic oscillator `u'' + ω₀² u = 0`, `ω₀² = −h/2` (`h < 0` the specific energy), so
//! `u(s) = u₀·cos(ω₀ s) + (u₀'/ω₀)·sin(ω₀ s)` — a constant-generator matrix exponential, exact to
//! round-off. Physical time is the closed-form integral `t(s) = ∫₀ˢ |u|² ds'`; a monotone Newton solve
//! inverts `t(s) = dt` (a well-conditioned Kepler-equation analogue).
//!
//! # References
//! * Stiefel, E. L. & Scheifele, G., *Linear and Regular Celestial Mechanics*, Springer (1971).
//! * Battin, R. H., *An Introduction to the Mathematics and Methods of Astrodynamics*, AIAA (1999).

use crate::PhysicsError;
use deep_causality_algebra::RealField;
use deep_causality_linear::dot_n;
use deep_causality_num::FromPrimitive;

/// An exact 3-D two-body (Kepler) propagator via KS regularisation, built from a single physical
/// state. Propagation is the constant-generator matrix exponential in the KS fictitious time; it is
/// exact to floating-point round-off for any `dt`, with no step-size error.
#[derive(Clone, Copy, Debug)]
pub struct KsPropagator<R> {
    gravitational_parameter: R,
    /// KS oscillator frequency `ω₀ = √(−h/2)`.
    omega: R,
    /// The epoch KS coordinate `A = u₀` (with `|A|² = |r₀|`).
    a: [R; 4],
    /// The scaled epoch KS velocity `B = u₀'/ω₀`.
    b: [R; 4],
    /// Cached `|A|²`, `|B|²`, `A·B` for the closed-form `t(s)`.
    aa: R,
    bb: R,
    ab: R,
    /// Semi-major axis `a = μ/(4 ω₀²)`.
    semi_major_axis: R,
}

impl<R> KsPropagator<R>
where
    R: RealField + FromPrimitive,
{
    /// Build the propagator from a 3-D physical state `(position, velocity)` about a primary with
    /// gravitational parameter `gm = GM`.
    ///
    /// # Errors
    /// [`PhysicsError::PhysicalInvariantBroken`] if `gm ≤ 0` or the state is not bound (`energy ≥ 0`);
    /// [`PhysicsError::Singularity`] on a non-positive radius; numeric-conversion failures.
    pub fn from_state(position: [R; 3], velocity: [R; 3], gm: R) -> Result<Self, PhysicsError> {
        if gm <= R::zero() {
            return Err(PhysicsError::PhysicalInvariantBroken(
                "Gravitational parameter GM must be positive".into(),
            ));
        }
        let two = Self::lit(2.0)?;
        let four = Self::lit(4.0)?;

        let radius =
            (position[0] * position[0] + position[1] * position[1] + position[2] * position[2])
                .sqrt();
        if radius <= R::zero() {
            return Err(PhysicsError::Singularity("Radius must be positive".into()));
        }
        let v2 = velocity[0] * velocity[0] + velocity[1] * velocity[1] + velocity[2] * velocity[2];
        let energy = v2 / two - gm / radius;
        if energy >= R::zero() {
            return Err(PhysicsError::PhysicalInvariantBroken(
                "State is not a bound orbit (energy >= 0); only ellipses are supported".into(),
            ));
        }
        let omega = (-energy / two).sqrt();
        let semi_major_axis = gm / (four * omega * omega);

        // KS lift r -> u (gauge branch keeps the leading component well away from zero).
        let a = Self::ks_lift(position, radius);
        // u' = ½ L(u)ᵀ (v; 0); then B = u'/ω₀.
        let uprime = Self::l_transpose_times_v(&a, velocity);
        let half = R::one() / two;
        let b = [
            uprime[0] * half / omega,
            uprime[1] * half / omega,
            uprime[2] * half / omega,
            uprime[3] * half / omega,
        ];

        let aa = Self::dot4(&a, &a);
        let bb = Self::dot4(&b, &b);
        let ab = Self::dot4(&a, &b);

        Ok(Self {
            gravitational_parameter: gm,
            omega,
            a,
            b,
            aa,
            bb,
            ab,
            semi_major_axis,
        })
    }

    /// The state `(position, velocity)` at time `dt` after the epoch — exact for any `dt`.
    ///
    /// # Errors
    /// Numeric-conversion failures.
    pub fn propagate(&self, dt: R) -> Result<([R; 3], [R; 3]), PhysicsError> {
        let s = self.solve_fictitious_time(dt)?;
        let phi = self.omega * s;
        let (cp, sp) = (phi.cos(), phi.sin());

        // u(s) and u'(s).
        let mut u = [R::zero(); 4];
        let mut up = [R::zero(); 4];
        for i in 0..4 {
            u[i] = self.a[i] * cp + self.b[i] * sp;
            up[i] = self.omega * (-self.a[i] * sp + self.b[i] * cp);
        }
        let radius = Self::dot4(&u, &u);

        // r = first three components of L(u)·u; dr/ds = 2·L(u)·u'; v = (dr/ds)/radius.
        let r3 = Self::l_times(&u, &u);
        let drds = Self::l_times(&u, &up);
        let two = Self::lit(2.0)?;
        let velocity = [
            two * drds[0] / radius,
            two * drds[1] / radius,
            two * drds[2] / radius,
        ];
        Ok((r3, velocity))
    }

    /// Semi-major axis `a` (m).
    pub fn semi_major_axis(&self) -> R {
        self.semi_major_axis
    }
    /// Gravitational parameter `GM` (m³/s²).
    pub fn gravitational_parameter(&self) -> R {
        self.gravitational_parameter
    }
    /// Mean motion `n = √(GM/a³)` (rad/s).
    pub fn mean_motion(&self) -> R {
        (self.gravitational_parameter
            / (self.semi_major_axis * self.semi_major_axis * self.semi_major_axis))
            .sqrt()
    }
    /// Orbital period `T = 2π/n` (s).
    ///
    /// # Errors
    /// Numeric-conversion failures.
    pub fn period(&self) -> Result<R, PhysicsError> {
        let two = Self::lit(2.0)?;
        Ok(two * R::pi() / self.mean_motion())
    }

    // ── helpers ─────────────────────────────────────────────────────────────────────────────────────

    /// `t(s) = ∫₀ˢ |u|² ds'` and the radius `|u(s)|²` at `s` — the value and derivative for Newton.
    fn t_and_radius(&self, s: R) -> Result<(R, R), PhysicsError> {
        let two = Self::lit(2.0)?;
        let four = Self::lit(4.0)?;
        let phi = self.omega * s;
        let (cp, sp) = (phi.cos(), phi.sin());
        let sin2 = two * sp * cp; // sin 2φ
        let cos2 = cp * cp - sp * sp; // cos 2φ
        let c_lin = (self.aa + self.bb) / two;
        let t = c_lin * s
            + (self.aa - self.bb) * sin2 / (four * self.omega)
            + self.ab * (R::one() - cos2) / (two * self.omega);
        let radius = self.aa * cp * cp + two * self.ab * cp * sp + self.bb * sp * sp;
        Ok((t, radius))
    }

    /// Invert `t(s) = dt` by a bracketed Newton iteration (`dt/ds = |u|² > 0`).
    ///
    /// # Two convergence tests, because each is unreachable in the other's regime
    ///
    /// The Newton step is `d = (t − dt) / radius`, so the step test and the residual test differ by
    /// a factor of `radius` and neither covers the whole input space:
    ///
    /// * **Near periapsis of an eccentric orbit** `radius` is small, so `d` is inflated and the
    ///   step test asks for an accuracy the arithmetic cannot deliver. Measured on
    ///   `e = 0.9999, a = 7000 km`, where the periapsis radius is **700 m**: the step test was
    ///   unreachable by factors of 9 to 1700 while `t` was already correct to the last bit.
    /// * **On a large, slow orbit** `radius` is huge — `1e9 m` — so a step of `1e-15` corresponds
    ///   to a *time* residual of `1e-6 s`, and the residual test is the one that cannot be met.
    ///
    /// Accepting on **either** covers both, and each is a sound statement of convergence on its own:
    /// one says the root has stopped moving, the other that `t(s)` equals `dt` to working precision.
    ///
    /// Newton proposals are constrained to a bracket of the root. The bracket follows directly
    /// from the monotonicity of `t(s)`, and its bisection fallback makes convergence independent of
    /// last-bit differences in the platform implementations of `sin` and `cos`.
    fn solve_fictitious_time(&self, dt: R) -> Result<R, PhysicsError> {
        let two = Self::lit(2.0)?;
        let c_lin = (self.aa + self.bb) / two; // the s-average of |u|²; exact asymptotic slope
        let mut s = dt / c_lin;
        // Scaled by the scalar's own precision, as `two_body::solve_kepler` is: a fixed `1e-15` is
        // below what `f32` can reach at all and meaningless in `BFloat16`.
        let tol = Self::lit(8.0)? * R::epsilon();
        let bound = tol * (dt.abs() + R::one());

        // `t(s)` is strictly monotone. Keep every Newton proposal inside a bracket so a last-bit
        // difference in sin/cos cannot send it into a different oscillation.
        let zero = R::zero();
        let mut lower = if dt >= zero { zero } else { s };
        let mut upper = if dt >= zero { s } else { zero };
        let mut lower_t = self.t_and_radius(lower)?.0;
        let mut upper_t = self.t_and_radius(upper)?.0;
        while lower_t > dt {
            upper = lower;
            lower *= two;
            lower_t = self.t_and_radius(lower)?.0;
        }
        while upper_t < dt {
            lower = upper;
            upper *= two;
            upper_t = self.t_and_radius(upper)?.0;
        }

        for _ in 0..100 {
            let (t, radius) = self.t_and_radius(s)?;
            let residual = (t - dt).abs();
            if residual <= bound {
                return Ok(s);
            }
            if t < dt {
                lower = s;
            } else {
                upper = s;
            }
            let d = (t - dt) / radius;
            let newton = s - d;
            s = if newton <= lower || newton >= upper {
                (lower + upper) / two
            } else {
                newton
            };
            if d.abs() < tol * (s.abs() + R::one()) {
                return Ok(s);
            }
        }
        Err(PhysicsError::NotConverged(
            "the fictitious-time inversion t(s) = dt did not converge in 100 iterations: no \
             iterate inverted t(s) = dt to the working tolerance"
                .into(),
        ))
    }

    /// KS lift `r → u` with a gauge choice that keeps the pivot component large (Stiefel–Scheifele).
    fn ks_lift(r: [R; 3], radius: R) -> [R; 4] {
        let two = Self::lit(2.0).expect("2.0 lifts into every real field");
        let half = R::one() / two;
        if r[0] >= R::zero() {
            let u1 = (half * (radius + r[0])).sqrt();
            [u1, r[1] / (two * u1), r[2] / (two * u1), R::zero()]
        } else {
            let u2 = (half * (radius - r[0])).sqrt();
            [r[1] / (two * u2), u2, R::zero(), r[2] / (two * u2)]
        }
    }

    /// First three components of `L(u)·w` (the KS matrix acting on a 4-vector).
    fn l_times(u: &[R; 4], w: &[R; 4]) -> [R; 3] {
        [
            u[0] * w[0] - u[1] * w[1] - u[2] * w[2] + u[3] * w[3],
            u[1] * w[0] + u[0] * w[1] - u[3] * w[2] - u[2] * w[3],
            u[2] * w[0] + u[3] * w[1] + u[0] * w[2] + u[1] * w[3],
        ]
    }

    /// `L(u)ᵀ · (v; 0)` — used to recover the KS velocity from the physical velocity.
    fn l_transpose_times_v(u: &[R; 4], v: [R; 3]) -> [R; 4] {
        [
            u[0] * v[0] + u[1] * v[1] + u[2] * v[2],
            -u[1] * v[0] + u[0] * v[1] + u[3] * v[2],
            -u[2] * v[0] - u[3] * v[1] + u[0] * v[2],
            u[3] * v[0] - u[2] * v[1] + u[1] * v[2],
        ]
    }

    /// The inner product of two KS 4-vectors.
    ///
    /// Dispatches to `deep_causality_linear::dot_n` (`unified-math-next` task 6.7). The
    /// const-generic form rather than the slice form, because `[R; 4]` already proves the lengths
    /// agree and the slice form's refusal would be an error arm no input here can reach.
    fn dot4(x: &[R; 4], y: &[R; 4]) -> R {
        dot_n(x, y)
    }

    fn lit(x: f64) -> Result<R, PhysicsError> {
        R::from_f64(x)
            .ok_or_else(|| PhysicsError::NumericalInstability("R::from_f64 failed".into()))
    }
}

/// One **Strang-split** perturbed step (the B1 between-step perturbation hook): half-kick the velocity
/// with the caller-supplied non-conformal Cartesian acceleration, drift exactly along the KS conformal
/// core for `dt`, then half-kick again. The perturbation is never expressed inside the KS algebra
/// (FS-2). 2nd-order accurate; the exact core is untouched when `accel` returns zero.
///
/// `accel(position, velocity) -> acceleration` is the non-conformal force per unit mass (e.g. the ④
/// aero force + J2).
///
/// # Errors
/// As [`KsPropagator::from_state`] / [`KsPropagator::propagate`] (the half-kicked state must stay bound).
pub fn ks_strang_step<R, A>(
    position: [R; 3],
    velocity: [R; 3],
    gm: R,
    dt: R,
    accel: A,
) -> Result<([R; 3], [R; 3]), PhysicsError>
where
    R: RealField + FromPrimitive,
    A: Fn([R; 3], [R; 3]) -> [R; 3],
{
    let two = R::from_f64(2.0)
        .ok_or_else(|| PhysicsError::NumericalInstability("R::from_f64 failed".into()))?;
    let half_dt = dt / two;

    // Half-kick in physical Cartesian velocity.
    let a0 = accel(position, velocity);
    let v_half = [
        velocity[0] + a0[0] * half_dt,
        velocity[1] + a0[1] * half_dt,
        velocity[2] + a0[2] * half_dt,
    ];

    // Exact conformal drift.
    let (r1, v1) = KsPropagator::from_state(position, v_half, gm)?.propagate(dt)?;

    // Half-kick.
    let a1 = accel(r1, v1);
    let v_out = [
        v1[0] + a1[0] * half_dt,
        v1[1] + a1[1] * half_dt,
        v1[2] + a1[2] * half_dt,
    ];
    Ok((r1, v_out))
}