xad-rs 0.1.1

//! `Dual2` - dedicated second-order forward-mode AD type.
//!
//! A `Dual2<T>` carries three values:
//!   - `value`: f(x)
//!   - `d1`:    f'(x)  (first derivative w.r.t. a seed direction)
//!   - `d2`:    f''(x) (second derivative w.r.t. the same seed direction)
//!
//! Unlike full Hessian machinery, `Dual2` tracks derivatives with respect to
//! a *single* scalar variable. This is exactly what you need to compute
//! diagonal elements of a Hessian (e.g. "own-gamma" in financial risk) in a
//! single forward pass, with no tape and no finite-difference error.
//!
//! Usage:
//! ```
//! use xad_rs::dual2::Dual2;
//! // Compute f(x) = x^3 and its first/second derivatives at x = 2.
//! let x = Dual2::variable(2.0_f64);
//! let y = x * x * x;
//! assert_eq!(y.value(), 8.0);
//! assert_eq!(y.first_derivative(), 12.0); // 3x^2
//! assert_eq!(y.second_derivative(), 12.0); // 6x
//! ```

use crate::traits::Scalar;
use std::fmt;
use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};

/// Second-order forward-mode dual number.
#[derive(Clone, Copy)]
pub struct Dual2<T: Scalar> {
    value: T,
    d1: T,
    d2: T,
}

impl<T: Scalar> Dual2<T> {
    /// Create a `Dual2` with explicit value, first and second derivative.
    #[inline]
    pub fn new(value: T, d1: T, d2: T) -> Self {
        Dual2 { value, d1, d2 }
    }

    /// Create a constant (derivative-free) `Dual2`.
    #[inline]
    pub fn constant(value: T) -> Self {
        Dual2 {
            value,
            d1: T::zero(),
            d2: T::zero(),
        }
    }

    /// Create the active variable: value with `d1 = 1`, `d2 = 0`.
    /// Use this to seed the single input direction to differentiate against.
    #[inline]
    pub fn variable(value: T) -> Self {
        Dual2 {
            value,
            d1: T::one(),
            d2: T::zero(),
        }
    }

    /// Underlying value.
    #[inline]
    pub fn value(&self) -> T {
        self.value
    }

    /// First derivative (tangent) along the seeded direction.
    #[inline]
    pub fn first_derivative(&self) -> T {
        self.d1
    }

    /// Second derivative along the seeded direction.
    #[inline]
    pub fn second_derivative(&self) -> T {
        self.d2
    }

    /// Raise `self` to a scalar power `n`: `self^n`.
    ///
    /// Applies the chain rule for 2nd order: for `g(u) = u^n`,
    /// `g'(u) = n u^{n-1}`, `g''(u) = n (n-1) u^{n-2}`, then
    /// `result.d1 = g'(v) * self.d1` and
    /// `result.d2 = g''(v) * self.d1^2 + g'(v) * self.d2`.
    pub fn powf(self, n: T) -> Dual2<T> {
        let v = self.value;
        let two = T::from(2.0).unwrap();
        let vn = v.powf(n);
        let gp = n * v.powf(n - T::one());
        let gpp = n * (n - T::one()) * v.powf(n - two);
        Dual2 {
            value: vn,
            d1: gp * self.d1,
            d2: gpp * self.d1 * self.d1 + gp * self.d2,
        }
    }

    /// Exponential `exp(self)`.
    pub fn exp(self) -> Dual2<T> {
        let e = self.value.exp();
        Dual2 {
            value: e,
            d1: e * self.d1,
            d2: e * self.d1 * self.d1 + e * self.d2,
        }
    }

    /// Natural logarithm `ln(self)`.
    pub fn ln(self) -> Dual2<T> {
        // g(u) = ln(u), g'(u) = 1/u, g''(u) = -1/u^2
        let v = self.value;
        let inv = T::one() / v;
        let inv2 = inv * inv;
        Dual2 {
            value: v.ln(),
            d1: inv * self.d1,
            d2: -inv2 * self.d1 * self.d1 + inv * self.d2,
        }
    }

    /// Square root.
    pub fn sqrt(self) -> Dual2<T> {
        // g(u) = u^{1/2}, g'(u) = 1/(2*sqrt(u)), g''(u) = -1/(4*u^{3/2})
        let v = self.value;
        let s = v.sqrt();
        let half = T::from(0.5).unwrap();
        let quarter = T::from(0.25).unwrap();
        let gp = half / s;
        let gpp = -quarter / (s * v);
        Dual2 {
            value: s,
            d1: gp * self.d1,
            d2: gpp * self.d1 * self.d1 + gp * self.d2,
        }
    }

