1use crate::*;
2use nalgebra::{Const, DMatrix, DVector, Dyn, OVector, SVector, U1};
3
4pub fn partial<G: Fn(X, &A) -> O, F, X, A: DualStruct<F>, O>(
7 g: G,
8 args: &A::Inner,
9) -> impl Fn(X) -> O {
10 let args = A::from_inner(args);
11 move |x| g(x, &args)
12}
13
14pub fn partial2<G: Fn(X, &A1, &A2) -> O, F, X, A1: DualStruct<F>, A2: DualStruct<F>, O>(
17 g: G,
18 args1: &A1::Inner,
19 args2: &A2::Inner,
20) -> impl Fn(X) -> O {
21 let args1 = A1::from_inner(args1);
22 let args2 = A2::from_inner(args2);
23 move |x| g(x, &args1, &args2)
24}
25
26pub fn partial3<
29 G: Fn(X, &A1, &A2, &A3) -> O,
30 F,
31 X,
32 A1: DualStruct<F>,
33 A2: DualStruct<F>,
34 A3: DualStruct<F>,
35 O,
36>(
37 g: G,
38 args1: &A1::Inner,
39 args2: &A2::Inner,
40 args3: &A3::Inner,
41) -> impl Fn(X) -> O {
42 let args1 = A1::from_inner(args1);
43 let args2 = A2::from_inner(args2);
44 let args3 = A3::from_inner(args3);
45 move |x| g(x, &args1, &args2, &args3)
46}
47
48pub fn zeroth_derivative<G, T: DualNum<F>, F: DualNumFloat, O: Mappable<Real<T, F>>>(
58 g: G,
59 x: T,
60) -> O::Output<T>
61where
62 G: Fn(Real<T, F>) -> O,
63{
64 let x = Real::from_re(x);
65 g(x).map_dual(|r| r.re)
66}
67
68pub fn first_derivative<G, T: DualNum<F>, F: DualNumFloat, O: Mappable<Dual<T, F>>>(
76 g: G,
77 x: T,
78) -> O::Output<(T, T)>
79where
80 G: Fn(Dual<T, F>) -> O,
81{
82 let x = Dual::from_re(x).derivative();
83 g(x).map_dual(|r| (r.re, r.eps))
84}
85
86pub fn gradient<G, T: DualNum<F>, F: DualNumFloat, D: Dim, O: Mappable<DualVec<T, F, D>>>(
114 g: G,
115 x: &OVector<T, D>,
116) -> O::Output<(T, OVector<T, D>)>
117where
118 G: Fn(OVector<DualVec<T, F, D>, D>) -> O,
119 DefaultAllocator: Allocator<D>,
120{
121 let mut x = x.map(DualVec::from_re);
122 let (r, c) = x.shape_generic();
123 for (i, xi) in x.iter_mut().enumerate() {
124 xi.eps = Derivative::derivative_generic(r, c, i);
125 }
126 g(x).map_dual(|res| (res.re, res.eps.unwrap_generic(r, c)))
127}
128
129#[expect(clippy::type_complexity)]
149pub fn jacobian<
150 G,
151 T: DualNum<F>,
152 F: DualNumFloat,
153 M: Dim,
154 N: Dim,
155 O: Mappable<OVector<DualVec<T, F, N>, M>>,
156>(
157 g: G,
158 x: &OVector<T, N>,
159) -> O::Output<(OVector<T, M>, OMatrix<T, M, N>)>
160where
161 G: FnOnce(OVector<DualVec<T, F, N>, N>) -> O,
162 DefaultAllocator: Allocator<M> + Allocator<N> + Allocator<M, N> + Allocator<U1, N>,
163{
164 let mut x = x.map(DualVec::from_re);
165 let (r, c) = x.shape_generic();
166 for (i, xi) in x.iter_mut().enumerate() {
167 xi.eps = Derivative::derivative_generic(r, c, i);
168 }
169 let res = g(x);
170 res.map_dual(|res| {
171 let eps = OMatrix::from_rows(
172 res.map(|res| res.eps.unwrap_generic(r, c).transpose())
173 .as_slice(),
174 );
175 (res.map(|r| r.re), eps)
176 })
177}
178
179pub fn second_derivative<G, T: DualNum<F>, F, O: Mappable<Dual2<T, F>>>(
201 g: G,
202 x: T,
203) -> O::Output<(T, T, T)>
204where
205 G: Fn(Dual2<T, F>) -> O,
206{
207 let x = Dual2::from_re(x).derivative();
208 g(x).map_dual(|r| (r.re, r.v1, r.v2))
209}
210
211pub fn second_partial_derivative<G, T: DualNum<F>, F, O: Mappable<HyperDual<T, F>>>(
224 g: G,
225 (x, y): (T, T),
226) -> O::Output<(T, T, T, T)>
227where
228 G: Fn((HyperDual<T, F>, HyperDual<T, F>)) -> O,
229{
230 let x = HyperDual::from_re(x).derivative1();
231 let y = HyperDual::from_re(y).derivative2();
232 g((x, y)).map_dual(|r| (r.re, r.eps1, r.eps2, r.eps1eps2))
233}
234
235#[expect(clippy::type_complexity)]
252pub fn hessian<G, T: DualNum<F>, F: DualNumFloat, D: Dim, O: Mappable<Dual2Vec<T, F, D>>>(
253 g: G,
254 x: &OVector<T, D>,
255) -> O::Output<(T, OVector<T, D>, OMatrix<T, D, D>)>
256where
257 G: Fn(OVector<Dual2Vec<T, F, D>, D>) -> O,
258 DefaultAllocator: Allocator<D> + Allocator<U1, D> + Allocator<D, D>,
259{
260 let mut x = x.map(Dual2Vec::from_re);
261 let (r, c) = x.shape_generic();
262 for (i, xi) in x.iter_mut().enumerate() {
263 xi.v1 = Derivative::derivative_generic(c, r, i)
264 }
265 g(x).map_dual(|res| {
266 (
267 res.re,
268 res.v1.unwrap_generic(c, r).transpose(),
269 res.v2.unwrap_generic(r, r),
270 )
271 })
272}
273
274#[expect(clippy::type_complexity)]
292pub fn partial_hessian<
293 G,
294 T: DualNum<F>,
295 F: DualNumFloat,
296 M: Dim,
297 N: Dim,
298 O: Mappable<HyperDualVec<T, F, M, N>>,
299>(
300 g: G,
301 (x, y): (&OVector<T, M>, &OVector<T, N>),
302) -> O::Output<(T, OVector<T, M>, OVector<T, N>, OMatrix<T, M, N>)>
303where
304 G: Fn(
305 (
306 OVector<HyperDualVec<T, F, M, N>, M>,
307 OVector<HyperDualVec<T, F, M, N>, N>,
308 ),
309 ) -> O,
310 DefaultAllocator: Allocator<N> + Allocator<M> + Allocator<M, N> + Allocator<U1, N>,
311{
312 let mut x = x.map(HyperDualVec::from_re);
313 let mut y = y.map(HyperDualVec::from_re);
314 let (m, _) = x.shape_generic();
315 for (i, xi) in x.iter_mut().enumerate() {
316 xi.eps1 = Derivative::derivative_generic(m, U1, i)
317 }
318 let (n, _) = y.shape_generic();
319 for (i, yi) in y.iter_mut().enumerate() {
320 yi.eps2 = Derivative::derivative_generic(U1, n, i)
321 }
322 g((x, y)).map_dual(|r| {
323 (
324 r.re,
325 r.eps1.unwrap_generic(m, U1),
326 r.eps2.unwrap_generic(U1, n).transpose(),
327 r.eps1eps2.unwrap_generic(m, n),
328 )
329 })
330}
331
332pub fn third_derivative<G, T: DualNum<F>, F, O: Mappable<Dual3<T, F>>>(
342 g: G,
343 x: T,
344) -> O::Output<(T, T, T, T)>
345where
346 G: Fn(Dual3<T, F>) -> O,
347{
348 let x = Dual3::from_re(x).derivative();
349 g(x).map_dual(|r| (r.re, r.v1, r.v2, r.v3))
350}
351
352#[expect(clippy::type_complexity)]
370pub fn third_partial_derivative<G, T: DualNum<F>, F, O: Mappable<HyperHyperDual<T, F>>>(
371 g: G,
372 (x, y, z): (T, T, T),
373) -> O::Output<(T, T, T, T, T, T, T, T)>
374where
375 G: Fn(
376 (
377 HyperHyperDual<T, F>,
378 HyperHyperDual<T, F>,
379 HyperHyperDual<T, F>,
380 ),
381 ) -> O,
382{
383 let x = HyperHyperDual::from_re(x).derivative1();
384 let y = HyperHyperDual::from_re(y).derivative2();
385 let z = HyperHyperDual::from_re(z).derivative3();
386 g((x, y, z)).map_dual(|r| {
387 (
388 r.re,
389 r.eps1,
390 r.eps2,
391 r.eps3,
392 r.eps1eps2,
393 r.eps1eps3,
394 r.eps2eps3,
395 r.eps1eps2eps3,
396 )
397 })
398}
399
400#[expect(clippy::type_complexity)]
419pub fn third_partial_derivative_vec<G, T: DualNum<F>, F, O: Mappable<HyperHyperDual<T, F>>>(
420 g: G,
421 x: &[T],
422 i: usize,
423 j: usize,
424 k: usize,
425) -> O::Output<(T, T, T, T, T, T, T, T)>
426where
427 G: Fn(&[HyperHyperDual<T, F>]) -> O,
428{
429 let mut x: Vec<_> = x
430 .iter()
431 .map(|x| HyperHyperDual::from_re(x.clone()))
432 .collect();
433 x[i].eps1 = T::one();
434 x[j].eps2 = T::one();
435 x[k].eps3 = T::one();
436 g(&x).map_dual(|r| {
437 (
438 r.re,
439 r.eps1,
440 r.eps2,
441 r.eps3,
442 r.eps1eps2,
443 r.eps1eps3,
444 r.eps2eps3,
445 r.eps1eps2eps3,
446 )
447 })
448}
449
450pub trait Gradients: Dim
453where
454 DefaultAllocator: Allocator<Self>,
455{
456 type Dual<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, InnerDual = T, Inner = T> + Copy;
457 type Dual2<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, InnerDual = T, Inner = T> + Copy;
458 type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, InnerDual = T, Inner = T>
459 + Copy;
460
461 fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
462 g: G,
463 x: &OVector<T, Self>,
464 args: &A::Inner,
465 ) -> (T, OVector<T, Self>)
466 where
467 G: Fn(OVector<Self::Dual<T, F>, Self>, &A) -> Self::Dual<T, F>;
468
469 fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
470 g: G,
471 x: &OVector<T, Self>,
472 args: &A::Inner,
473 ) -> (T, OVector<T, Self>, OMatrix<T, Self, Self>)
474 where
475 G: Fn(OVector<Self::Dual2<T, F>, Self>, &A) -> Self::Dual2<T, F>,
476 DefaultAllocator: Allocator<Self, Self>;
477
478 fn partial_hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
479 g: G,
480 x: &OVector<T, Self>,
481 y: T,
482 args: &A::Inner,
483 ) -> (T, OVector<T, Self>, T, OVector<T, Self>)
484 where
485 G: Fn(
486 OVector<Self::HyperDual<T, F>, Self>,
487 Self::HyperDual<T, F>,
488 &A,
489 ) -> Self::HyperDual<T, F>;
490
491 fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
492 g: G,
493 x: &OVector<T, Self>,
494 args: &A::Inner,
495 ) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
496 where
497 G: Fn(OVector<Self::Dual<T, F>, Self>, &A) -> OVector<Self::Dual<T, F>, Self>,
498 DefaultAllocator: Allocator<Self, Self>;
499}
500
501impl<const N: usize> Gradients for Const<N> {
502 type Dual<T: DualNum<F> + Copy, F: DualNumFloat> = DualSVec<T, F, N>;
503 type Dual2<T: DualNum<F> + Copy, F: DualNumFloat> = Dual2Vec<T, F, Const<N>>;
504 type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDualVec<T, F, Const<N>, U1>;
505
506 fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
507 g: G,
508 x: &SVector<T, N>,
509 args: &A::Inner,
510 ) -> (T, SVector<T, N>)
511 where
512 G: Fn(SVector<DualSVec<T, F, N>, N>, &A) -> DualSVec<T, F, N>,
513 {
514 gradient(partial(g, args), x)
515 }
516
517 fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
518 g: G,
519 x: &OVector<T, Self>,
520 args: &A::Inner,
521 ) -> (T, OVector<T, Self>, OMatrix<T, Self, Self>)
522 where
523 G: Fn(OVector<Self::Dual2<T, F>, Self>, &A) -> Self::Dual2<T, F>,
524 {
525 hessian(partial(g, args), x)
526 }
527
528 fn partial_hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
529 g: G,
530 x: &OVector<T, Self>,
531 y: T,
532 args: &A::Inner,
533 ) -> (T, OVector<T, Self>, T, OVector<T, Self>)
534 where
535 G: Fn(
536 OVector<Self::HyperDual<T, F>, Self>,
537 Self::HyperDual<T, F>,
538 &A,
539 ) -> Self::HyperDual<T, F>,
540 {
541 let (a, b, c, d) = partial_hessian(
542 |(x, y)| {
543 let [[y]] = y.data.0;
544 g(x, y, &A::from_inner(args))
545 },
546 (x, &SVector::from([y])),
547 );
548 let [[c]] = c.data.0;
549 (a, b, c, d)
550 }
551
552 fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
553 g: G,
554 x: &OVector<T, Self>,
555 args: &A::Inner,
556 ) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
557 where
558 G: Fn(OVector<DualVec<T, F, Self>, Self>, &A) -> OVector<DualVec<T, F, Self>, Self>,
559 {
560 jacobian(partial(g, args), x)
561 }
562}
563
564impl Gradients for Dyn {
565 type Dual<T: DualNum<F> + Copy, F: DualNumFloat> = Dual<T, F>;
566 type Dual2<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDual<T, F>;
567 type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDual<T, F>;
568
569 fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
570 g: G,
571 x: &DVector<T>,
572 args: &A::Inner,
573 ) -> (T, DVector<T>)
574 where
575 G: Fn(OVector<Dual<T, F>, Dyn>, &A) -> Dual<T, F>,
576 {
577 let mut re = T::zero();
578 let n = x.len();
579 let args = A::from_inner(args);
580 let grad = DVector::from_fn(n, |i, _| {
581 let mut x = x.map(Dual::from_re);
582 x[i].eps = T::one();
583 let res = g(x, &args);
584 re = res.re;
585 res.eps
586 });
587 (re, grad)
588 }
589
590 fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
591 g: G,
592 x: &DVector<T>,
593 args: &A::Inner,
594 ) -> (T, DVector<T>, DMatrix<T>)
595 where
596 G: Fn(DVector<HyperDual<T, F>>, &A) -> HyperDual<T, F>,
597 {
598 let mut re = T::zero();
599 let n = x.len();
600 let args = A::from_inner(args);
601 let mut grad = DVector::zeros(n);
602 let hessian = DMatrix::from_fn(n, n, |i, j| {
603 let mut x = x.map(HyperDual::from_re);
604 x[i].eps1 = T::one();
605 x[j].eps2 = T::one();
606 let res = g(x, &args);
607 re = res.re;
608 grad[i] = res.eps1;
609 grad[j] = res.eps2;
610 res.eps1eps2
611 });
612 (re, grad, hessian)
613 }
614
615 fn partial_hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
616 g: G,
617 x: &DVector<T>,
618 y: T,
619 args: &A::Inner,
620 ) -> (T, DVector<T>, T, DVector<T>)
621 where
622 G: Fn(DVector<HyperDual<T, F>>, HyperDual<T, F>, &A) -> HyperDual<T, F>,
623 {
624 let mut re = T::zero();
625 let n = x.len();
626 let args = A::from_inner(args);
627 let y = HyperDual::from_re(y).derivative2();
628 let mut grad_x = DVector::zeros(n);
629 let mut grad_y = T::zero();
630 let hessian = DVector::from_fn(n, |i, _| {
631 let mut x = x.map(HyperDual::from_re);
632 x[i].eps1 = T::one();
633 let res = g(x, y, &args);
634 re = res.re;
635 grad_x[i] = res.eps1;
636 grad_y = res.eps2;
637 res.eps1eps2
638 });
639 (re, grad_x, grad_y, hessian)
640 }
641
642 fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<F>>(
643 g: G,
644 x: &OVector<T, Self>,
645 args: &A::Inner,
646 ) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
647 where
648 G: Fn(OVector<Dual<T, F>, Self>, &A) -> OVector<Dual<T, F>, Self>,
649 DefaultAllocator: Allocator<Self, Self>,
650 {
651 let n = x.len();
652 let args = A::from_inner(args);
653 let mut f = DVector::zeros(n);
654 let columns: Vec<_> = (0..n)
655 .map(|i| {
656 let mut x = x.map(Dual::from_re);
657 x[i].eps = T::one();
658 let res = g(x, &args);
659 f = res.map(|r| r.re);
660 res.map(|r| r.eps)
661 })
662 .collect();
663 let jac = DMatrix::from_columns(&columns);
664 (f, jac)
665 }
666}