astrodynamics 0.12.1

//! Generic weighted least-squares substrate.
//!
//! Domain-free numerical building blocks for nonlinear least-squares fitting:
//! a forward-difference Jacobian and a trust-region (trf-style) Gauss-Newton
//! solver. Nothing here knows about GNSS, orbits, or any physical units; the
//! caller supplies a residual closure `r: R^n -> R^m` and optional diagonal
//! weights.
//!
//! Two distinct numerical regimes live in this module:
//!
//! - The residual evaluation and the finite-difference Jacobian are pure
//!   `f64` arithmetic plus libm `exp`/etc. inside the caller's closure. Their
//!   operation order is fixed and reproducible, so they reproduce a reference
//!   implementation (e.g. scipy `approx_derivative`) bit-for-bit when the same
//!   recipe and the same libm are used.
//! - The solver's trust-region step solves a linear subproblem via a matrix
//!   factorization. That factorization goes through dense linear algebra whose
//!   last bits depend on the BLAS/LAPACK backend, so the converged solution is
//!   reproducible only to a tight tolerance, not bit-for-bit.
//!
//! Keeping the finite-difference primitive separate from the linear-algebra
//! step lets callers assert the former to the bit while treating the latter as
//! a tolerance-bound agreement.

use nalgebra::{DMatrix, DVector};

/// Relative finite-difference step for a 2-point (forward) scheme: `sqrt(eps)`
/// for `f64`, i.e. `2^-26`. This matches scipy's `_eps_for_method` choice for
/// the `"2-point"` method.
pub const FD_REL_STEP_2POINT: f64 = 1.4901161193847656e-8; // 0x1.0p-26 == sqrt(2^-52)

/// Default first-order optimality tolerance (scipy `least_squares` `gtol`).
const TRF_DEFAULT_GTOL: f64 = 1e-10;
/// Default relative-cost-reduction tolerance (scipy `least_squares` `ftol`).
const TRF_DEFAULT_FTOL: f64 = 1e-8;
/// Default relative-step tolerance (scipy `least_squares` `xtol`).
const TRF_DEFAULT_XTOL: f64 = 1e-8;
/// Default maximum residual evaluations.
const TRF_DEFAULT_MAX_NFEV: usize = 300;
/// Initial Levenberg damping as a fraction of the largest Gauss-Newton normal
/// diagonal: `mu0 = TRF_INITIAL_DAMPING_SCALE * max_i (J^T J)_ii`.
const TRF_INITIAL_DAMPING_SCALE: f64 = 1e-3;

/// Per-parameter step pieces for a single forward-difference column, recorded
/// in evaluation order so they can be inspected or compared against a
/// reference trace.
#[derive(Debug, Clone, PartialEq)]
pub struct FdStep {
    /// Index of the perturbed parameter.
    pub param_index: usize,
    /// `+1.0` if `x0[i] >= 0`, else `-1.0` (the `(x0>=0)*2 - 1` convention;
    /// note `x0[i] == 0` yields `+1.0`).
    pub sign_x0: f64,
    /// Nominal step `rel_step * sign_x0 * max(1, |x0[i]|)`.
    pub h: f64,
    /// Effective step after rounding: `(x0[i] + h) - x0[i]`. This is the
    /// denominator actually used for the column, recomputed rather than reused
    /// from `h`.
    pub dx: f64,
    /// The perturbed parameter vector (only component `i` bumped by `h`).
    pub x_perturbed: DVector<f64>,
}

/// Compute the per-parameter forward-difference step pieces for `x0`.
///
/// `sign_x0[i] = +1 if x0[i] >= 0 else -1`, `h[i] = rel_step * sign_x0[i] *
/// max(1, |x0[i]|)`, and the effective step `dx[i] = (x0[i] + h[i]) - x0[i]`.
/// The post-rounding `dx` is the value used as the column denominator.
pub fn fd_steps(x0: &DVector<f64>, rel_step: f64) -> Vec<FdStep> {
    (0..x0.len())
        .map(|i| {
            let xi = x0[i];
            let sign_x0 = if xi >= 0.0 { 1.0 } else { -1.0 };
            let h = rel_step * sign_x0 * xi.abs().max(1.0);
            let mut x_perturbed = x0.clone();
            x_perturbed[i] = xi + h;
            let dx = x_perturbed[i] - xi;
            FdStep {
                param_index: i,
                sign_x0,
                h,
                dx,
                x_perturbed,
            }
        })
        .collect()
}

