use crate::*;
use nalgebra::{Const, DMatrix, DVector, Dyn, OVector, SVector, U1};
pub fn partial<G: Fn(X, &A) -> O, T: DualNum<F>, F, X, A: DualStruct<T, F>, O>(
g: G,
args: &A::Inner,
) -> impl Fn(X) -> O {
let args = A::from_inner(args);
move |x| g(x, &args)
}
pub fn partial2<
G: Fn(X, &A1, &A2) -> O,
T: DualNum<F>,
F,
X,
A1: DualStruct<T, F>,
A2: DualStruct<T, F>,
O,
>(
g: G,
args1: &A1::Inner,
args2: &A2::Inner,
) -> impl Fn(X) -> O {
let args1 = A1::from_inner(args1);
let args2 = A2::from_inner(args2);
move |x| g(x, &args1, &args2)
}
pub fn partial3<
G: Fn(X, &A1, &A2, &A3) -> O,
T: DualNum<F>,
F,
X,
A1: DualStruct<T, F>,
A2: DualStruct<T, F>,
A3: DualStruct<T, F>,
O,
>(
g: G,
args1: &A1::Inner,
args2: &A2::Inner,
args3: &A3::Inner,
) -> impl Fn(X) -> O {
let args1 = A1::from_inner(args1);
let args2 = A2::from_inner(args2);
let args3 = A3::from_inner(args3);
move |x| g(x, &args1, &args2, &args3)
}
pub fn zeroth_derivative<G, T: DualNum<F>, F: DualNumFloat, O: Mappable<Real<T, F>>>(
g: G,
x: T,
) -> O::Output<T>
where
G: Fn(Real<T, F>) -> O,
{
let x = Real::from_re(x);
g(x).map_dual(|r| r.re)
}
pub fn first_derivative<G, T: DualNum<F>, F: DualNumFloat, O: Mappable<Dual<T, F>>>(
g: G,
x: T,
) -> O::Output<(T, T)>
where
G: Fn(Dual<T, F>) -> O,
{
let x = Dual::from_re(x).derivative();
g(x).map_dual(|r| (r.re, r.eps))
}
pub fn gradient<G, T: DualNum<F>, F: DualNumFloat, D: Dim, O: Mappable<DualVec<T, F, D>>>(
g: G,
x: &OVector<T, D>,
) -> O::Output<(T, OVector<T, D>)>
where
G: Fn(OVector<DualVec<T, F, D>, D>) -> O,
DefaultAllocator: Allocator<D>,
{
let mut x = x.map(DualVec::from_re);
let (r, c) = x.shape_generic();
for (i, xi) in x.iter_mut().enumerate() {
xi.eps = Derivative::derivative_generic(r, c, i);
}
g(x).map_dual(|res| (res.re, res.eps.unwrap_generic(r, c)))
}
#[expect(clippy::type_complexity)]
pub fn jacobian<
G,
T: DualNum<F>,
F: DualNumFloat,
M: Dim,
N: Dim,
O: Mappable<OVector<DualVec<T, F, N>, M>>,
>(
g: G,
x: &OVector<T, N>,
) -> O::Output<(OVector<T, M>, OMatrix<T, M, N>)>
where
G: FnOnce(OVector<DualVec<T, F, N>, N>) -> O,
DefaultAllocator: Allocator<M> + Allocator<N> + Allocator<M, N> + Allocator<U1, N>,
{
let mut x = x.map(DualVec::from_re);
let (r, c) = x.shape_generic();
for (i, xi) in x.iter_mut().enumerate() {
xi.eps = Derivative::derivative_generic(r, c, i);
}
let res = g(x);
res.map_dual(|res| {
let eps = OMatrix::from_rows(
res.map(|res| res.eps.unwrap_generic(r, c).transpose())
.as_slice(),
);
(res.map(|r| r.re), eps)
})
}
pub fn second_derivative<G, T: DualNum<F>, F, O: Mappable<Dual2<T, F>>>(
g: G,
x: T,
) -> O::Output<(T, T, T)>
where
G: Fn(Dual2<T, F>) -> O,
{
let x = Dual2::from_re(x).derivative();
g(x).map_dual(|r| (r.re, r.v1, r.v2))
}
pub fn second_partial_derivative<G, T: DualNum<F>, F, O: Mappable<HyperDual<T, F>>>(
g: G,
(x, y): (T, T),
) -> O::Output<(T, T, T, T)>
where
G: Fn((HyperDual<T, F>, HyperDual<T, F>)) -> O,
{
let x = HyperDual::from_re(x).derivative1();
let y = HyperDual::from_re(y).derivative2();
g((x, y)).map_dual(|r| (r.re, r.eps1, r.eps2, r.eps1eps2))
}
#[expect(clippy::type_complexity)]
pub fn hessian<G, T: DualNum<F>, F: DualNumFloat, D: Dim, O: Mappable<Dual2Vec<T, F, D>>>(
g: G,
x: &OVector<T, D>,
) -> O::Output<(T, OVector<T, D>, OMatrix<T, D, D>)>
where
G: Fn(OVector<Dual2Vec<T, F, D>, D>) -> O,
DefaultAllocator: Allocator<D> + Allocator<U1, D> + Allocator<D, D>,
{
let mut x = x.map(Dual2Vec::from_re);
let (r, c) = x.shape_generic();
for (i, xi) in x.iter_mut().enumerate() {
xi.v1 = Derivative::derivative_generic(c, r, i)
}
g(x).map_dual(|res| {
(
res.re,
res.v1.unwrap_generic(c, r).transpose(),
res.v2.unwrap_generic(r, r),
)
})
}
#[expect(clippy::type_complexity)]
pub fn partial_hessian<
G,
T: DualNum<F>,
F: DualNumFloat,
M: Dim,
N: Dim,
O: Mappable<HyperDualVec<T, F, M, N>>,
>(
g: G,
(x, y): (&OVector<T, M>, &OVector<T, N>),
) -> O::Output<(T, OVector<T, M>, OVector<T, N>, OMatrix<T, M, N>)>
where
G: Fn(
(
OVector<HyperDualVec<T, F, M, N>, M>,
OVector<HyperDualVec<T, F, M, N>, N>,
),
) -> O,
DefaultAllocator: Allocator<N> + Allocator<M> + Allocator<M, N> + Allocator<U1, N>,
{
let mut x = x.map(HyperDualVec::from_re);
let mut y = y.map(HyperDualVec::from_re);
let (m, _) = x.shape_generic();
for (i, xi) in x.iter_mut().enumerate() {
xi.eps1 = Derivative::derivative_generic(m, U1, i)
}
let (n, _) = y.shape_generic();
for (i, yi) in y.iter_mut().enumerate() {
yi.eps2 = Derivative::derivative_generic(U1, n, i)
}
g((x, y)).map_dual(|r| {
(
r.re,
r.eps1.unwrap_generic(m, U1),
r.eps2.unwrap_generic(U1, n).transpose(),
r.eps1eps2.unwrap_generic(m, n),
)
})
}
pub fn third_derivative<G, T: DualNum<F>, F, O: Mappable<Dual3<T, F>>>(
g: G,
x: T,
) -> O::Output<(T, T, T, T)>
where
G: Fn(Dual3<T, F>) -> O,
{
let x = Dual3::from_re(x).derivative();
g(x).map_dual(|r| (r.re, r.v1, r.v2, r.v3))
}
#[expect(clippy::type_complexity)]
pub fn third_partial_derivative<G, T: DualNum<F>, F, O: Mappable<HyperHyperDual<T, F>>>(
g: G,
(x, y, z): (T, T, T),
) -> O::Output<(T, T, T, T, T, T, T, T)>
where
G: Fn(
(
HyperHyperDual<T, F>,
HyperHyperDual<T, F>,
HyperHyperDual<T, F>,
),
) -> O,
{
let x = HyperHyperDual::from_re(x).derivative1();
let y = HyperHyperDual::from_re(y).derivative2();
let z = HyperHyperDual::from_re(z).derivative3();
g((x, y, z)).map_dual(|r| {
(
r.re,
r.eps1,
r.eps2,
r.eps3,
r.eps1eps2,
r.eps1eps3,
r.eps2eps3,
r.eps1eps2eps3,
)
})
}
#[expect(clippy::type_complexity)]
pub fn third_partial_derivative_vec<G, T: DualNum<F>, F, O: Mappable<HyperHyperDual<T, F>>>(
g: G,
x: &[T],
i: usize,
j: usize,
k: usize,
) -> O::Output<(T, T, T, T, T, T, T, T)>
where
G: Fn(&[HyperHyperDual<T, F>]) -> O,
{
let mut x: Vec<_> = x
.iter()
.map(|x| HyperHyperDual::from_re(x.clone()))
.collect();
x[i].eps1 = T::one();
x[j].eps2 = T::one();
x[k].eps3 = T::one();
g(&x).map_dual(|r| {
(
r.re,
r.eps1,
r.eps2,
r.eps3,
r.eps1eps2,
r.eps1eps3,
r.eps2eps3,
r.eps1eps2eps3,
)
})
}
pub trait Gradients: Dim
where
DefaultAllocator: Allocator<Self>,
{
type Dual<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, Inner = T> + Copy;
type Dual2<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, Inner = T> + Copy;
type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, Inner = T> + Copy;
fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
g: G,
x: &OVector<T, Self>,
args: &A::Inner,
) -> (T, OVector<T, Self>)
where
G: Fn(OVector<Self::Dual<T, F>, Self>, &A) -> Self::Dual<T, F>;
fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual2<T, F>, F>>(
g: G,
x: &OVector<T, Self>,
args: &A::Inner,
) -> (T, OVector<T, Self>, OMatrix<T, Self, Self>)
where
G: Fn(OVector<Self::Dual2<T, F>, Self>, &A) -> Self::Dual2<T, F>,
DefaultAllocator: Allocator<Self, Self>;
fn partial_hessian<
G,
T: DualNum<F> + Copy,
F: DualNumFloat,
A: DualStruct<Self::HyperDual<T, F>, F>,
>(
g: G,
x: &OVector<T, Self>,
y: T,
args: &A::Inner,
) -> (T, OVector<T, Self>, T, OVector<T, Self>)
where
G: Fn(
OVector<Self::HyperDual<T, F>, Self>,
Self::HyperDual<T, F>,
&A,
) -> Self::HyperDual<T, F>;
fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
g: G,
x: &OVector<T, Self>,
args: &A::Inner,
) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
where
G: Fn(OVector<Self::Dual<T, F>, Self>, &A) -> OVector<Self::Dual<T, F>, Self>,
DefaultAllocator: Allocator<Self, Self>;
}
impl<const N: usize> Gradients for Const<N> {
type Dual<T: DualNum<F> + Copy, F: DualNumFloat> = DualSVec<T, F, N>;
type Dual2<T: DualNum<F> + Copy, F: DualNumFloat> = Dual2Vec<T, F, Const<N>>;
type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDualVec<T, F, Const<N>, U1>;
fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<DualSVec<T, F, N>, F>>(
g: G,
x: &SVector<T, N>,
args: &A::Inner,
) -> (T, SVector<T, N>)
where
G: Fn(SVector<DualSVec<T, F, N>, N>, &A) -> DualSVec<T, F, N>,
{
gradient(partial(g, args), x)
}
fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual2<T, F>, F>>(
g: G,
x: &OVector<T, Self>,
args: &A::Inner,
) -> (T, OVector<T, Self>, OMatrix<T, Self, Self>)
where
G: Fn(OVector<Self::Dual2<T, F>, Self>, &A) -> Self::Dual2<T, F>,
{
hessian(partial(g, args), x)
}
fn partial_hessian<
G,
T: DualNum<F> + Copy,
F: DualNumFloat,
A: DualStruct<Self::HyperDual<T, F>, F>,
>(
g: G,
x: &OVector<T, Self>,
y: T,
args: &A::Inner,
) -> (T, OVector<T, Self>, T, OVector<T, Self>)
where
G: Fn(
OVector<Self::HyperDual<T, F>, Self>,
Self::HyperDual<T, F>,
&A,
) -> Self::HyperDual<T, F>,
{
let (a, b, c, d) = partial_hessian(
|(x, y)| {
let [[y]] = y.data.0;
g(x, y, &A::from_inner(args))
},
(x, &SVector::from([y])),
);
let [[c]] = c.data.0;
(a, b, c, d)
}
fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
g: G,
x: &OVector<T, Self>,
args: &A::Inner,
) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
where
G: Fn(OVector<DualVec<T, F, Self>, Self>, &A) -> OVector<DualVec<T, F, Self>, Self>,
{
jacobian(partial(g, args), x)
}
}
impl Gradients for Dyn {
type Dual<T: DualNum<F> + Copy, F: DualNumFloat> = Dual<T, F>;
type Dual2<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDual<T, F>;
type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDual<T, F>;
fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Dual<T, F>, F>>(
g: G,
x: &DVector<T>,
args: &A::Inner,
) -> (T, DVector<T>)
where
G: Fn(OVector<Dual<T, F>, Dyn>, &A) -> Dual<T, F>,
{
let mut re = T::zero();
let n = x.len();
let args = A::from_inner(args);
let grad = DVector::from_fn(n, |i, _| {
let mut x = x.map(Dual::from_re);
x[i].eps = T::one();
let res = g(x, &args);
re = res.re;
res.eps
});
(re, grad)
}
fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<HyperDual<T, F>, F>>(
g: G,
x: &DVector<T>,
args: &A::Inner,
) -> (T, DVector<T>, DMatrix<T>)
where
G: Fn(DVector<HyperDual<T, F>>, &A) -> HyperDual<T, F>,
{
let mut re = T::zero();
let n = x.len();
let args = A::from_inner(args);
let mut grad = DVector::zeros(n);
let hessian = DMatrix::from_fn(n, n, |i, j| {
let mut x = x.map(HyperDual::from_re);
x[i].eps1 = T::one();
x[j].eps2 = T::one();
let res = g(x, &args);
re = res.re;
grad[i] = res.eps1;
grad[j] = res.eps2;
res.eps1eps2
});
(re, grad, hessian)
}
fn partial_hessian<
G,
T: DualNum<F> + Copy,
F: DualNumFloat,
A: DualStruct<HyperDual<T, F>, F>,
>(
g: G,
x: &DVector<T>,
y: T,
args: &A::Inner,
) -> (T, DVector<T>, T, DVector<T>)
where
G: Fn(DVector<HyperDual<T, F>>, HyperDual<T, F>, &A) -> HyperDual<T, F>,
{
let mut re = T::zero();
let n = x.len();
let args = A::from_inner(args);
let y = HyperDual::from_re(y).derivative2();
let mut grad_x = DVector::zeros(n);
let mut grad_y = T::zero();
let hessian = DVector::from_fn(n, |i, _| {
let mut x = x.map(HyperDual::from_re);
x[i].eps1 = T::one();
let res = g(x, y, &args);
re = res.re;
grad_x[i] = res.eps1;
grad_y = res.eps2;
res.eps1eps2
});
(re, grad_x, grad_y, hessian)
}
fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
g: G,
x: &OVector<T, Self>,
args: &A::Inner,
) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
where
G: Fn(OVector<Dual<T, F>, Self>, &A) -> OVector<Dual<T, F>, Self>,
DefaultAllocator: Allocator<Self, Self>,
{
let n = x.len();
let args = A::from_inner(args);
let mut f = DVector::zeros(n);
let columns: Vec<_> = (0..n)
.map(|i| {
let mut x = x.map(Dual::from_re);
x[i].eps = T::one();
let res = g(x, &args);
f = res.map(|r| r.re);
res.map(|r| r.eps)
})
.collect();
let jac = DMatrix::from_columns(&columns);
(f, jac)
}
}