oxiproj-core 0.1.2

//! Second-order forward-mode dual numbers for automatic differentiation.
//!
//! `Dual2<N>` extends the first-order approach with diagonal second-derivative
//! tracking: `d2[i] = partial^2 f / partial xi^2`. This enables computing
//! Hessian diagonal entries (and curvatures) alongside the gradient in a
//! single forward pass.
//!
//! # Arithmetic second-derivative rules
//!
//! - Add/Sub: trivial componentwise.
//! - Mul: `(a*b).d2[i] = a.d2[i]*b.v + 2*a.d1[i]*b.d1[i] + a.v*b.d2[i]`
//! - Div: `(a/b).d2[i] = (a.d2[i]*b.v - a.v*b.d2[i]) / b.v^2`
//!   `- 2*b.d1[i]*(a.d1[i]*b.v - a.v*b.d1[i]) / b.v^3`

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

/// Second-order forward-mode dual number.
///
/// Tracks: primal value `v`, first derivatives `d1[i] = partial f / partial xi`,
/// and diagonal second derivatives `d2[i] = partial^2 f / partial xi^2`.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Dual2<const N: usize> {
    /// Primal value.
    pub v: f64,
    /// First derivatives (gradient).
    pub d1: [f64; N],
    /// Second derivatives (diagonal Hessian).
    pub d2: [f64; N],
}

/// Apply a unary chain-rule: `d1_out[i] = fp * d1[i]`, `d2_out[i] = fpp * d1[i]^2 + fp * d2[i]`.
#[inline]
fn chain1<const N: usize>(d1: &[f64; N], d2: &[f64; N], fp: f64, fpp: f64) -> ([f64; N], [f64; N]) {
    let mut out_d1 = [0.0; N];
    let mut out_d2 = [0.0; N];
    for ((o1, o2), (&g1, &g2)) in out_d1
        .iter_mut()
        .zip(out_d2.iter_mut())
        .zip(d1.iter().zip(d2.iter()))
    {
        *o1 = fp * g1;
        *o2 = fpp * g1 * g1 + fp * g2;
    }
    (out_d1, out_d2)
}

impl<const N: usize> Dual2<N> {
    /// Construct a constant (zero first and second derivatives).
    #[inline]
    pub fn constant(v: f64) -> Self {
        Self {
            v,
            d1: [0.0; N],
            d2: [0.0; N],
        }
    }

    /// Construct an independent variable (`d1[idx] = 1.0`, `d2 = 0` everywhere).
    ///
    /// # Panics
    ///
    /// Panics if `idx >= N`.
    #[inline]
    pub fn variable(v: f64, idx: usize) -> Self {
        assert!(
            idx < N,
            "Dual2::variable: idx {idx} is out of bounds for N={N}"
        );
        let mut d1 = [0.0; N];
        d1[idx] = 1.0;
        Self {
            v,
            d1,
            d2: [0.0; N],
        }
    }

    // ---- Transcendental methods ----
    // Convention: g = self; fp = f'(g.v); fpp = f''(g.v)
    // First:  (f(g)).d1[i] = fp * g.d1[i]
    // Second: (f(g)).d2[i] = fpp * g.d1[i]^2 + fp * g.d2[i]

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

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

    /// tan(self)
    #[inline]
    pub fn tan(self) -> Self {
        let tan_v = self.v.tan();
        let cos_v = self.v.cos();
        let sec2 = 1.0 / (cos_v * cos_v);
        let f_pp = 2.0 * sec2 * tan_v;
        let (d1, d2) = chain1(&self.d1, &self.d2, sec2, f_pp);
        Self { v: tan_v, d1, d2 }
    }

    /// asin(self)
    #[inline]
    pub fn asin(self) -> Self {
        let one_m_v2 = 1.0 - self.v * self.v;
        let sqrt_v = one_m_v2.sqrt();
        let fp = 1.0 / sqrt_v;
        let fpp = self.v / (one_m_v2 * sqrt_v);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self {
            v: self.v.asin(),
            d1,
            d2,
        }
    }