/// Forward (2-point) finite-difference Jacobian of `residual` at `x0`, given a
/// precomputed `f0 = residual(x0)`.
///
/// Column `i` is `(residual(x0 + h_i e_i) - f0) / dx_i`, where `h_i` and `dx_i`
/// come from [`fd_steps`]. The arithmetic is plain `f64` (no fused multiply-add):
/// each entry is one subtraction followed by one division, matching scipy's
/// `approx_derivative` operation order.
///
/// `f0` is passed in (rather than recomputed) so the caller controls the base
/// evaluation and so the same `f0` used elsewhere is reused exactly.
pub fn jacobian_2point<F>(residual: F, x0: &DVector<f64>, f0: &DVector<f64>) -> DMatrix<f64>
where
    F: Fn(&DVector<f64>) -> DVector<f64>,
{
    let m = f0.len();
    let n = x0.len();
    let steps = fd_steps(x0, FD_REL_STEP_2POINT);
    let mut jac = DMatrix::zeros(m, n);
    for step in &steps {
        let f1 = residual(&step.x_perturbed);
        let i = step.param_index;
        for row in 0..m {
            jac[(row, i)] = (f1[row] - f0[row]) / step.dx;
        }
    }
    jac
}

/// Termination state of a [`solve_trf`] run, mirroring the scipy
/// `least_squares` status codes for the conditions this solver detects.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Status {
    /// `||J^T r||_inf` fell below `gtol` (first-order optimality).
    GradientTolerance,
    /// The relative cost reduction fell below `ftol`.
    CostTolerance,
    /// The relative step size fell below `xtol`.
    StepTolerance,
    /// The maximum number of residual evaluations was reached.
    MaxEvaluations,
}

/// Stopping tolerances and evaluation budget for [`solve_trf`].
#[derive(Debug, Clone, Copy)]
pub struct SolveOptions {
    /// First-order optimality tolerance on `||J^T r||_inf`.
    pub gtol: f64,
    /// Relative-cost-reduction tolerance.
    pub ftol: f64,
    /// Relative-step tolerance.
    pub xtol: f64,
    /// Maximum number of residual evaluations.
    pub max_nfev: usize,
}

impl Default for SolveOptions {
    fn default() -> Self {
        // scipy's defaults for these tolerances.
        Self {
            gtol: TRF_DEFAULT_GTOL,
            ftol: TRF_DEFAULT_FTOL,
            xtol: TRF_DEFAULT_XTOL,
            max_nfev: TRF_DEFAULT_MAX_NFEV,
        }
    }
}

/// Result of a [`solve_trf`] run.
#[derive(Debug, Clone)]
pub struct LeastSquaresReport {
    /// Converged parameter vector.
    pub x: DVector<f64>,
    /// Residual at `x`.
    pub residual: DVector<f64>,
    /// Cost `0.5 * dot(r, r)` at `x`.
    pub cost: f64,
    /// Finite-difference Jacobian at `x`.
    pub jacobian: DMatrix<f64>,
    /// First-order optimality `||J^T r||_inf` at `x`.
    pub optimality_inf: f64,
    /// Number of accepted iterations.
    pub iterations: usize,
    /// Why the solve stopped.
    pub status: Status,
}

/// Error from [`solve_trf`].
#[derive(Debug, Clone, thiserror::Error)]
pub enum SolveError {
    /// The Jacobian is rank-deficient / the trust-region subproblem has no
    /// usable descent direction (degenerate geometry).
    #[error("singular or rank-deficient Jacobian: no usable descent direction")]
    SingularJacobian,
}