    /// Sine.
    pub fn sin(self) -> Dual2<T> {
        // g(u) = sin(u), g'(u) = cos(u), g''(u) = -sin(u)
        let v = self.value;
        let (s, c) = (v.sin(), v.cos());
        Dual2 {
            value: s,
            d1: c * self.d1,
            d2: -s * self.d1 * self.d1 + c * self.d2,
        }
    }

    /// Cosine.
    pub fn cos(self) -> Dual2<T> {
        // g(u) = cos(u), g'(u) = -sin(u), g''(u) = -cos(u)
        let v = self.value;
        let (s, c) = (v.sin(), v.cos());
        Dual2 {
            value: c,
            d1: -s * self.d1,
            d2: -c * self.d1 * self.d1 - s * self.d2,
        }
    }

    /// Error function `erf(self)`.
    ///
    /// Derivatives:
    ///   `erf'(x)  = (2/√π) · exp(-x²)`
    ///   `erf''(x) = -2x · erf'(x)`
    pub fn erf(self) -> Dual2<T> {
        let v = self.value;
        let r = crate::math::erf(v);
        let two = T::from(2.0).unwrap();
        // Convert the compile-time `f64` constant `FRAC_2_SQRT_PI` into the
        // generic scalar `T` exactly once — cheaper than computing a sqrt.
        let two_over_sqrt_pi =
            T::from(std::f64::consts::FRAC_2_SQRT_PI).unwrap();
        let gp = two_over_sqrt_pi * (-v * v).exp();
        let gpp = -two * v * gp;
        Dual2 {
            value: r,
            d1: gp * self.d1,
            d2: gpp * self.d1 * self.d1 + gp * self.d2,
        }
    }
}

// ============================================================================
// Operator implementations
// ============================================================================

// --- Add ---
impl<T: Scalar> Add for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn add(self, rhs: Self) -> Self {
        Dual2 {
            value: self.value + rhs.value,
            d1: self.d1 + rhs.d1,
            d2: self.d2 + rhs.d2,
        }
    }
}

impl<T: Scalar> Add<T> for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn add(self, rhs: T) -> Self {
        Dual2 {
            value: self.value + rhs,
            d1: self.d1,
            d2: self.d2,
        }
    }
}

impl Add<Dual2<f64>> for f64 {
    type Output = Dual2<f64>;
    #[inline]
    fn add(self, rhs: Dual2<f64>) -> Dual2<f64> {
        Dual2 {
            value: self + rhs.value,
            d1: rhs.d1,
            d2: rhs.d2,
        }
    }
}

impl Add<Dual2<f32>> for f32 {
    type Output = Dual2<f32>;
    #[inline]
    fn add(self, rhs: Dual2<f32>) -> Dual2<f32> {
        Dual2 {
            value: self + rhs.value,
            d1: rhs.d1,
            d2: rhs.d2,
        }
    }
}

// --- Sub ---
impl<T: Scalar> Sub for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn sub(self, rhs: Self) -> Self {
        Dual2 {
            value: self.value - rhs.value,
            d1: self.d1 - rhs.d1,
            d2: self.d2 - rhs.d2,
        }
    }
}

impl<T: Scalar> Sub<T> for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn sub(self, rhs: T) -> Self {
        Dual2 {
            value: self.value - rhs,
            d1: self.d1,
            d2: self.d2,
        }
    }
}

impl Sub<Dual2<f64>> for f64 {
    type Output = Dual2<f64>;
    #[inline]
    fn sub(self, rhs: Dual2<f64>) -> Dual2<f64> {
        Dual2 {
            value: self - rhs.value,
            d1: -rhs.d1,
            d2: -rhs.d2,
        }
    }
}

impl Sub<Dual2<f32>> for f32 {
    type Output = Dual2<f32>;
    #[inline]
    fn sub(self, rhs: Dual2<f32>) -> Dual2<f32> {
        Dual2 {
            value: self - rhs.value,
            d1: -rhs.d1,
            d2: -rhs.d2,
        }
    }
}

// --- Mul ---
// (a*b)' = a'*b + a*b'
// (a*b)'' = a''*b + 2*a'*b' + a*b''
impl<T: Scalar> Mul for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn mul(self, rhs: Self) -> Self {
        let two = T::from(2.0).unwrap();
        Dual2 {
            value: self.value * rhs.value,
            d1: self.d1 * rhs.value + self.value * rhs.d1,
            d2: self.d2 * rhs.value + two * self.d1 * rhs.d1 + self.value * rhs.d2,
        }
    }
}

