kshana 0.37.0

// SPDX-License-Identifier: AGPL-3.0-only
//! Gauss–Newton weighted batch least squares — the batch differential corrector.
//!
//! This is the estimation core a batch *orbit determination* (or any parameter-fit)
//! uses: given measurements `z` that depend on an unknown state `x` through a model
//! `h(x)`, it linearises `h` about the current estimate (a finite-difference Jacobian
//! `H = ∂h/∂x`), forms and solves the weighted normal equations
//! `(HᵀWH)·Δx = HᵀW·(z − h(x))`, and iterates to convergence.
//!
//! It is deliberately generic over the model closure, so it is independent of any
//! particular measurement set. The orbit-specific measurement models live in
//! [`crate::orbit_determination`]: its range-only path linearises through
//! [`gauss_newton`] (finite-difference Jacobian, kept as it was), and its
//! range / range-rate / azimuth-elevation tracking path supplies an analytic Jacobian
//! `H = (∂h/∂x(t))·Φ(t, t0)` from the variational state-transition matrix to
//! [`gauss_newton_analytic`], which does no finite differencing. Both are compared with
//! Orekit 12.2's `BatchLSEstimator` (see `docs/VALIDATION.md`).
//!
//! **References.** The weighted-normal-equation batch corrector is the standard
//! batch processor of Tapley, Schutz & Born, *Statistical Orbit Determination*
//! (2004), §4; the underlying Gauss–Newton / weighted-least-squares numerics follow
//! Björck, *Numerical Methods for Least Squares Problems* (1996).

use crate::fusion::ukf::inverse;
use crate::portable_math::PortableFloat;

/// The outcome of a batch least-squares solve.
#[derive(Clone, Debug)]
pub struct LsqResult {
    /// Converged (or last) state estimate.
    pub x: Vec<f64>,
    /// Iterations actually run.
    pub iterations: usize,
    /// RMS of the post-fit measurement residual `z − h(x)`.
    pub rms_residual: f64,
    /// `true` when the step norm fell below `tol` before `max_iter`.
    pub converged: bool,
}

/// Central finite-difference Jacobian `H` (`m × n`) of `h` at `x`, the one [`gauss_newton`]
/// linearises with (step `1e-6 · max(|x_p|, 1)` per parameter).
pub fn fd_jacobian<H>(h: &H, x: &[f64], m: usize) -> Vec<Vec<f64>>
where
    H: Fn(&[f64]) -> Vec<f64>,
{
    let n = x.len();
    let mut jac = vec![vec![0.0; n]; m];
    for (p, &xp_val) in x.iter().enumerate() {
        let step = 1e-6 * xp_val.abs().max(1.0);
        let mut xp = x.to_vec();
        let mut xm = x.to_vec();
        xp[p] += step;
        xm[p] -= step;
        let hp = h(&xp);
        let hm = h(&xm);
        for i in 0..m {
            jac[i][p] = (hp[i] - hm[i]) / (2.0 * step);
        }
    }
    jac
}

