oxiproj-core 0.1.2

//! Forward-mode dual numbers for automatic differentiation.
//!
//! `Dual1<N>` carries a primal value and an N-dimensional gradient vector,
//! enabling forward-mode AD through any computation expressed via the
//! standard arithmetic operators and transcendental functions defined here.
//!
//! # Typical use
//!
//! For 2-D projection Jacobian (partial/partial lambda, partial/partial phi),
//! set `N = 2` and seed:
//! - longitude variable: `Dual1::<2>::variable(lam, 0)`
//! - latitude variable:  `Dual1::<2>::variable(phi, 1)`
//!
//! After passing both through a generic projection function the output
//! dual numbers carry the full 2x2 Jacobian in their `.d` arrays.

use core::fmt;
use core::ops::{Add, Div, Mul, Neg, Sub};

/// Forward-mode dual number: primal value + N-dimensional gradient.
///
/// `N = 2` covers the 2-D projection Jacobian (partial/partial lambda,
/// partial/partial phi). Larger `N` is supported for other uses.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Dual1<const N: usize> {
    /// Primal value.
    pub v: f64,
    /// Partial derivatives: `d[i] = partial f / partial xi`.
    pub d: [f64; N],
}

/// Apply a unary chain-rule elementwise: `d_out[i] = scale * d_in[i]`.
#[inline]
fn map_scale<const N: usize>(d: &[f64; N], scale: f64) -> [f64; N] {
    let mut out = [0.0; N];
    for (o, &di) in out.iter_mut().zip(d.iter()) {
        *o = scale * di;
    }
    out
}

/// Apply a binary elementwise operation: `d_out[i] = f(a[i], b[i])`.
#[inline]
fn map_binary<const N: usize, F: Fn(f64, f64) -> f64>(
    a: &[f64; N],
    b: &[f64; N],
    f: F,
) -> [f64; N] {
    let mut out = [0.0; N];
    for (o, (&ai, &bi)) in out.iter_mut().zip(a.iter().zip(b.iter())) {
        *o = f(ai, bi);
    }
    out
}

impl<const N: usize> Dual1<N> {
    /// Construct a constant (zero gradient).
    #[inline]
    pub fn constant(v: f64) -> Self {
        Self { v, d: [0.0; N] }
    }

    /// Construct an independent variable (1.0 in the `idx`-th slot, 0.0 elsewhere).
    ///
    /// # Panics
    ///
    /// Panics if `idx >= N`.
    #[inline]
    pub fn variable(v: f64, idx: usize) -> Self {
        assert!(
            idx < N,
            "Dual1::variable: idx {idx} is out of bounds for N={N}"
        );
        let mut d = [0.0; N];
        d[idx] = 1.0;
        Self { v, d }
    }

    // ---- Transcendental methods (all chain-rule, no libm) ----

    /// sin(self)
    #[inline]
    pub fn sin(self) -> Self {
        let cos_v = self.v.cos();
        Self {
            v: self.v.sin(),
            d: map_scale(&self.d, cos_v),
        }
    }

    /// cos(self)
    #[inline]
    pub fn cos(self) -> Self {
        let neg_sin_v = -self.v.sin();
        Self {
            v: self.v.cos(),
            d: map_scale(&self.d, neg_sin_v),
        }
    }

    /// tan(self)
    #[inline]
    pub fn tan(self) -> Self {
        let cos_v = self.v.cos();
        let sec2 = 1.0 / (cos_v * cos_v);
        Self {
            v: self.v.tan(),
            d: map_scale(&self.d, sec2),
        }
    }

    /// asin(self) — chain rule: d/dx asin(f) = f' / sqrt(1 - f^2)
    #[inline]
    pub fn asin(self) -> Self {
        let scale = 1.0 / (1.0 - self.v * self.v).sqrt();
        Self {
            v: self.v.asin(),
            d: map_scale(&self.d, scale),
        }
    }