    /// acos(self)
    #[inline]
    pub fn acos(self) -> Self {
        let one_m_v2 = 1.0 - self.v * self.v;
        let sqrt_v = one_m_v2.sqrt();
        let fp = -1.0 / sqrt_v;
        let fpp = -self.v / (one_m_v2 * sqrt_v);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self {
            v: self.v.acos(),
            d1,
            d2,
        }
    }

    /// atan(self)
    #[inline]
    pub fn atan(self) -> Self {
        let one_p_v2 = 1.0 + self.v * self.v;
        let fp = 1.0 / one_p_v2;
        let fpp = -2.0 * self.v / (one_p_v2 * one_p_v2);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self {
            v: self.v.atan(),
            d1,
            d2,
        }
    }

    /// atan2(y=self, x=rhs)
    ///
    /// For scalar-function derivatives we treat atan2 as f(y,x) and propagate
    /// through the combined expression for each direction i independently.
    #[inline]
    pub fn atan2(self, rhs: Dual2<N>) -> Self {
        let r2 = rhs.v * rhs.v + self.v * self.v;
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), ((&yd1, &yd2), (&xd1, &xd2))) in d1.iter_mut().zip(d2.iter_mut()).zip(
            self.d1
                .iter()
                .zip(self.d2.iter())
                .zip(rhs.d1.iter().zip(rhs.d2.iter())),
        ) {
            let num = rhs.v * yd1 - self.v * xd1;
            *o1 = num / r2;
            let num2 = rhs.v * yd2 - self.v * xd2;
            let denom_d = 2.0 * (rhs.v * xd1 + self.v * yd1);
            *o2 = (num2 * r2 - num * denom_d) / (r2 * r2);
        }
        Self {
            v: self.v.atan2(rhs.v),
            d1,
            d2,
        }
    }

    /// sqrt(self)
    #[inline]
    pub fn sqrt(self) -> Self {
        let sqrt_v = self.v.sqrt();
        let fp = 1.0 / (2.0 * sqrt_v);
        let fpp = -1.0 / (4.0 * self.v * sqrt_v);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self { v: sqrt_v, d1, d2 }
    }

    /// cbrt(self)
    #[inline]
    pub fn cbrt(self) -> Self {
        let cbrt_v = self.v.cbrt();
        let fp = 1.0 / (3.0 * cbrt_v * cbrt_v);
        let fpp = -2.0 / (9.0 * self.v * cbrt_v * cbrt_v);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self { v: cbrt_v, d1, d2 }
    }

    /// exp(self)
    #[inline]
    pub fn exp(self) -> Self {
        let exp_v = self.v.exp();
        // fp = fpp = exp(v)
        let (d1, d2) = chain1(&self.d1, &self.d2, exp_v, exp_v);
        Self { v: exp_v, d1, d2 }
    }

    /// ln(self)
    #[inline]
    pub fn ln(self) -> Self {
        let fp = 1.0 / self.v;
        let fpp = -1.0 / (self.v * self.v);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self {
            v: self.v.ln(),
            d1,
            d2,
        }
    }

    /// abs(self) — chain rule: d/dx |f| = f'*signum(f); d2 ignores the delta at 0
    #[inline]
    pub fn abs(self) -> Self {
        let sign = self.v.signum();
        let (d1, d2) = chain1(&self.d1, &self.d2, sign, 0.0);
        Self {
            v: self.v.abs(),
            d1,
            d2,
        }
    }

    /// powi(self, n)
    #[inline]
    pub fn powi(self, n: i32) -> Self {
        let fp = n as f64 * self.v.powi(n - 1);
        let fpp = (n as f64) * ((n - 1) as f64) * self.v.powi(n - 2);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self {
            v: self.v.powi(n),
            d1,
            d2,
        }
    }

    /// powf(self, e)
    #[inline]
    pub fn powf(self, e: f64) -> Self {
        let fp = e * self.v.powf(e - 1.0);
        let fpp = e * (e - 1.0) * self.v.powf(e - 2.0);
        let (d1, d2) = chain1(&self.d1, &self.d2, fp, fpp);
        Self {
            v: self.v.powf(e),
            d1,
            d2,
        }
    }

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

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

    /// tanh(self)
    #[inline]
    pub fn tanh(self) -> Self {
        let tanh_v = self.v.tanh();
        let cosh_v = self.v.cosh();
        let sech2 = 1.0 / (cosh_v * cosh_v);
        let fpp = -2.0 * sech2 * tanh_v;
        let (d1, d2) = chain1(&self.d1, &self.d2, sech2, fpp);
        Self { v: tanh_v, d1, d2 }
    }

    /// hypot(a=self, b=rhs)
    #[inline]
    pub fn hypot(self, rhs: Dual2<N>) -> Self {
        let h = self.v.hypot(rhs.v);
        let h3 = h * h * h;
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), ((&ad1, &ad2), (&bd1, &bd2))) in d1.iter_mut().zip(d2.iter_mut()).zip(
            self.d1
                .iter()
                .zip(self.d2.iter())
                .zip(rhs.d1.iter().zip(rhs.d2.iter())),
        ) {
            let num1 = self.v * ad1 + rhs.v * bd1;
            *o1 = num1 / h;
            let d_num1 = ad1 * ad1 + self.v * ad2 + bd1 * bd1 + rhs.v * bd2;
            *o2 = (d_num1 * h * h - num1 * num1) / h3;
        }
        Self { v: h, d1, d2 }
    }
}

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