/// Solve `z ≈ h(x)` by Gauss–Newton iteration with per-measurement weights
/// `weights` (e.g. `1/σ²`), starting from `x0`. Returns `None` on a dimension
/// mismatch or a singular normal matrix (rank-deficient geometry).
///
/// For a *linear* model `h(x) = Ax` this reaches the exact weighted-least-squares
/// solution `(AᵀWA)⁻¹AᵀWz` in a single step (the second iteration then confirms a
/// zero update); for a nonlinear model it converges quadratically near the solution.
pub fn gauss_newton<H>(
    h: H,
    z: &[f64],
    weights: &[f64],
    x0: &[f64],
    max_iter: usize,
    tol: f64,
) -> Option<LsqResult>
where
    H: Fn(&[f64]) -> Vec<f64>,
{
    let n = x0.len();
    let m = z.len();
    if weights.len() != m || n == 0 || m < n {
        return None;
    }
    let mut x = x0.to_vec();
    let mut iterations = 0;
    let mut converged = false;
    for it in 0..max_iter {
        iterations = it + 1;
        let hx = h(&x);
        if hx.len() != m {
            return None;
        }
        let r: Vec<f64> = (0..m).map(|i| z[i] - hx[i]).collect();
        let jac = fd_jacobian(&h, &x, m);
        // Weighted normal equations: a = HᵀWH (n×n), b = HᵀW r (n).
        let mut a = vec![vec![0.0; n]; n];
        let mut b = vec![0.0; n];
        for i in 0..m {
            let w = weights[i];
            for p in 0..n {
                b[p] += jac[i][p] * w * r[i];
                for q in 0..n {
                    a[p][q] += jac[i][p] * w * jac[i][q];
                }
            }
        }
        let a_inv = inverse(&a)?;
        let dx: Vec<f64> = (0..n)
            .map(|p| (0..n).map(|q| a_inv[p][q] * b[q]).sum())
            .collect();
        for (xp, &d) in x.iter_mut().zip(&dx) {
            *xp += d;
        }
        let dx_norm = dx.iter().map(|v| v * v).sum::<f64>().sqrt();
        if dx_norm < tol {
            converged = true;
            break;
        }
    }
    let hx = h(&x);
    let rms = (z
        .iter()
        .zip(&hx)
        .map(|(&zi, &hi)| (zi - hi).ppowi(2))
        .sum::<f64>()
        / m as f64)
        .sqrt();
    Some(LsqResult {
        x,
        iterations,
        rms_residual: rms,
        converged,
    })
}

/// [`gauss_newton`] with each step solved by Householder QR of the whitened Jacobian
/// (`min ‖√W (H dx − r)‖`, [`crate::leo_navmsg::fit::lstsq`]) instead of inverting the normal
/// matrix `HᵀWH`. The normal equations square the condition number; QR keeps the step's
/// forward error near `κ ε` rather than `κ² ε`, which matters for ill-conditioned networks.
/// Same arguments, convergence rule and result as [`gauss_newton`].
pub fn gauss_newton_qr<H>(
    h: H,
    z: &[f64],
    weights: &[f64],
    x0: &[f64],
    max_iter: usize,
    tol: f64,
) -> Option<LsqResult>
where
    H: Fn(&[f64]) -> Vec<f64>,
{
    let n = x0.len();
    let m = z.len();
    if weights.len() != m || n == 0 || m < n || weights.iter().any(|w| w.is_nan() || *w < 0.0) {
        return None;
    }
    let sw: Vec<f64> = weights.iter().map(|w| w.sqrt()).collect();
    let mut x = x0.to_vec();
    let mut iterations = 0;
    let mut converged = false;
    for it in 0..max_iter {
        iterations = it + 1;
        let hx = h(&x);
        if hx.len() != m {
            return None;
        }
        let jac = fd_jacobian(&h, &x, m);
        let a: Vec<Vec<f64>> = jac
            .iter()
            .zip(&sw)
            .map(|(row, s)| row.iter().map(|v| v * s).collect())
            .collect();
        let b: Vec<f64> = (0..m).map(|i| sw[i] * (z[i] - hx[i])).collect();
        let dx = crate::leo_navmsg::fit::lstsq(&a, &b)?;
        for (xp, &d) in x.iter_mut().zip(&dx) {
            *xp += d;
        }
        let dx_norm = dx.iter().map(|v| v * v).sum::<f64>().sqrt();
        if dx_norm < tol {
            converged = true;
            break;
        }
    }
    let hx = h(&x);
    let rms = (z
        .iter()
        .zip(&hx)
        .map(|(&zi, &hi)| (zi - hi).ppowi(2))
        .sum::<f64>()
        / m as f64)
        .sqrt();
    Some(LsqResult {
        x,
        iterations,
        rms_residual: rms,
        converged,
    })
}