/// Cost `0.5 * dot(r, r)`, a plain fold of `f64` operations.
pub fn cost(residual: &DVector<f64>) -> f64 {
    0.5 * residual.dot(residual)
}

/// A nonlinear least-squares problem: a residual closure, optional diagonal
/// weights, and a starting point. The weighted form scales both the residual
/// and the Jacobian rows by `sqrt(weight)`; with all weights `1` it reduces to
/// the ordinary (unweighted) least-squares problem.
pub struct LeastSquaresProblem<F> {
    residual: F,
    /// `sqrt` of the diagonal weights, or `None` for the identity weighting.
    sqrt_weights: Option<DVector<f64>>,
    x0: DVector<f64>,
}

impl<F> LeastSquaresProblem<F>
where
    F: Fn(&DVector<f64>) -> DVector<f64>,
{
    /// An unweighted problem (identity weighting).
    pub fn new(residual: F, x0: DVector<f64>) -> Self {
        Self {
            residual,
            sqrt_weights: None,
            x0,
        }
    }

    /// A problem with diagonal weights `W`; residual and Jacobian rows are
    /// scaled by `sqrt(W)`.
    pub fn with_weights(residual: F, x0: DVector<f64>, weights: DVector<f64>) -> Self {
        let sqrt_weights = weights.map(f64::sqrt);
        Self {
            residual,
            sqrt_weights: Some(sqrt_weights),
            x0,
        }
    }

    /// Weighted residual at `x`.
    fn weighted_residual(&self, x: &DVector<f64>) -> DVector<f64> {
        let r = (self.residual)(x);
        match &self.sqrt_weights {
            Some(sw) => r.component_mul(sw),
            None => r,
        }
    }
}

/// Trust-region (trf-style) Gauss-Newton solve.
///
/// At each iterate the weighted residual and its forward-difference Jacobian
/// are formed, then a Levenberg-damped Gauss-Newton step `(J^T J + mu I) dx =
/// -J^T r` is taken inside a trust region; the damping `mu` is grown on a
/// rejected step and shrunk on an accepted one. The linear solve uses a dense
/// factorization, so the converged solution is reproducible to a tight
/// tolerance rather than to the bit.
///
/// Returns [`SolveError::SingularJacobian`] if the normal-equation system
/// cannot be solved (degenerate geometry).
pub fn solve_trf<F>(
    problem: &LeastSquaresProblem<F>,
    opts: &SolveOptions,
) -> Result<LeastSquaresReport, SolveError>
where
    F: Fn(&DVector<f64>) -> DVector<f64>,
{
    let n = problem.x0.len();

    let mut x = problem.x0.clone();
    let mut r = problem.weighted_residual(&x);
    let mut f0 = r.clone();
    let mut jac = jacobian_2point(|p| problem.weighted_residual(p), &x, &f0);
    let mut nfev = 1usize; // the f0 above
    let mut cur_cost = cost(&r);

    // Initial Levenberg damping scaled to the Gauss-Newton normal matrix.
    let jtj0 = jac.transpose() * &jac;
    let mut mu = TRF_INITIAL_DAMPING_SCALE
        * (0..n)
            .map(|i| jtj0[(i, i)])
            .fold(0.0_f64, f64::max)
            .max(1.0);

    let mut iterations = 0usize;

    loop {
        let jt = jac.transpose();
        let grad = &jt * &r;
        let optimality_inf = grad.amax();

        if optimality_inf < opts.gtol {
            return Ok(finish(
                x,
                r,
                cur_cost,
                jac,
                iterations,
                Status::GradientTolerance,
            ));
        }
        if nfev >= opts.max_nfev {
            return Ok(finish(
                x,
                r,
                cur_cost,
                jac,
                iterations,
                Status::MaxEvaluations,
            ));
        }

        let jtj = &jt * &jac;

        // Levenberg-damped Gauss-Newton subproblem.
        let mut accepted = false;
        for _ in 0..30 {
            let mut lhs = jtj.clone();
            for i in 0..n {
                lhs[(i, i)] += mu;
            }
            let rhs = -&grad;
            let step = match lhs.clone().lu().solve(&rhs) {
                Some(s) => s,
                None => return Err(SolveError::SingularJacobian),
            };

            let x_trial = &x + &step;
            let r_trial = problem.weighted_residual(&x_trial);
            nfev += 1;
            let cost_trial = cost(&r_trial);

            if cost_trial < cur_cost {
                // Accept; relative-cost and relative-step stopping checks.
                let cost_reduction = (cur_cost - cost_trial) / cur_cost.max(f64::MIN_POSITIVE);
                let step_norm = step.norm();
                let x_norm = x.norm();
                let rel_step = step_norm / x_norm.max(f64::MIN_POSITIVE);

                x = x_trial;
                r = r_trial;
                cur_cost = cost_trial;
                f0 = r.clone();
                jac = jacobian_2point(|p| problem.weighted_residual(p), &x, &f0);
                nfev += n; // FD probes for the new Jacobian
                iterations += 1;
                mu *= 0.5;
                accepted = true;

                if cost_reduction < opts.ftol {
                    return Ok(finish(
                        x,
                        r,
                        cur_cost,
                        jac,
                        iterations,
                        Status::CostTolerance,
                    ));
                }
                if rel_step < opts.xtol {
                    return Ok(finish(
                        x,
                        r,
                        cur_cost,
                        jac,
                        iterations,
                        Status::StepTolerance,
                    ));
                }
                break;
            } else {
                // Reject: grow damping and retry the subproblem.
                mu *= 2.0;
            }
        }

        if !accepted {
            // Could not find an improving step within the damping sweep.
            return Ok(finish(
                x,
                r,
                cur_cost,
                jac,
                iterations,
                Status::StepTolerance,
            ));
        }
    }
}