    /// acos(self) — chain rule: d/dx acos(f) = -f' / sqrt(1 - f^2)
    #[inline]
    pub fn acos(self) -> Self {
        let scale = -1.0 / (1.0 - self.v * self.v).sqrt();
        Self {
            v: self.v.acos(),
            d: map_scale(&self.d, scale),
        }
    }

    /// atan(self) — chain rule: d/dx atan(f) = f' / (1 + f^2)
    #[inline]
    pub fn atan(self) -> Self {
        let scale = 1.0 / (1.0 + self.v * self.v);
        Self {
            v: self.v.atan(),
            d: map_scale(&self.d, scale),
        }
    }

    /// atan2(y=self, x=rhs) — chain rule:
    /// d/dxi atan2(y,x) = (x·y.d\[i\] - y·x.d\[i\]) / (x^2+y^2)
    #[inline]
    pub fn atan2(self, rhs: Dual1<N>) -> Self {
        let denom = rhs.v * rhs.v + self.v * self.v;
        let y = self.v;
        let x = rhs.v;
        let d = map_binary(&self.d, &rhs.d, |yd, xd| (x * yd - y * xd) / denom);
        Self {
            v: self.v.atan2(rhs.v),
            d,
        }
    }

    /// sqrt(self) — chain rule: d/dx sqrt(f) = f' / (2*sqrt(f))
    #[inline]
    pub fn sqrt(self) -> Self {
        let sqrt_v = self.v.sqrt();
        let scale = 1.0 / (2.0 * sqrt_v);
        Self {
            v: sqrt_v,
            d: map_scale(&self.d, scale),
        }
    }

    /// cbrt(self) — chain rule: d/dx cbrt(f) = f' / (3*cbrt(f)^2)
    #[inline]
    pub fn cbrt(self) -> Self {
        let cbrt_v = self.v.cbrt();
        let scale = 1.0 / (3.0 * cbrt_v * cbrt_v);
        Self {
            v: cbrt_v,
            d: map_scale(&self.d, scale),
        }
    }

    /// exp(self)
    #[inline]
    pub fn exp(self) -> Self {
        let exp_v = self.v.exp();
        Self {
            v: exp_v,
            d: map_scale(&self.d, exp_v),
        }
    }

    /// ln(self)
    #[inline]
    pub fn ln(self) -> Self {
        let scale = 1.0 / self.v;
        Self {
            v: self.v.ln(),
            d: map_scale(&self.d, scale),
        }
    }

    /// abs(self) — chain rule: d/dx |f| = f' * signum(f)
    #[inline]
    pub fn abs(self) -> Self {
        let sign = self.v.signum();
        Self {
            v: self.v.abs(),
            d: map_scale(&self.d, sign),
        }
    }

    /// powi(self, n) — chain rule: d/dx f^n = n*f^(n-1)*f'
    #[inline]
    pub fn powi(self, n: i32) -> Self {
        let scale = n as f64 * self.v.powi(n - 1);
        Self {
            v: self.v.powi(n),
            d: map_scale(&self.d, scale),
        }
    }

    /// powf(self, e) — chain rule: d/dx f^e = e*f^(e-1)*f'
    #[inline]
    pub fn powf(self, e: f64) -> Self {
        let scale = e * self.v.powf(e - 1.0);
        Self {
            v: self.v.powf(e),
            d: map_scale(&self.d, scale),
        }
    }

    /// sinh(self)
    #[inline]
    pub fn sinh(self) -> Self {
        let cosh_v = self.v.cosh();
        Self {
            v: self.v.sinh(),
            d: map_scale(&self.d, cosh_v),
        }
    }

    /// cosh(self)
    #[inline]
    pub fn cosh(self) -> Self {
        let sinh_v = self.v.sinh();
        Self {
            v: self.v.cosh(),
            d: map_scale(&self.d, sinh_v),
        }
    }