impl<T: Scalar> Mul<T> for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn mul(self, rhs: T) -> Self {
        Dual2 {
            value: self.value * rhs,
            d1: self.d1 * rhs,
            d2: self.d2 * rhs,
        }
    }
}

impl Mul<Dual2<f64>> for f64 {
    type Output = Dual2<f64>;
    #[inline]
    fn mul(self, rhs: Dual2<f64>) -> Dual2<f64> {
        Dual2 {
            value: self * rhs.value,
            d1: self * rhs.d1,
            d2: self * rhs.d2,
        }
    }
}

impl Mul<Dual2<f32>> for f32 {
    type Output = Dual2<f32>;
    #[inline]
    fn mul(self, rhs: Dual2<f32>) -> Dual2<f32> {
        Dual2 {
            value: self * rhs.value,
            d1: self * rhs.d1,
            d2: self * rhs.d2,
        }
    }
}

// --- Div ---
// a/b = a * (1/b)
// (1/b)'  = -b'/b^2
// (1/b)'' = 2 b'^2/b^3 - b''/b^2
impl<T: Scalar> Div for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn div(self, rhs: Self) -> Self {
        let two = T::from(2.0).unwrap();
        let inv_b = T::one() / rhs.value;
        let inv_b2 = inv_b * inv_b;
        let inv_b3 = inv_b2 * inv_b;
        let recip = Dual2 {
            value: inv_b,
            d1: -rhs.d1 * inv_b2,
            d2: two * rhs.d1 * rhs.d1 * inv_b3 - rhs.d2 * inv_b2,
        };
        self * recip
    }
}

impl<T: Scalar> Div<T> for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn div(self, rhs: T) -> Self {
        let inv = T::one() / rhs;
        Dual2 {
            value: self.value * inv,
            d1: self.d1 * inv,
            d2: self.d2 * inv,
        }
    }
}

impl Div<Dual2<f64>> for f64 {
    type Output = Dual2<f64>;
    #[inline]
    fn div(self, rhs: Dual2<f64>) -> Dual2<f64> {
        Dual2::constant(self) / rhs
    }
}

impl Div<Dual2<f32>> for f32 {
    type Output = Dual2<f32>;
    #[inline]
    fn div(self, rhs: Dual2<f32>) -> Dual2<f32> {
        Dual2::constant(self) / rhs
    }
}

// --- Neg ---
impl<T: Scalar> Neg for Dual2<T> {
    type Output = Dual2<T>;
    #[inline]
    fn neg(self) -> Self {
        Dual2 {
            value: -self.value,
            d1: -self.d1,
            d2: -self.d2,
        }
    }
}

// --- Compound assignment ---
impl<T: Scalar> AddAssign for Dual2<T> {
    #[inline]
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl<T: Scalar> AddAssign<T> for Dual2<T> {
    #[inline]
    fn add_assign(&mut self, rhs: T) {
        self.value = self.value + rhs;
    }
}

impl<T: Scalar> SubAssign for Dual2<T> {
    #[inline]
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl<T: Scalar> SubAssign<T> for Dual2<T> {
    #[inline]
    fn sub_assign(&mut self, rhs: T) {
        self.value = self.value - rhs;
    }
}

impl<T: Scalar> MulAssign for Dual2<T> {
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl<T: Scalar> MulAssign<T> for Dual2<T> {
    #[inline]
    fn mul_assign(&mut self, rhs: T) {
        *self = *self * rhs;
    }
}

impl<T: Scalar> DivAssign for Dual2<T> {
    #[inline]
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs;
    }
}

impl<T: Scalar> DivAssign<T> for Dual2<T> {
    #[inline]
    fn div_assign(&mut self, rhs: T) {
        *self = *self / rhs;
    }
}

// --- Display / Debug / Default ---
impl<T: Scalar> fmt::Display for Dual2<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.value)
    }
}

impl<T: Scalar> fmt::Debug for Dual2<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Dual2(v={}, d1={}, d2={})",
            self.value, self.d1, self.d2
        )
    }
}

impl<T: Scalar> Default for Dual2<T> {
    fn default() -> Self {
        Dual2::constant(T::zero())
    }
}

impl<T: Scalar> From<T> for Dual2<T> {
    fn from(value: T) -> Self {
        Dual2::constant(value)
    }
}