use crate::{DualNum, DualNumFloat, IsDerivativeZero, StaticMat, StaticVec};
use num_traits::{Float, FloatConst, FromPrimitive, Inv, Num, One, Signed, Zero};
use std::fmt;
use std::iter::{Product, Sum};
use std::marker::PhantomData;
use std::ops::{
Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, RemAssign, Sub, SubAssign,
};
#[derive(PartialEq, Eq, Copy, Clone, Debug)]
pub struct HyperDualVec<T, F, const M: usize, const N: usize> {
pub re: T,
pub eps1: StaticVec<T, M>,
pub eps2: StaticVec<T, N>,
pub eps1eps2: StaticMat<T, M, N>,
f: PhantomData<F>,
}
pub type HyperDualVec32<const M: usize, const N: usize> = HyperDualVec<f32, f32, M, N>;
pub type HyperDualVec64<const M: usize, const N: usize> = HyperDualVec<f64, f64, M, N>;
pub type HyperDual<T, F> = HyperDualVec<T, F, 1, 1>;
pub type HyperDual32 = HyperDual<f32, f32>;
pub type HyperDual64 = HyperDual<f64, f64>;
impl<T, F, const M: usize, const N: usize> HyperDualVec<T, F, M, N> {
#[inline]
pub fn new(
re: T,
eps1: StaticVec<T, M>,
eps2: StaticVec<T, N>,
eps1eps2: StaticMat<T, M, N>,
) -> Self {
Self {
re,
eps1,
eps2,
eps1eps2,
f: PhantomData,
}
}
}
impl<T, F> HyperDual<T, F> {
#[inline]
pub fn new_scalar(re: T, eps1: T, eps2: T, eps1eps2: T) -> Self {
Self::new(
re,
StaticVec::new_vec([eps1]),
StaticVec::new_vec([eps2]),
StaticMat::new([[eps1eps2]]),
)
}
}
impl<T: Copy + Zero + AddAssign, F, const M: usize, const N: usize> HyperDualVec<T, F, M, N> {
#[inline]
pub fn from_re(re: T) -> Self {
HyperDualVec::new(re, StaticVec::zero(), StaticVec::zero(), StaticMat::zero())
}
}
impl<T: One, F, const N: usize> HyperDualVec<T, F, 1, N> {
#[inline]
pub fn derive1(mut self) -> Self {
self.eps1[0] = T::one();
self
}
}
impl<T: One, F, const M: usize> HyperDualVec<T, F, M, 1> {
#[inline]
pub fn derive2(mut self) -> Self {
self.eps2[0] = T::one();
self
}
}
impl<T: One, F, const M: usize, const N: usize> StaticVec<HyperDualVec<T, F, M, N>, M> {
#[inline]
pub fn derive1(mut self) -> Self {
for i in 0..M {
self[i].eps1[i] = T::one();
}
self
}
}
impl<T: One, F, const M: usize, const N: usize> StaticVec<HyperDualVec<T, F, M, N>, N> {
#[inline]
pub fn derive2(mut self) -> Self {
for i in 0..N {
self[i].eps2[i] = T::one();
}
self
}
}
impl<T: DualNum<F>, F: Float, const M: usize, const N: usize> HyperDualVec<T, F, M, N> {
#[inline]
fn chain_rule(&self, f0: T, f1: T, f2: T) -> Self {
Self::new(
f0,
self.eps1 * f1,
self.eps2 * f1,
self.eps1eps2 * f1 + self.eps1.transpose_matmul(&self.eps2) * f2,
)
}
}
impl<'a, 'b, T: DualNum<F>, F: Float, const M: usize, const N: usize>
Mul<&'a HyperDualVec<T, F, M, N>> for &'b HyperDualVec<T, F, M, N>
{
type Output = HyperDualVec<T, F, M, N>;
#[inline]
fn mul(self, other: &HyperDualVec<T, F, M, N>) -> HyperDualVec<T, F, M, N> {
HyperDualVec::new(
self.re * other.re,
other.eps1 * self.re + self.eps1 * other.re,
other.eps2 * self.re + self.eps2 * other.re,
other.eps1eps2 * self.re
+ self.eps1.transpose_matmul(&other.eps2)
+ other.eps1.transpose_matmul(&self.eps2)
+ self.eps1eps2 * other.re,
)
}
}
impl<'a, 'b, T: DualNum<F>, F: Float, const M: usize, const N: usize>
Div<&'a HyperDualVec<T, F, M, N>> for &'b HyperDualVec<T, F, M, N>
{
type Output = HyperDualVec<T, F, M, N>;
#[inline]
fn div(self, other: &HyperDualVec<T, F, M, N>) -> HyperDualVec<T, F, M, N> {
let inv = other.re.recip();
let inv2 = inv * inv;
HyperDualVec::new(
self.re * inv,
(self.eps1 * other.re - other.eps1 * self.re) * inv2,
(self.eps2 * other.re - other.eps2 * self.re) * inv2,
self.eps1eps2 * inv
- (other.eps1eps2 * self.re
+ self.eps1.transpose_matmul(&other.eps2)
+ other.eps1.transpose_matmul(&self.eps2))
* inv2
+ other.eps1.transpose_matmul(&other.eps2)
* ((T::one() + T::one()) * self.re * inv2 * inv),
)
}
}
impl<T: fmt::Display, F: fmt::Display, const M: usize, const N: usize> fmt::Display
for HyperDualVec<T, F, M, N>
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"{} + {}ε1 + {}ε2 + {}ε1ε2",
self.re, self.eps1, self.eps2, self.eps1eps2
)
}
}
impl_second_derivatives!(HyperDualVec, [M, N], [eps1, eps2, eps1eps2]);
impl_dual!(HyperDualVec, [M, N], [eps1, eps2, eps1eps2]);