impl<const N: usize> Add for Dual2<N> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: Self) -> Self {
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), ((&a1, &a2), (&b1, &b2))) in d1.iter_mut().zip(d2.iter_mut()).zip(
            self.d1
                .iter()
                .zip(self.d2.iter())
                .zip(rhs.d1.iter().zip(rhs.d2.iter())),
        ) {
            *o1 = a1 + b1;
            *o2 = a2 + b2;
        }
        Self {
            v: self.v + rhs.v,
            d1,
            d2,
        }
    }
}

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

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

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

impl<const N: usize> Sub for Dual2<N> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: Self) -> Self {
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), ((&a1, &a2), (&b1, &b2))) in d1.iter_mut().zip(d2.iter_mut()).zip(
            self.d1
                .iter()
                .zip(self.d2.iter())
                .zip(rhs.d1.iter().zip(rhs.d2.iter())),
        ) {
            *o1 = a1 - b1;
            *o2 = a2 - b2;
        }
        Self {
            v: self.v - rhs.v,
            d1,
            d2,
        }
    }
}

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

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

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

impl<const N: usize> Mul for Dual2<N> {
    type Output = Self;
    /// `(a*b).d2[i] = a.d2[i]*b.v + 2*a.d1[i]*b.d1[i] + a.v*b.d2[i]`
    #[inline]
    fn mul(self, rhs: Self) -> Self {
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), ((&a1, &a2), (&b1, &b2))) in d1.iter_mut().zip(d2.iter_mut()).zip(
            self.d1
                .iter()
                .zip(self.d2.iter())
                .zip(rhs.d1.iter().zip(rhs.d2.iter())),
        ) {
            *o1 = a1 * rhs.v + self.v * b1;
            *o2 = a2 * rhs.v + 2.0 * a1 * b1 + self.v * b2;
        }
        Self {
            v: self.v * rhs.v,
            d1,
            d2,
        }
    }
}

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

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

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