/// [`gauss_newton`] with each step solved by the opt-in square-root route of
/// [`crate::linalg_sr::weighted_lstsq`] (Householder QR of `[√W H | √W r]`, never forming `HᵀWH`),
/// returning with the result the formal covariance `R⁻¹R⁻ᵀ` at the converged state, read off the
/// same triangular factor rather than an inverted normal matrix. Same arguments and convergence
/// rule as [`gauss_newton`]; `None` on a dimension mismatch, a negative weight or a rank-deficient
/// Jacobian. The default solvers are unchanged.
pub fn gauss_newton_srif<H>(
    h: H,
    z: &[f64],
    weights: &[f64],
    x0: &[f64],
    max_iter: usize,
    tol: f64,
) -> Option<(LsqResult, Vec<Vec<f64>>)>
where
    H: Fn(&[f64]) -> Vec<f64>,
{
    let n = x0.len();
    let m = z.len();
    if weights.len() != m || n == 0 || m < n {
        return None;
    }
    let mut x = x0.to_vec();
    let mut iterations = 0;
    let mut converged = false;
    for it in 0..max_iter {
        iterations = it + 1;
        let hx = h(&x);
        if hx.len() != m {
            return None;
        }
        let r: Vec<f64> = (0..m).map(|i| z[i] - hx[i]).collect();
        let jac = fd_jacobian(&h, &x, m);
        let dx = crate::linalg_sr::weighted_lstsq(&jac, weights, &r)?.x;
        for (xp, &d) in x.iter_mut().zip(&dx) {
            *xp += d;
        }
        let dx_norm = dx.iter().map(|v| v * v).sum::<f64>().sqrt();
        if dx_norm < tol {
            converged = true;
            break;
        }
    }
    let hx = h(&x);
    let jac = fd_jacobian(&h, &x, m);
    let resid: Vec<f64> = (0..m).map(|i| z[i] - hx[i]).collect();
    let sol = crate::linalg_sr::weighted_lstsq(&jac, weights, &resid)?;
    let cov = crate::linalg_sr::covariance_from_sqrt_information(&sol.r)?;
    let rms = (resid.iter().map(|v| v * v).sum::<f64>() / m as f64).sqrt();
    Some((
        LsqResult {
            x,
            iterations,
            rms_residual: rms,
            converged,
        },
        cov,
    ))
}

