#![allow(clippy::suspicious_arithmetic_impl)]
use crate::passive::Passive;
use ndarray::{Array1, Array2};
use std::ops::{Add, Div, Mul, Sub};
#[derive(Clone, Debug, PartialEq)]
pub struct Jet2Vec<T: Passive = f64> {
pub value: T,
pub grad: Array1<T>,
pub hess: Array2<T>,
}
impl<T: Passive> Jet2Vec<T> {
#[inline]
pub fn variable(value: T, i: usize, n: usize) -> Self {
let mut grad = Array1::<T>::zeros(n);
grad[i] = T::one();
let hess = Array2::<T>::zeros((n, n));
Self { value, grad, hess }
}
#[inline]
pub fn constant(value: T, n: usize) -> Self {
Self {
value,
grad: Array1::<T>::zeros(n),
hess: Array2::<T>::zeros((n, n)),
}
}
#[inline]
pub fn value(&self) -> T {
self.value
}
#[inline]
pub fn gradient(&self) -> &Array1<T> {
&self.grad
}
#[inline]
pub fn hessian(&self) -> &Array2<T> {
&self.hess
}
#[inline]
pub fn len(&self) -> usize {
self.grad.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.grad.is_empty()
}
}
#[allow(dead_code)]
#[inline]
pub(crate) fn mirror_upper_triangle<T: Passive>(h: &mut Array2<T>) {
let n = h.nrows();
debug_assert_eq!(n, h.ncols(), "mirror_upper_triangle on non-square matrix");
for i in 0..n {
for j in (i + 1)..n {
h[[j, i]] = h[[i, j]];
}
}
}
#[inline]
fn shape_check<T: Passive>(a: &Jet2Vec<T>, b: &Jet2Vec<T>) {
debug_assert_eq!(
a.grad.len(),
b.grad.len(),
"Jet2Vec gradient length mismatch: {} vs {}",
a.grad.len(),
b.grad.len()
);
debug_assert_eq!(
a.hess.shape(),
b.hess.shape(),
"Jet2Vec hessian shape mismatch"
);
}
impl<T: Passive> Add<&Jet2Vec<T>> for &Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn add(self, rhs: &Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(self, rhs);
Jet2Vec {
value: self.value + rhs.value,
grad: &self.grad + &rhs.grad,
hess: &self.hess + &rhs.hess,
}
}
}
impl<T: Passive> Add for Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn add(mut self, rhs: Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(&self, &rhs);
self.value += rhs.value;
self.grad += &rhs.grad;
self.hess += &rhs.hess;
self
}
}
impl<T: Passive> Sub<&Jet2Vec<T>> for &Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn sub(self, rhs: &Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(self, rhs);
Jet2Vec {
value: self.value - rhs.value,
grad: &self.grad - &rhs.grad,
hess: &self.hess - &rhs.hess,
}
}
}
impl<T: Passive> Sub for Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn sub(mut self, rhs: Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(&self, &rhs);
self.value -= rhs.value;
self.grad -= &rhs.grad;
self.hess -= &rhs.hess;
self
}
}
impl<T: Passive> Mul<&Jet2Vec<T>> for &Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn mul(self, rhs: &Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(self, rhs);
let a_val = self.value;
let b_val = rhs.value;
let n = self.grad.len();
let value = a_val * b_val;
let grad: Array1<T> = self
.grad
.iter()
.zip(rhs.grad.iter())
.map(|(&ga, &gb)| a_val * gb + b_val * ga)
.collect();
let mut hess = Array2::<T>::zeros((n, n));
write_mul_hess(
&mut hess,
&self.hess,
&rhs.hess,
&self.grad,
&rhs.grad,
a_val,
b_val,
n,
);
Jet2Vec { value, grad, hess }
}
}
impl<T: Passive> Mul for Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn mul(mut self, rhs: Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(&self, &rhs);
let a_val = self.value;
let b_val = rhs.value;
let n = self.grad.len();
self.value = a_val * b_val;
accumulate_mul_hess(
&mut self.hess,
&rhs.hess,
&self.grad,
&rhs.grad,
a_val,
b_val,
n,
);
for (sg, &rg) in self.grad.iter_mut().zip(rhs.grad.iter()) {
*sg = a_val * rg + b_val * *sg;
}
self
}
}
#[inline]
#[allow(clippy::too_many_arguments)]
fn write_mul_hess<T: Passive>(
out: &mut Array2<T>,
sa: &Array2<T>,
sb: &Array2<T>,
ga: &Array1<T>,
gb: &Array1<T>,
a_val: T,
b_val: T,
n: usize,
) {
let h = out
.as_slice_mut()
.expect("Jet2Vec hess must be standard-layout");
let sa = sa
.as_slice()
.expect("Jet2Vec hess must be standard-layout");
let sb = sb
.as_slice()
.expect("Jet2Vec hess must be standard-layout");
let ga = ga.as_slice().expect("Jet2Vec grad must be contiguous");
let gb = gb.as_slice().expect("Jet2Vec grad must be contiguous");
for i in 0..n {
let off = i * n;
let ga_i = ga[i];
let gb_i = gb[i];
for j in i..n {
let cross = ga_i * gb[j] + gb_i * ga[j];
let v = a_val * sb[off + j] + b_val * sa[off + j] + cross;
h[off + j] = v;
h[j * n + i] = v;
}
}
}
#[inline]
fn accumulate_mul_hess<T: Passive>(
lhs: &mut Array2<T>,
rhs: &Array2<T>,
ga: &Array1<T>,
gb: &Array1<T>,
a_val: T,
b_val: T,
n: usize,
) {
let h = lhs
.as_slice_mut()
.expect("Jet2Vec hess must be standard-layout");
let sb = rhs
.as_slice()
.expect("Jet2Vec hess must be standard-layout");
let ga = ga.as_slice().expect("Jet2Vec grad must be contiguous");
let gb = gb.as_slice().expect("Jet2Vec grad must be contiguous");
for i in 0..n {
let off = i * n;
let ga_i = ga[i];
let gb_i = gb[i];
for j in i..n {
let cross = ga_i * gb[j] + gb_i * ga[j];
let v = a_val * sb[off + j] + b_val * h[off + j] + cross;
h[off + j] = v;
h[j * n + i] = v;
}
}
}
impl<T: Passive> Div<&Jet2Vec<T>> for &Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn div(self, rhs: &Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(self, rhs);
let b_val = rhs.value;
debug_assert!(b_val != T::zero(), "Jet2Vec::div by zero value");
let inv_b = T::one() / b_val;
let n = self.grad.len();
let value = self.value * inv_b;
let grad: Array1<T> = self
.grad
.iter()
.zip(rhs.grad.iter())
.map(|(&ga, &gb)| (ga - value * gb) * inv_b)
.collect();
let mut hess = Array2::<T>::zeros((n, n));
write_div_hess(
&mut hess,
&self.hess,
&rhs.hess,
&grad,
&rhs.grad,
value,
inv_b,
n,
);
Jet2Vec { value, grad, hess }
}
}
impl<T: Passive> Div for Jet2Vec<T> {
type Output = Jet2Vec<T>;
#[inline]
fn div(mut self, rhs: Jet2Vec<T>) -> Jet2Vec<T> {
shape_check(&self, &rhs);
let b_val = rhs.value;
debug_assert!(b_val != T::zero(), "Jet2Vec::div by zero value");
let inv_b = T::one() / b_val;
let n = self.grad.len();
let value = self.value * inv_b;
self.value = value;
for (sg, &rg) in self.grad.iter_mut().zip(rhs.grad.iter()) {
*sg = (*sg - value * rg) * inv_b;
}
accumulate_div_hess(
&mut self.hess,
&rhs.hess,
&self.grad,
&rhs.grad,
value,
inv_b,
n,
);
self
}
}
#[inline]
#[allow(clippy::too_many_arguments)]
fn write_div_hess<T: Passive>(
out: &mut Array2<T>,
sa: &Array2<T>,
sb: &Array2<T>,
g_quot: &Array1<T>,
gb: &Array1<T>,
value: T,
inv_b: T,
n: usize,
) {
let h = out
.as_slice_mut()
.expect("Jet2Vec hess must be standard-layout");
let sa = sa
.as_slice()
.expect("Jet2Vec hess must be standard-layout");
let sb = sb
.as_slice()
.expect("Jet2Vec hess must be standard-layout");
let gq = g_quot.as_slice().expect("Jet2Vec grad must be contiguous");
let gb = gb.as_slice().expect("Jet2Vec grad must be contiguous");
for i in 0..n {
let off = i * n;
let gq_i = gq[i];
let gb_i = gb[i];
for j in i..n {
let cross = gq_i * gb[j] + gb_i * gq[j];
let v = (sa[off + j] - value * sb[off + j] - cross) * inv_b;
h[off + j] = v;
h[j * n + i] = v;
}
}
}
#[inline]
fn accumulate_div_hess<T: Passive>(
lhs: &mut Array2<T>,
rhs: &Array2<T>,
g_quot: &Array1<T>,
gb: &Array1<T>,
value: T,
inv_b: T,
n: usize,
) {
let h = lhs
.as_slice_mut()
.expect("Jet2Vec hess must be standard-layout");
let sb = rhs
.as_slice()
.expect("Jet2Vec hess must be standard-layout");
let gq = g_quot.as_slice().expect("Jet2Vec grad must be contiguous");
let gb = gb.as_slice().expect("Jet2Vec grad must be contiguous");
for i in 0..n {
let off = i * n;
let gq_i = gq[i];
let gb_i = gb[i];
for j in i..n {
let cross = gq_i * gb[j] + gb_i * gq[j];
let v = (h[off + j] - value * sb[off + j] - cross) * inv_b;
h[off + j] = v;
h[j * n + i] = v;
}
}
}
macro_rules! impl_unary_dual2vec {
($name:ident, $f:expr, $fp:expr, $fpp:expr) => {
impl<T: Passive> Jet2Vec<T> {
#[doc = concat!("Applies `", stringify!($name), "` elementwise via the multi-variable chain rule.")]
#[doc = ""]
#[doc = "Propagates value, gradient, and Hessian in one forward pass. See"]
#[doc = "`src/forward/jet2vec.rs` module docs for the chain-rule derivation."]
#[inline]
pub fn $name(mut self) -> Self {
let u = self.value;
let fv: T = ($f)(u);
let fp: T = ($fp)(u);
let fpp: T = ($fpp)(u);
let n = self.grad.len();
self.value = fv;
accumulate_unary_hess(&mut self.hess, &self.grad, fp, fpp, n);
for d in self.grad.iter_mut() {
*d *= fp;
}
self
}
}
};
}
#[inline]
fn accumulate_unary_hess<T: Passive>(
lhs: &mut Array2<T>,
grad: &Array1<T>,
fp: T,
fpp: T,
n: usize,
) {
let h = lhs
.as_slice_mut()
.expect("Jet2Vec hess must be standard-layout");
let g = grad.as_slice().expect("Jet2Vec grad must be contiguous");
for i in 0..n {
let off = i * n;
let g_i = g[i];
for j in i..n {
let v = fp * h[off + j] + fpp * g_i * g[j];
h[off + j] = v;
h[j * n + i] = v;
}
}
}
#[allow(unused_imports)]
pub(crate) use impl_unary_dual2vec;
impl_unary_dual2vec!(sin, |u: T| u.sin(), |u: T| u.cos(), |u: T| -u.sin());
impl_unary_dual2vec!(cos, |u: T| u.cos(), |u: T| -u.sin(), |u: T| -u.cos());
impl_unary_dual2vec!(
tan,
|u: T| u.tan(),
|u: T| {
let t = u.tan();
T::one() + t * t
},
|u: T| {
let t = u.tan();
T::from(2.0).unwrap() * t * (T::one() + t * t)
}
);
impl_unary_dual2vec!(exp, |u: T| u.exp(), |u: T| u.exp(), |u: T| u.exp());
impl_unary_dual2vec!(
ln,
|u: T| u.ln(),
|u: T| T::one() / u,
|u: T| -T::one() / (u * u)
);
impl_unary_dual2vec!(
sqrt,
|u: T| u.sqrt(),
|u: T| T::from(0.5).unwrap() / u.sqrt(),
|u: T| -T::from(0.25).unwrap() / (u * u.sqrt())
);
impl_unary_dual2vec!(
tanh,
|u: T| u.tanh(),
|u: T| {
let t = u.tanh();
T::one() - t * t
},
|u: T| {
let t = u.tanh();
-T::from(2.0).unwrap() * t * (T::one() - t * t)
}
);
impl_unary_dual2vec!(
atan,
|u: T| u.atan(),
|u: T| T::one() / (T::one() + u * u),
|u: T| {
let d = T::one() + u * u;
-T::from(2.0).unwrap() * u / (d * d)
}
);
impl_unary_dual2vec!(
asinh,
|u: T| u.asinh(),
|u: T| T::one() / (T::one() + u * u).sqrt(),
|u: T| {
let d = T::one() + u * u;
-u / (d * d.sqrt())
}
);
impl_unary_dual2vec!(
erf,
|u: T| crate::math::erf::<T>(u),
|u: T| T::from(std::f64::consts::FRAC_2_SQRT_PI).unwrap() * (-u * u).exp(),
|u: T| {
let fp = T::from(std::f64::consts::FRAC_2_SQRT_PI).unwrap() * (-u * u).exp();
-T::from(2.0).unwrap() * u * fp
}
);
impl_unary_dual2vec!(
norm_cdf,
|u: T| crate::math::norm_cdf::<T>(u),
|u: T| crate::math::norm_pdf::<T>(u),
|u: T| -u * crate::math::norm_pdf::<T>(u)
);
impl_unary_dual2vec!(
inv_norm_cdf,
|u: T| crate::math::inv_norm_cdf::<T>(u),
|u: T| {
let r = crate::math::inv_norm_cdf::<T>(u);
T::one() / crate::math::norm_pdf::<T>(r)
},
|u: T| {
let r = crate::math::inv_norm_cdf::<T>(u);
let gp = T::one() / crate::math::norm_pdf::<T>(r);
r * gp * gp
}
);
impl<T: Passive> Jet2Vec<T> {
#[inline]
pub fn powf(mut self, k: T) -> Self {
let u = self.value;
let two = T::from(2.0).unwrap();
let fv = u.powf(k);
let fp = k * u.powf(k - T::one());
let fpp = k * (k - T::one()) * u.powf(k - two);
let n = self.grad.len();
self.value = fv;
accumulate_unary_hess(&mut self.hess, &self.grad, fp, fpp, n);
for d in self.grad.iter_mut() {
*d *= fp;
}
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::array;
#[test]
fn test_variable_constructor() {
let x = Jet2Vec::variable(3.0_f64, 0, 2);
assert_eq!(x.value(), 3.0);
assert_eq!(x.gradient(), &array![1.0, 0.0]);
assert_eq!(x.hessian(), &Array2::<f64>::zeros((2, 2)));
}
#[test]
fn test_variable_constructor_index_1() {
let y = Jet2Vec::variable(5.0_f64, 1, 3);
assert_eq!(y.value(), 5.0);
assert_eq!(y.gradient(), &array![0.0, 1.0, 0.0]);
assert_eq!(y.hessian(), &Array2::<f64>::zeros((3, 3)));
}
#[test]
fn test_constant_constructor() {
let c = Jet2Vec::constant(7.5_f64, 4);
assert_eq!(c.value(), 7.5);
assert_eq!(c.gradient(), &Array1::<f64>::zeros(4));
assert_eq!(c.hessian(), &Array2::<f64>::zeros((4, 4)));
assert_eq!(c.len(), 4);
assert!(!c.is_empty());
}
#[test]
fn test_mirror_upper_triangle_3x3() {
let mut h = Array2::<f64>::zeros((3, 3));
h[[0, 0]] = 1.0;
h[[1, 1]] = 2.0;
h[[2, 2]] = 3.0;
h[[0, 1]] = 4.0;
h[[0, 2]] = 5.0;
h[[1, 2]] = 6.0;
mirror_upper_triangle(&mut h);
assert_eq!(h[[1, 0]], 4.0);
assert_eq!(h[[2, 0]], 5.0);
assert_eq!(h[[2, 1]], 6.0);
assert_eq!(h[[0, 0]], 1.0);
assert_eq!(h[[1, 1]], 2.0);
assert_eq!(h[[2, 2]], 3.0);
assert_eq!(h, h.t());
}
#[test]
fn test_add_by_value() {
let x = Jet2Vec::variable(3.0_f64, 0, 2);
let y = Jet2Vec::variable(4.0_f64, 1, 2);
let s = x + y;
assert_eq!(s.value(), 7.0);
assert_eq!(s.gradient(), &array![1.0, 1.0]);
assert_eq!(s.hessian(), &Array2::<f64>::zeros((2, 2)));
assert_eq!(s.hess, s.hess.t()); }
#[test]
fn test_sub_by_value() {
let x = Jet2Vec::variable(5.0_f64, 0, 2);
let y = Jet2Vec::variable(2.0_f64, 1, 2);
let d = x - y;
assert_eq!(d.value(), 3.0);
assert_eq!(d.gradient(), &array![1.0, -1.0]);
assert_eq!(d.hessian(), &Array2::<f64>::zeros((2, 2)));
assert_eq!(d.hess, d.hess.t());
}
#[test]
fn test_add_by_ref_matches_by_value() {
let x = Jet2Vec::variable(3.0_f64, 0, 2);
let y = Jet2Vec::variable(4.0_f64, 1, 2);
let s_ref = &x + &y;
let s_val = x.clone() + y.clone();
assert_eq!(s_ref, s_val);
}
#[test]
fn test_sub_by_ref_matches_by_value() {
let x = Jet2Vec::variable(5.0_f64, 0, 2);
let y = Jet2Vec::variable(2.0_f64, 1, 2);
let d_ref = &x - &y;
let d_val = x.clone() - y.clone();
assert_eq!(d_ref, d_val);
}
#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "gradient length mismatch")]
fn test_add_len_mismatch_panics() {
let a = Jet2Vec::variable(1.0_f64, 0, 2);
let b = Jet2Vec::variable(1.0_f64, 0, 3);
let _ = &a + &b;
}
#[test]
fn test_mul_cross_term_symmetric() {
let x = Jet2Vec::variable(3.0_f64, 0, 2);
let y = Jet2Vec::variable(4.0_f64, 1, 2);
let f = &x * &y; assert_eq!(f.value(), 12.0);
assert_eq!(f.gradient()[0], 4.0); assert_eq!(f.gradient()[1], 3.0); assert_eq!(f.hessian()[[0, 0]], 0.0);
assert_eq!(f.hessian()[[1, 1]], 0.0);
assert_eq!(f.hessian()[[0, 1]], 1.0); assert_eq!(f.hessian()[[1, 0]], 1.0); }
#[test]
fn test_mul_square_diagonal() {
let x = Jet2Vec::variable(5.0_f64, 0, 2);
let f = &x * &x;
assert_eq!(f.value(), 25.0);
assert_eq!(f.gradient(), &array![10.0, 0.0]);
assert_eq!(f.hess, f.hess.t());
assert_eq!(f.hessian()[[0, 0]], 2.0);
assert_eq!(f.hessian()[[1, 1]], 0.0);
assert_eq!(f.hessian()[[0, 1]], 0.0);
}
#[test]
fn test_mul_symmetry_bit_exact_on_3_vars() {
let x = Jet2Vec::variable(1.0_f64, 0, 3);
let y = Jet2Vec::variable(2.0_f64, 1, 3);
let z = Jet2Vec::variable(3.0_f64, 2, 3);
let a = &x + &y;
let b = &y + &z;
let f = &a * &b;
assert_eq!(f.value(), 15.0);
assert_eq!(f.gradient(), &array![5.0, 8.0, 3.0]);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_div_direct_form_at_b_near_one() {
let x = Jet2Vec::variable(7.5_f64, 0, 2);
let y = Jet2Vec::variable(1.0_f64, 1, 2);
let f = &x / &y;
assert_eq!(f.value(), 7.5);
assert_eq!(f.gradient()[0], 1.0);
assert_eq!(f.gradient()[1], -7.5);
assert_eq!(f.hess, f.hess.t());
assert_eq!(f.hessian()[[0, 0]], 0.0);
assert_eq!(f.hessian()[[1, 1]], 15.0);
assert_eq!(f.hessian()[[0, 1]], -1.0);
assert_eq!(f.hessian()[[1, 0]], -1.0);
}
#[test]
fn test_div_by_constant() {
let x = Jet2Vec::variable(6.0_f64, 0, 1);
let two = Jet2Vec::constant(2.0_f64, 1);
let f = &x / &two;
assert_eq!(f.value(), 3.0);
assert_eq!(f.gradient()[0], 0.5);
assert_eq!(f.hess, f.hess.t());
assert_eq!(f.hessian()[[0, 0]], 0.0);
}
use approx::assert_abs_diff_eq;
use std::f64::consts::PI;
#[test]
fn test_elementary_sin_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).sin();
let u: f64 = 0.7;
assert_abs_diff_eq!(f.value(), u.sin(), epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], u.cos(), epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[1], u.cos(), epsilon = 1e-15);
let expected_hess = -u.sin();
assert_abs_diff_eq!(f.hessian()[[0, 0]], expected_hess, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[1, 1]], expected_hess, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 1]], expected_hess, epsilon = 1e-15);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_cos_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).cos();
let u: f64 = 0.7;
assert_abs_diff_eq!(f.value(), u.cos(), epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], -u.sin(), epsilon = 1e-15);
let expected_hess = -u.cos();
assert_abs_diff_eq!(f.hessian()[[0, 1]], expected_hess, epsilon = 1e-15);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_tan_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).tan();
let u: f64 = 0.7;
let t = u.tan();
assert_abs_diff_eq!(f.value(), t, epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], 1.0 + t * t, epsilon = 1e-14);
let expected_hess = 2.0 * t * (1.0 + t * t);
assert_abs_diff_eq!(f.hessian()[[0, 1]], expected_hess, epsilon = 1e-14);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_exp_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).exp();
let e = 0.7_f64.exp();
assert_abs_diff_eq!(f.value(), e, epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], e, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 0]], e, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[1, 1]], e, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 1]], e, epsilon = 1e-15);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_ln_at_xy() {
let x = Jet2Vec::variable(2.0_f64, 0, 2);
let y = Jet2Vec::variable(3.0_f64, 1, 2);
let f = (&x * &y).ln();
assert_abs_diff_eq!(f.value(), 6.0_f64.ln(), epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], 0.5, epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[1], 1.0 / 3.0, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 0]], -0.25, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[1, 1]], -1.0 / 9.0, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 1]], 0.0, epsilon = 1e-14);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_sqrt_at_xy() {
let x = Jet2Vec::variable(1.0_f64, 0, 2);
let y = Jet2Vec::variable(3.0_f64, 1, 2);
let f = (&x + &y).sqrt();
assert_abs_diff_eq!(f.value(), 2.0, epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], 0.25, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 0]], -1.0 / 32.0, epsilon = 1e-15);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_tanh_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).tanh();
let u: f64 = 0.7;
let t = u.tanh();
assert_abs_diff_eq!(f.value(), t, epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], 1.0 - t * t, epsilon = 1e-14);
let expected_hess = -2.0 * t * (1.0 - t * t);
assert_abs_diff_eq!(f.hessian()[[0, 1]], expected_hess, epsilon = 1e-14);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_atan_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).atan();
let u: f64 = 0.7;
assert_abs_diff_eq!(f.value(), u.atan(), epsilon = 1e-15);
let fp = 1.0 / (1.0 + u * u);
assert_abs_diff_eq!(f.gradient()[0], fp, epsilon = 1e-15);
let fpp = -2.0 * u / ((1.0 + u * u) * (1.0 + u * u));
assert_abs_diff_eq!(f.hessian()[[0, 1]], fpp, epsilon = 1e-14);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_asinh_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).asinh();
let u: f64 = 0.7;
assert_abs_diff_eq!(f.value(), u.asinh(), epsilon = 1e-15);
let fp = 1.0 / (1.0 + u * u).sqrt();
assert_abs_diff_eq!(f.gradient()[0], fp, epsilon = 1e-15);
let d = 1.0 + u * u;
let fpp = -u / (d * d.sqrt());
assert_abs_diff_eq!(f.hessian()[[0, 1]], fpp, epsilon = 1e-14);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_elementary_erf_at_xy() {
let x = Jet2Vec::variable(0.3_f64, 0, 2);
let y = Jet2Vec::variable(0.4_f64, 1, 2);
let f = (&x + &y).erf();
let u: f64 = 0.7;
let two_over_sqrt_pi = std::f64::consts::FRAC_2_SQRT_PI;
let expected_fp = two_over_sqrt_pi * (-u * u).exp();
let expected_fpp = -2.0 * u * expected_fp;
assert_abs_diff_eq!(f.value(), crate::math::erf::<f64>(u), epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[0], expected_fp, epsilon = 1e-15);
assert_abs_diff_eq!(f.gradient()[1], expected_fp, epsilon = 1e-15);
assert_abs_diff_eq!(f.hessian()[[0, 0]], expected_fpp, epsilon = 1e-14);
assert_abs_diff_eq!(f.hessian()[[1, 1]], expected_fpp, epsilon = 1e-14);
assert_abs_diff_eq!(f.hessian()[[0, 1]], expected_fpp, epsilon = 1e-14);
assert_eq!(f.hess, f.hess.t());
}
#[test]
fn test_powf_direct_form_at_xy() {
let x = Jet2Vec::variable(1.0_f64, 0, 2);
let y = Jet2Vec::variable(2.0_f64, 1, 2);
let f = (&x + &y).powf(3.0);
assert_abs_diff_eq!(f.value(), 27.0, epsilon = 1e-13);
assert_abs_diff_eq!(f.gradient()[0], 27.0, epsilon = 1e-13);
assert_abs_diff_eq!(f.gradient()[1], 27.0, epsilon = 1e-13);
assert_abs_diff_eq!(f.hessian()[[0, 0]], 18.0, epsilon = 1e-13);
assert_abs_diff_eq!(f.hessian()[[1, 1]], 18.0, epsilon = 1e-13);
assert_abs_diff_eq!(f.hessian()[[0, 1]], 18.0, epsilon = 1e-13);
assert_eq!(f.hess, f.hess.t());
}
#[allow(dead_code)]
fn _use_pi() -> f64 {
PI
}
}