impl<const N: usize> Div for Dual2<N> {
    type Output = Self;
    /// `(a/b).d2[i] = (a.d2[i]*b.v - a.v*b.d2[i]) / b.v^2`
    ///   `- 2*b.d1[i]*(a.d1[i]*b.v - a.v*b.d1[i]) / b.v^3`
    #[inline]
    fn div(self, rhs: Self) -> Self {
        let bv = rhs.v;
        let b2 = bv * bv;
        let b3 = b2 * bv;
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), ((&a1, &a2), (&b1, &b2_d))) in d1.iter_mut().zip(d2.iter_mut()).zip(
            self.d1
                .iter()
                .zip(self.d2.iter())
                .zip(rhs.d1.iter().zip(rhs.d2.iter())),
        ) {
            *o1 = (a1 * bv - self.v * b1) / b2;
            let t1 = (a2 * bv - self.v * b2_d) / b2;
            let t2 = 2.0 * b1 * (a1 * bv - self.v * b1) / b3;
            *o2 = t1 - t2;
        }
        Self {
            v: self.v / rhs.v,
            d1,
            d2,
        }
    }
}

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

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

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

impl<const N: usize> Neg for Dual2<N> {
    type Output = Self;
    #[inline]
    fn neg(self) -> Self {
        let mut d1 = [0.0; N];
        let mut d2 = [0.0; N];
        for ((o1, o2), (&g1, &g2)) in d1
            .iter_mut()
            .zip(d2.iter_mut())
            .zip(self.d1.iter().zip(self.d2.iter()))
        {
            *o1 = -g1;
            *o2 = -g2;
        }
        Self { v: -self.v, d1, d2 }
    }
}

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

impl<const N: usize> PartialOrd for Dual2<N> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        self.v.partial_cmp(&other.v)
    }
}