/// Gauss–Newton with a caller-supplied **analytic Jacobian**: `model(x)` returns the predicted
/// measurements `h(x)` and their Jacobian `H = ∂h/∂x` (`m × n`) together, so no finite
/// differencing is done. Each step is solved by the square-root route of
/// [`crate::linalg_sr::weighted_lstsq`] (Householder QR of the whitened Jacobian), and the formal
/// covariance `(HᵀWH)⁻¹` at the final state is returned from the same triangular factor. Same
/// arguments and step-norm convergence rule as [`gauss_newton`]; `None` on a dimension mismatch,
/// a model failure (`model` returning `None`), a negative weight or a rank-deficient Jacobian.
///
/// This is the corrector the state-transition-matrix orbit-determination path
/// ([`crate::orbit_determination::determine_orbit_batch_tracking`]) uses, with
/// `H = (∂h/∂x(t))·Φ(t, t0)`.
pub fn gauss_newton_analytic<F>(
    model: F,
    z: &[f64],
    weights: &[f64],
    x0: &[f64],
    max_iter: usize,
    tol: f64,
) -> Option<(LsqResult, Vec<Vec<f64>>)>
where
    F: Fn(&[f64]) -> Option<(Vec<f64>, Vec<Vec<f64>>)>,
{
    let n = x0.len();
    let m = z.len();
    if weights.len() != m || n == 0 || m < n {
        return None;
    }
    let mut x = x0.to_vec();
    let mut iterations = 0;
    let mut converged = false;
    for it in 0..max_iter {
        iterations = it + 1;
        let (hx, jac) = model(&x)?;
        if hx.len() != m || jac.len() != m {
            return None;
        }
        let r: Vec<f64> = (0..m).map(|i| z[i] - hx[i]).collect();
        let dx = crate::linalg_sr::weighted_lstsq(&jac, weights, &r)?.x;
        for (xp, &d) in x.iter_mut().zip(&dx) {
            *xp += d;
        }
        let dx_norm = dx.iter().map(|v| v * v).sum::<f64>().sqrt();
        if dx_norm < tol {
            converged = true;
            break;
        }
    }
    let (hx, jac) = model(&x)?;
    if hx.len() != m || jac.len() != m {
        return None;
    }
    let resid: Vec<f64> = (0..m).map(|i| z[i] - hx[i]).collect();
    let sol = crate::linalg_sr::weighted_lstsq(&jac, weights, &resid)?;
    let cov = crate::linalg_sr::covariance_from_sqrt_information(&sol.r)?;
    let rms = (resid.iter().map(|v| v * v).sum::<f64>() / m as f64).sqrt();
    Some((
        LsqResult {
            x,
            iterations,
            rms_residual: rms,
            converged,
        },
        cov,
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_analytic_jacobian_corrector_matches_the_finite_difference_corrector() {
        let ts = [0.0f64, 0.5, 1.0, 1.5, 2.0, 2.5];
        let z: Vec<f64> = ts.iter().map(|&t| 1.5 * (0.5 * t).pexp()).collect();
        let w = vec![1.0, 2.0, 1.0, 0.5, 1.0, 3.0];
        let h = move |x: &[f64]| {
            ts.iter()
                .map(|&t| x[0] * (x[1] * t).pexp())
                .collect::<Vec<_>>()
        };
        let model = move |x: &[f64]| {
            let hx = h(x);
            let jac = ts
                .iter()
                .map(|&t| vec![(x[1] * t).pexp(), x[0] * t * (x[1] * t).pexp()])
                .collect();
            Some((hx, jac))
        };
        let (a, cov) = gauss_newton_analytic(model, &z, &w, &[1.0, 1.0], 50, 1e-12).expect("ok");
        let (b, cov_fd) = gauss_newton_srif(h, &z, &w, &[1.0, 1.0], 50, 1e-12).expect("ok");
        assert!(a.converged);
        for (p, q) in a.x.iter().zip(&b.x) {
            assert!((p - q).abs() < 1e-9, "{p} vs {q}");
        }
        for i in 0..2 {
            for j in 0..2 {
                assert!((cov[i][j] - cov_fd[i][j]).abs() < 1e-6 * cov[i][i].abs().max(1e-12));
            }
        }
        // A model failure or a dimension mismatch is reported, not hidden.
        assert!(gauss_newton_analytic(|_: &[f64]| None, &z, &w, &[1.0, 1.0], 5, 1e-9).is_none());
        assert!(gauss_newton_analytic(model, &z, &w[..3], &[1.0, 1.0], 5, 1e-9).is_none());
    }

    #[test]
    fn linear_line_fit_recovers_exact_solution() {
        // Fit y = a + b·t to points exactly on the line a=2, b=3. A linear model must
        // reach the exact weighted-least-squares answer (residual 0).
        let ts = [0.0, 1.0, 2.0, 3.0, 4.0];
        let z: Vec<f64> = ts.iter().map(|&t| 2.0 + 3.0 * t).collect();
        let w = vec![1.0; ts.len()];
        let model = move |x: &[f64]| ts.iter().map(|&t| x[0] + x[1] * t).collect::<Vec<_>>();
        let r = gauss_newton(model, &z, &w, &[0.0, 0.0], 10, 1e-12).expect("solves");
        assert!((r.x[0] - 2.0).abs() < 1e-9, "a = {}", r.x[0]);
        assert!((r.x[1] - 3.0).abs() < 1e-9, "b = {}", r.x[1]);
        assert!(r.rms_residual < 1e-9 && r.converged);
    }

    #[test]
    fn nonlinear_exponential_fit_converges() {
        // Fit y = a·exp(b·t); recover a=1.5, b=0.5 from noise-free samples.
        let ts = [0.0f64, 0.5, 1.0, 1.5, 2.0, 2.5];
        let z: Vec<f64> = ts.iter().map(|&t| 1.5 * (0.5 * t).pexp()).collect();
        let w = vec![1.0; ts.len()];
        let model = move |x: &[f64]| {
            ts.iter()
                .map(|&t| x[0] * (x[1] * t).pexp())
                .collect::<Vec<_>>()
        };
        let r = gauss_newton(model, &z, &w, &[1.0, 1.0], 50, 1e-12).expect("solves");
        assert!((r.x[0] - 1.5).abs() < 1e-6, "a = {}", r.x[0]);
        assert!((r.x[1] - 0.5).abs() < 1e-6, "b = {}", r.x[1]);
        assert!(r.rms_residual < 1e-6, "rms = {}", r.rms_residual);
    }

    #[test]
    fn range_multilateration_recovers_a_known_position() {
        // Orbit-determination flavour: recover a 3-D position from ranges to four known
        // stations. Noise-free measurements ⇒ exact recovery (residual ≈ 0).
        let stations = [
            [0.0, 0.0, 0.0],
            [1000.0, 0.0, 0.0],
            [0.0, 1000.0, 0.0],
            [0.0, 0.0, 1000.0],
            [1000.0, 1000.0, 1000.0],
        ];
        let truth = [350.0, -120.0, 640.0];
        let range = |s: &[f64; 3], p: &[f64]| {
            ((p[0] - s[0]).ppowi(2) + (p[1] - s[1]).ppowi(2) + (p[2] - s[2]).ppowi(2)).sqrt()
        };
        let z: Vec<f64> = stations.iter().map(|s| range(s, &truth)).collect();
        let w = vec![1.0; stations.len()];
        let model = move |x: &[f64]| stations.iter().map(|s| range(s, x)).collect::<Vec<_>>();
        let r = gauss_newton(model, &z, &w, &[0.0, 0.0, 0.0], 50, 1e-10).expect("solves");
        for (k, (&got, &want)) in r.x.iter().zip(&truth).enumerate() {
            assert!((got - want).abs() < 1e-4, "x[{k}] = {got} vs {want}");
        }
        assert!(r.rms_residual < 1e-4, "rms = {}", r.rms_residual);
    }

    #[test]
    #[allow(clippy::needless_range_loop)] // the 2x2 product mirrors its formula
    fn the_srif_corrector_matches_the_normal_equation_corrector_and_returns_its_covariance() {
        let ts = [0.0f64, 0.5, 1.0, 1.5, 2.0, 2.5];
        let z: Vec<f64> = ts.iter().map(|&t| 1.5 * (0.5 * t).pexp()).collect();
        let w = vec![1.0, 2.0, 1.0, 0.5, 1.0, 3.0];
        let model = move |x: &[f64]| {
            ts.iter()
                .map(|&t| x[0] * (x[1] * t).pexp())
                .collect::<Vec<_>>()
        };
        let a = gauss_newton(model, &z, &w, &[1.0, 1.0], 50, 1e-12).expect("solves");
        let (b, cov) = gauss_newton_srif(model, &z, &w, &[1.0, 1.0], 50, 1e-12).expect("solves");
        assert!(b.converged);
        for (p, q) in a.x.iter().zip(&b.x) {
            assert!((p - q).abs() < 1e-9, "{p} vs {q}");
        }
        // The covariance inverts the normal matrix at the solution.
        let jac = fd_jacobian(&model, &b.x, ts.len());
        let mut info = [[0.0; 2]; 2];
        for (row, wi) in jac.iter().zip(&w) {
            for p in 0..2 {
                for q in 0..2 {
                    info[p][q] += row[p] * wi * row[q];
                }
            }
        }
        for i in 0..2 {
            for j in 0..2 {
                let e: f64 = (0..2).map(|k| info[i][k] * cov[k][j]).sum();
                assert!((e - if i == j { 1.0 } else { 0.0 }).abs() < 1e-9);
            }
        }
    }

    #[test]
    fn rejects_underdetermined_or_mismatched_inputs() {
        let model = |x: &[f64]| vec![x[0], x[0] + x[1]];
        // Fewer measurements than unknowns (m < n).
        assert!(gauss_newton(model, &[1.0], &[1.0], &[0.0, 0.0], 5, 1e-9).is_none());
        // Weight/measurement length mismatch.
        assert!(gauss_newton(model, &[1.0, 2.0], &[1.0], &[0.0, 0.0], 5, 1e-9).is_none());
    }
}