fn finish(
    x: DVector<f64>,
    residual: DVector<f64>,
    cost_value: f64,
    jacobian: DMatrix<f64>,
    iterations: usize,
    status: Status,
) -> LeastSquaresReport {
    let optimality_inf = (jacobian.transpose() * &residual).amax();
    LeastSquaresReport {
        x,
        residual,
        cost: cost_value,
        jacobian,
        optimality_inf,
        iterations,
        status,
    }
}

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

    #[test]
    fn fd_rel_step_is_sqrt_eps() {
        assert_eq!(FD_REL_STEP_2POINT, (2.0_f64.powi(-52)).sqrt());
        assert_eq!(FD_REL_STEP_2POINT, 2.0_f64.powi(-26));
    }

    #[test]
    fn fd_step_sign_convention() {
        let x0 = DVector::from_vec(vec![5.0, -2.0, 0.0]);
        let steps = fd_steps(&x0, FD_REL_STEP_2POINT);
        assert_eq!(steps[0].sign_x0, 1.0);
        assert_eq!(steps[1].sign_x0, -1.0);
        assert_eq!(steps[2].sign_x0, 1.0); // x == 0 -> +1
    }

    #[test]
    fn exp_fit_converges() {
        // a*exp(b*t) + c with a known minimum near the generated data.
        let t = vec![0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 4.0];
        let y = vec![
            3.0123, 2.2083, 1.6889, 1.3713, 1.0903, 0.9302, 0.8104, 0.6303,
        ];
        let tt = t.clone();
        let yy = y.clone();
        let residual = move |p: &DVector<f64>| {
            let (a, b, c) = (p[0], p[1], p[2]);
            DVector::from_iterator(
                tt.len(),
                tt.iter()
                    .zip(&yy)
                    .map(|(&tk, &yk)| a * (b * tk).exp() + c - yk),
            )
        };
        let problem = LeastSquaresProblem::new(residual, DVector::from_vec(vec![5.0, -2.0, 2.0]));
        let report = solve_trf(&problem, &SolveOptions::default()).unwrap();
        assert!(report.cost < 1.0, "cost did not reduce: {}", report.cost);
    }
}