impl<const N: usize> fmt::Display for Dual2<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;

    #[test]
    fn dual2_second_derivative_cubic() {
        // d^2/dx^2 x^3 = 6x, at x=2: = 12
        // d/dx x^3 = 3x^2, at x=2: = 12
        let x = Dual2::<1>::variable(2.0, 0);
        let x3 = x * x * x;
        assert!((x3.v - 8.0).abs() < 1e-12);
        assert!((x3.d1[0] - 12.0).abs() < 1e-12);
        assert!((x3.d2[0] - 12.0).abs() < 1e-12);
    }

    #[test]
    fn dual2_sin_second_derivative() {
        // d^2/dx^2 sin(x) = -sin(x)
        let x = Dual2::<1>::variable(0.5, 0);
        let s = x.sin();
        assert!((s.v - 0.5_f64.sin()).abs() < 1e-12);
        assert!((s.d1[0] - 0.5_f64.cos()).abs() < 1e-12);
        assert!((s.d2[0] - (-0.5_f64.sin())).abs() < 1e-12);
    }

    #[test]
    fn dual2_exp_second_derivative() {
        // d^2/dx^2 exp(x) = exp(x)
        let x = Dual2::<1>::variable(1.0, 0);
        let e = x.exp();
        assert!((e.v - 1.0_f64.exp()).abs() < 1e-12);
        assert!((e.d1[0] - 1.0_f64.exp()).abs() < 1e-12);
        assert!((e.d2[0] - 1.0_f64.exp()).abs() < 1e-12);
    }

    #[test]
    fn dual2_product_rule_second_order() {
        // (x * x).d2 = 2.0 (second derivative of x^2 is 2)
        let x = Dual2::<1>::variable(3.0, 0);
        let x2 = x * x;
        assert!((x2.v - 9.0).abs() < 1e-12);
        assert!((x2.d1[0] - 6.0).abs() < 1e-12);
        assert!((x2.d2[0] - 2.0).abs() < 1e-12);
    }

    #[test]
    fn dual2_ln_second_derivative() {
        // d^2/dx^2 ln(x) = -1/x^2, at x=2 => -0.25
        let x = Dual2::<1>::variable(2.0, 0);
        let l = x.ln();
        assert!((l.v - 2.0_f64.ln()).abs() < 1e-12);
        assert!((l.d1[0] - 0.5).abs() < 1e-12);
        assert!((l.d2[0] - (-0.25)).abs() < 1e-12);
    }

    #[test]
    fn dual2_sqrt_second_derivative() {
        // d^2/dx^2 sqrt(x) = -1/(4*x^(3/2)), at x=4 => -1/32
        let x = Dual2::<1>::variable(4.0, 0);
        let s = x.sqrt();
        assert!((s.v - 2.0).abs() < 1e-12);
        assert!((s.d1[0] - 0.25).abs() < 1e-12);
        assert!((s.d2[0] - (-1.0 / 32.0)).abs() < 1e-12);
    }

    #[test]
    fn dual2_constant_zero_derivatives() {
        let c = Dual2::<2>::constant(7.0);
        assert!((c.v - 7.0).abs() < 1e-14);
        assert_eq!(c.d1, [0.0; 2]);
        assert_eq!(c.d2, [0.0; 2]);
    }

    #[test]
    #[cfg(not(feature = "no_std"))]
    fn dual2_variable_panics_out_of_bounds() {
        let result = std::panic::catch_unwind(|| {
            Dual2::<2>::variable(1.0, 2);
        });
        assert!(result.is_err());
    }

    #[test]
    fn dual2_div_second_derivative() {
        // d^2/dx^2 (x / c) = 0 for constant c
        let x = Dual2::<1>::variable(5.0, 0);
        let c = Dual2::<1>::constant(2.0);
        let r = x / c;
        assert!((r.v - 2.5).abs() < 1e-12);
        assert!((r.d1[0] - 0.5).abs() < 1e-12);
        assert!((r.d2[0]).abs() < 1e-12);
    }

    #[test]
    fn dual2_cos_second_derivative() {
        // d^2/dx^2 cos(x) = -cos(x)
        let x = Dual2::<1>::variable(1.0, 0);
        let c = x.cos();
        assert!((c.v - 1.0_f64.cos()).abs() < 1e-12);
        assert!((c.d1[0] - (-1.0_f64.sin())).abs() < 1e-12);
        assert!((c.d2[0] - (-1.0_f64.cos())).abs() < 1e-12);
    }

    #[test]
    fn dual2_display() {
        use core::f64::consts::E;
        let x = Dual2::<1>::variable(E, 0);
        let s = format!("{x}");
        assert!(!s.is_empty());
    }

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

    #[test]
    fn dual2_neg() {
        let x = Dual2::<1>::variable(3.0, 0);
        let nx = -x;
        assert!((nx.v - (-3.0)).abs() < 1e-14);
        assert!((nx.d1[0] - (-1.0)).abs() < 1e-14);
        assert!(nx.d2[0].abs() < 1e-14);
    }

    #[test]
    fn dual2_sub() {
        let x = Dual2::<1>::variable(5.0, 0);
        let y = Dual2::<1>::constant(2.0);
        let r = x - y;
        assert!((r.v - 3.0).abs() < 1e-14);
        assert!((r.d1[0] - 1.0).abs() < 1e-14);
        assert!(r.d2[0].abs() < 1e-14);
    }

    #[test]
    fn dual2_ref_ops_compile() {
        // Exercise reference-form operators via explicit trait method calls to avoid
        // triggering clippy::op_ref on the operator syntax (Copy types).
        use core::ops::Add;
        let a = Dual2::<1>::variable(2.0, 0);
        let b = Dual2::<1>::constant(3.0);
        let add_vv = a + b;
        let add_vr: Dual2<1> = Add::add(a, &b);
        let add_rv: Dual2<1> = Add::add(&a, b);
        let add_rr: Dual2<1> = Add::add(&a, &b);
        assert_eq!(add_vv, add_vr);
        assert_eq!(add_vv, add_rv);
        assert_eq!(add_vv, add_rr);
    }
}