    /// tanh(self) — chain rule: d/dx tanh(f) = f' / cosh(f)^2
    #[inline]
    pub fn tanh(self) -> Self {
        let cosh_v = self.v.cosh();
        let sech2 = 1.0 / (cosh_v * cosh_v);
        Self {
            v: self.v.tanh(),
            d: map_scale(&self.d, sech2),
        }
    }

    /// hypot(a=self, b=rhs) — chain rule:
    /// d/dxi hypot(a,b) = (a·a.d\[i\] + b·b.d\[i\]) / hypot(a,b)
    #[inline]
    pub fn hypot(self, rhs: Dual1<N>) -> Self {
        let h = self.v.hypot(rhs.v);
        let a = self.v;
        let b = rhs.v;
        let d = map_binary(&self.d, &rhs.d, |ad, bd| (a * ad + b * bd) / h);
        Self { v: h, d }
    }
}

// ---- Arithmetic operators ----

impl<const N: usize> Add for Dual1<N> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: Self) -> Self {
        let d = map_binary(&self.d, &rhs.d, |a, b| a + b);
        Self {
            v: self.v + rhs.v,
            d,
        }
    }
}

impl<const N: usize> Add for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn add(self, rhs: Self) -> Dual1<N> {
        *self + *rhs
    }
}

impl<const N: usize> Add<Dual1<N>> for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn add(self, rhs: Dual1<N>) -> Dual1<N> {
        *self + rhs
    }
}

impl<const N: usize> Add<&Dual1<N>> for Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn add(self, rhs: &Dual1<N>) -> Dual1<N> {
        self + *rhs
    }
}

impl<const N: usize> Sub for Dual1<N> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: Self) -> Self {
        let d = map_binary(&self.d, &rhs.d, |a, b| a - b);
        Self {
            v: self.v - rhs.v,
            d,
        }
    }
}

impl<const N: usize> Sub for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn sub(self, rhs: Self) -> Dual1<N> {
        *self - *rhs
    }
}

impl<const N: usize> Sub<Dual1<N>> for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn sub(self, rhs: Dual1<N>) -> Dual1<N> {
        *self - rhs
    }
}

impl<const N: usize> Sub<&Dual1<N>> for Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn sub(self, rhs: &Dual1<N>) -> Dual1<N> {
        self - *rhs
    }
}

impl<const N: usize> Mul for Dual1<N> {
    type Output = Self;
    /// Product rule: (a*b)' = a'*b + a*b'
    #[inline]
    fn mul(self, rhs: Self) -> Self {
        let av = self.v;
        let bv = rhs.v;
        // Inline to avoid clippy::suspicious_arithmetic_impl on `+` in closure.
        let mut d = [0.0; N];
        for (o, (&ad, &bd)) in d.iter_mut().zip(self.d.iter().zip(rhs.d.iter())) {
            *o = ad * bv + av * bd;
        }
        Self { v: av * bv, d }
    }
}

impl<const N: usize> Mul for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn mul(self, rhs: Self) -> Dual1<N> {
        *self * *rhs
    }
}

impl<const N: usize> Mul<Dual1<N>> for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn mul(self, rhs: Dual1<N>) -> Dual1<N> {
        *self * rhs
    }
}

impl<const N: usize> Mul<&Dual1<N>> for Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn mul(self, rhs: &Dual1<N>) -> Dual1<N> {
        self * *rhs
    }
}

impl<const N: usize> Div for Dual1<N> {
    type Output = Self;
    /// Quotient rule: (a/b)' = (a'*b - a*b') / b^2
    #[inline]
    fn div(self, rhs: Self) -> Self {
        let av = self.v;
        let bv = rhs.v;
        let b2 = bv * bv;
        let d = map_binary(&self.d, &rhs.d, |ad, bd| (ad * bv - av * bd) / b2);
        Self { v: av / bv, d }
    }
}

impl<const N: usize> Div for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn div(self, rhs: Self) -> Dual1<N> {
        *self / *rhs
    }
}

impl<const N: usize> Div<Dual1<N>> for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn div(self, rhs: Dual1<N>) -> Dual1<N> {
        *self / rhs
    }
}

impl<const N: usize> Div<&Dual1<N>> for Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn div(self, rhs: &Dual1<N>) -> Dual1<N> {
        self / *rhs
    }
}

impl<const N: usize> Neg for Dual1<N> {
    type Output = Self;
    #[inline]
    fn neg(self) -> Self {
        Self {
            v: -self.v,
            d: map_scale(&self.d, -1.0),
        }
    }
}

impl<const N: usize> Neg for &Dual1<N> {
    type Output = Dual1<N>;
    #[inline]
    fn neg(self) -> Dual1<N> {
        -*self
    }
}

impl<const N: usize> PartialOrd for Dual1<N> {
    /// Compares on primal value only.
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        self.v.partial_cmp(&other.v)
    }
}

impl<const N: usize> fmt::Display for Dual1<N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.v)
    }
}

// ---- Tests ----

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(feature = "no_std")]
    use alloc::format;
    use core::f64::consts::{FRAC_PI_2, FRAC_PI_4, PI};

    #[test]
    fn dual1_add_mul() {
        // (3 + e0) * (2 + e1) = 6 + 2*e0 + 3*e1
        let a = Dual1::<2>::variable(3.0, 0);
        let b = Dual1::<2>::variable(2.0, 1);
        let c = a * b;
        assert!((c.v - 6.0).abs() < 1e-14);
        assert!((c.d[0] - 2.0).abs() < 1e-14);
        assert!((c.d[1] - 3.0).abs() < 1e-14);
    }

    #[test]
    fn dual1_sin_chain_rule() {
        // d/dx sin(x^2) at x=pi/4 = 2x*cos(x^2) = 2*(pi/4)*cos(pi^2/16)
        let x = Dual1::<1>::variable(FRAC_PI_4, 0);
        let x2 = x * x;
        let s = x2.sin();
        let expected_d = 2.0 * FRAC_PI_4 * (PI * PI / 16.0).cos();
        assert!((s.v - (PI * PI / 16.0).sin()).abs() < 1e-12);
        assert!((s.d[0] - expected_d).abs() < 1e-12);
    }

    #[test]
    fn dual1_mercator_jacobian() {
        // Mercator: x=lam, y=ln(tan(pi/4 + phi/2))
        // partial x / partial lam = 1, partial x / partial phi = 0
        // partial y / partial lam = 0, partial y / partial phi = sec(phi)
        let phi0 = 0.5_f64; // ~28.6 deg
        let lam = Dual1::<2>::variable(1.0, 0);
        let phi = Dual1::<2>::variable(phi0, 1);
        let x = lam;
        let y = (Dual1::constant(FRAC_PI_4) + phi * Dual1::constant(0.5))
            .tan()
            .ln();
        assert!((x.d[0] - 1.0).abs() < 1e-12, "partial x/partial lam = 1");
        assert!(x.d[1].abs() < 1e-12, "partial x/partial phi = 0");
        assert!(y.d[0].abs() < 1e-12, "partial y/partial lam = 0");
        let sec_phi = 1.0 / phi0.cos();
        assert!(
            (y.d[1] - sec_phi).abs() < 1e-9,
            "partial y/partial phi = sec(phi)"
        );
    }

    #[test]
    #[cfg(not(feature = "no_std"))]
    fn dual1_variable_panics_out_of_bounds() {
        // This should panic since N=2 and idx=2 is out of bounds
        let result = std::panic::catch_unwind(|| {
            Dual1::<2>::variable(1.0, 2);
        });
        assert!(result.is_err());
    }

    #[test]
    fn dual1_chain_rules_sqrt() {
        // d/dx sqrt(x^2) = x/sqrt(x^2) = sign(x) for x > 0
        let x = Dual1::<1>::variable(4.0, 0);
        let result = (x * x).sqrt();
        assert!((result.v - 4.0).abs() < 1e-12);
        assert!((result.d[0] - 1.0).abs() < 1e-12); // d/dx sqrt(x^2) = 2x/(2*sqrt(x^2)) = 1 at x=4
    }

    #[test]
    fn dual1_atan2_gradient() {
        // atan2(y, x) at (1, 1): gradient w.r.t. y = x/(x^2+y^2) = 0.5, w.r.t. x = -y/(x^2+y^2) = -0.5
        let y = Dual1::<2>::variable(1.0, 0);
        let x = Dual1::<2>::variable(1.0, 1);
        let result = y.atan2(x);
        assert!((result.v - FRAC_PI_4).abs() < 1e-12);
        assert!((result.d[0] - 0.5).abs() < 1e-12); // partial/partial y atan2(y, x) = x/(x^2+y^2)
        assert!((result.d[1] - (-0.5)).abs() < 1e-12); // partial/partial x atan2(y, x) = -y/(x^2+y^2)
    }

    #[test]
    fn dual1_constant_zero_gradient() {
        let c = Dual1::<3>::constant(42.0);
        assert!((c.v - 42.0).abs() < 1e-14);
        assert_eq!(c.d, [0.0; 3]);
    }

    #[test]
    fn dual1_sub_neg() {
        let a = Dual1::<2>::variable(5.0, 0);
        let b = Dual1::<2>::variable(3.0, 1);
        let c = a - b;
        assert!((c.v - 2.0).abs() < 1e-14);
        assert!((c.d[0] - 1.0).abs() < 1e-14);
        assert!((c.d[1] - (-1.0)).abs() < 1e-14);
        let neg_a = -a;
        assert!((neg_a.v - (-5.0)).abs() < 1e-14);
        assert!((neg_a.d[0] - (-1.0)).abs() < 1e-14);
    }

    #[test]
    fn dual1_div() {
        // d/dx (x / c) = 1/c
        let x = Dual1::<1>::variable(6.0, 0);
        let c = Dual1::<1>::constant(3.0);
        let r = x / c;
        assert!((r.v - 2.0).abs() < 1e-14);
        assert!((r.d[0] - (1.0 / 3.0)).abs() < 1e-14);
    }

    #[test]
    fn dual1_display() {
        let x = Dual1::<2>::variable(PI, 0);
        let s = format!("{x}");
        // Display shows the primal value
        assert!(!s.is_empty());
    }

    #[test]
    fn dual1_partial_ord() {
        let a = Dual1::<1>::constant(1.0);
        let b = Dual1::<1>::constant(2.0);
        assert!(a < b);
        assert!(b > a);
    }

    #[test]
    fn dual1_exp_chain_rule() {
        // d/dx exp(2x) at x=1 = 2*exp(2)
        let x = Dual1::<1>::variable(1.0, 0);
        let two_x = x * Dual1::constant(2.0);
        let r = two_x.exp();
        assert!((r.v - 2.0_f64.exp()).abs() < 1e-12);
        assert!((r.d[0] - 2.0 * 2.0_f64.exp()).abs() < 1e-12);
    }

    #[test]
    fn dual1_ln_chain_rule() {
        // d/dx ln(x^2) = 2/x  at x=3 => 2/3
        let x = Dual1::<1>::variable(3.0, 0);
        let r = (x * x).ln();
        assert!((r.v - (9.0_f64).ln()).abs() < 1e-12);
        assert!((r.d[0] - (2.0 / 3.0)).abs() < 1e-12);
    }

    #[test]
    fn dual1_powi_powf() {
        // powi: d/dx x^3 at x=2 = 3*4 = 12
        let x = Dual1::<1>::variable(2.0, 0);
        let r = x.powi(3);
        assert!((r.v - 8.0).abs() < 1e-12);
        assert!((r.d[0] - 12.0).abs() < 1e-12);
        // powf: same
        let r2 = x.powf(3.0);
        assert!((r2.v - 8.0).abs() < 1e-12);
        assert!((r2.d[0] - 12.0).abs() < 1e-12);
    }

    #[test]
    fn dual1_sinh_cosh_tanh() {
        let x = Dual1::<1>::variable(0.5, 0);
        let s = x.sinh();
        assert!((s.v - 0.5_f64.sinh()).abs() < 1e-12);
        assert!((s.d[0] - 0.5_f64.cosh()).abs() < 1e-12);
        let c = x.cosh();
        assert!((c.v - 0.5_f64.cosh()).abs() < 1e-12);
        assert!((c.d[0] - 0.5_f64.sinh()).abs() < 1e-12);
        let t = x.tanh();
        let sech2 = 1.0 / 0.5_f64.cosh().powi(2);
        assert!((t.v - 0.5_f64.tanh()).abs() < 1e-12);
        assert!((t.d[0] - sech2).abs() < 1e-12);
    }

    #[test]
    fn dual1_hypot_gradient() {
        // hypot(3,4) = 5; d/dx hypot(x,4) = x/5 = 0.6
        let x = Dual1::<2>::variable(3.0, 0);
        let y = Dual1::<2>::variable(4.0, 1);
        let r = x.hypot(y);
        assert!((r.v - 5.0).abs() < 1e-12);
        assert!((r.d[0] - 0.6).abs() < 1e-12);
        assert!((r.d[1] - 0.8).abs() < 1e-12);
    }

    #[test]
    fn dual1_asin_acos_atan() {
        let x = Dual1::<1>::variable(0.5, 0);
        let s = x.asin();
        let expected = 1.0 / (1.0 - 0.25_f64).sqrt();
        assert!((s.v - 0.5_f64.asin()).abs() < 1e-12);
        assert!((s.d[0] - expected).abs() < 1e-12);
        let c = x.acos();
        assert!((c.v - 0.5_f64.acos()).abs() < 1e-12);
        assert!((c.d[0] - (-expected)).abs() < 1e-12);
        let t = x.atan();
        let at_d = 1.0 / (1.0 + 0.25);
        assert!((t.v - 0.5_f64.atan()).abs() < 1e-12);
        assert!((t.d[0] - at_d).abs() < 1e-12);
    }

    #[test]
    fn dual1_abs_cbrt() {
        let x = Dual1::<1>::variable(-3.0, 0);
        let a = x.abs();
        assert!((a.v - 3.0).abs() < 1e-12);
        assert!((a.d[0] - (-1.0)).abs() < 1e-12); // signum(-3) = -1
        let y = Dual1::<1>::variable(8.0, 0);
        let cb = y.cbrt();
        // d/dx cbrt(x) = 1 / (3 * cbrt(x)^2) = 1 / (3 * 4) = 1/12
        assert!((cb.v - 2.0).abs() < 1e-12);
        assert!((cb.d[0] - (1.0 / 12.0)).abs() < 1e-12);
    }

    #[test]
    fn dual1_ref_ops_compile() {
        // Verify reference-form operators compile and produce correct values.
        // We call the Add::add trait method explicitly to exercise reference impls
        // without triggering clippy::op_ref on the operator syntax.
        use core::ops::Add;
        let a = Dual1::<1>::variable(2.0, 0);
        let b = Dual1::<1>::constant(3.0);
        let add_vv = a + b;
        let add_vr: Dual1<1> = Add::add(a, &b);
        let add_rv: Dual1<1> = Add::add(&a, b);
        let add_rr: Dual1<1> = Add::add(&a, &b);
        assert_eq!(add_vv, add_vr);
        assert_eq!(add_vv, add_rv);
        assert_eq!(add_vv, add_rr);
    }

    #[test]
    fn dual1_half_pi_approx() {
        // Verify atan2 at the right angle
        let y = Dual1::<1>::variable(1.0, 0);
        let x = Dual1::<1>::constant(0.0_f64);
        let r = y.atan2(x);
        assert!((r.v - FRAC_PI_2).abs() < 1e-12);
    }
}