1use crate::*;
2use nalgebra::{Const, DMatrix, DVector, Dyn, OVector, SVector, U1};
3
4pub fn partial<G: Fn(X, &A) -> O, T: DualNum<F>, F, X, A: DualStruct<T, 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<
17 G: Fn(X, &A1, &A2) -> O,
18 T: DualNum<F>,
19 F,
20 X,
21 A1: DualStruct<T, F>,
22 A2: DualStruct<T, F>,
23 O,
24>(
25 g: G,
26 args1: &A1::Inner,
27 args2: &A2::Inner,
28) -> impl Fn(X) -> O {
29 let args1 = A1::from_inner(args1);
30 let args2 = A2::from_inner(args2);
31 move |x| g(x, &args1, &args2)
32}
33
34pub fn zeroth_derivative<G, T: DualNum<F>, F: DualNumFloat, O: Mappable<Real<T, F>>>(
44 g: G,
45 x: T,
46) -> O::Output<T>
47where
48 G: Fn(Real<T, F>) -> O,
49{
50 let x = Real::from_re(x);
51 g(x).map_dual(|r| r.re)
52}
53
54pub fn first_derivative<G, T: DualNum<F>, F: DualNumFloat, O: Mappable<Dual<T, F>>>(
62 g: G,
63 x: T,
64) -> O::Output<(T, T)>
65where
66 G: Fn(Dual<T, F>) -> O,
67{
68 let x = Dual::from_re(x).derivative();
69 g(x).map_dual(|r| (r.re, r.eps))
70}
71
72pub fn gradient<G, T: DualNum<F>, F: DualNumFloat, D: Dim, O: Mappable<DualVec<T, F, D>>>(
100 g: G,
101 x: &OVector<T, D>,
102) -> O::Output<(T, OVector<T, D>)>
103where
104 G: Fn(OVector<DualVec<T, F, D>, D>) -> O,
105 DefaultAllocator: Allocator<D>,
106{
107 let mut x = x.map(DualVec::from_re);
108 let (r, c) = x.shape_generic();
109 for (i, xi) in x.iter_mut().enumerate() {
110 xi.eps = Derivative::derivative_generic(r, c, i);
111 }
112 g(x).map_dual(|res| (res.re, res.eps.unwrap_generic(r, c)))
113}
114
115#[expect(clippy::type_complexity)]
135pub fn jacobian<
136 G,
137 T: DualNum<F>,
138 F: DualNumFloat,
139 M: Dim,
140 N: Dim,
141 O: Mappable<OVector<DualVec<T, F, N>, M>>,
142>(
143 g: G,
144 x: &OVector<T, N>,
145) -> O::Output<(OVector<T, M>, OMatrix<T, M, N>)>
146where
147 G: FnOnce(OVector<DualVec<T, F, N>, N>) -> O,
148 DefaultAllocator: Allocator<M> + Allocator<N> + Allocator<M, N> + Allocator<U1, N>,
149{
150 let mut x = x.map(DualVec::from_re);
151 let (r, c) = x.shape_generic();
152 for (i, xi) in x.iter_mut().enumerate() {
153 xi.eps = Derivative::derivative_generic(r, c, i);
154 }
155 let res = g(x);
156 res.map_dual(|res| {
157 let eps = OMatrix::from_rows(
158 res.map(|res| res.eps.unwrap_generic(r, c).transpose())
159 .as_slice(),
160 );
161 (res.map(|r| r.re), eps)
162 })
163}
164
165pub fn second_derivative<G, T: DualNum<F>, F, O: Mappable<Dual2<T, F>>>(
187 g: G,
188 x: T,
189) -> O::Output<(T, T, T)>
190where
191 G: Fn(Dual2<T, F>) -> O,
192{
193 let x = Dual2::from_re(x).derivative();
194 g(x).map_dual(|r| (r.re, r.v1, r.v2))
195}
196
197pub fn second_partial_derivative<G, T: DualNum<F>, F, O: Mappable<HyperDual<T, F>>>(
210 g: G,
211 (x, y): (T, T),
212) -> O::Output<(T, T, T, T)>
213where
214 G: Fn((HyperDual<T, F>, HyperDual<T, F>)) -> O,
215{
216 let x = HyperDual::from_re(x).derivative1();
217 let y = HyperDual::from_re(y).derivative2();
218 g((x, y)).map_dual(|r| (r.re, r.eps1, r.eps2, r.eps1eps2))
219}
220
221#[expect(clippy::type_complexity)]
238pub fn hessian<G, T: DualNum<F>, F: DualNumFloat, D: Dim, O: Mappable<Dual2Vec<T, F, D>>>(
239 g: G,
240 x: &OVector<T, D>,
241) -> O::Output<(T, OVector<T, D>, OMatrix<T, D, D>)>
242where
243 G: Fn(OVector<Dual2Vec<T, F, D>, D>) -> O,
244 DefaultAllocator: Allocator<D> + Allocator<U1, D> + Allocator<D, D>,
245{
246 let mut x = x.map(Dual2Vec::from_re);
247 let (r, c) = x.shape_generic();
248 for (i, xi) in x.iter_mut().enumerate() {
249 xi.v1 = Derivative::derivative_generic(c, r, i)
250 }
251 g(x).map_dual(|res| {
252 (
253 res.re,
254 res.v1.unwrap_generic(c, r).transpose(),
255 res.v2.unwrap_generic(r, r),
256 )
257 })
258}
259
260#[expect(clippy::type_complexity)]
278pub fn partial_hessian<
279 G,
280 T: DualNum<F>,
281 F: DualNumFloat,
282 M: Dim,
283 N: Dim,
284 O: Mappable<HyperDualVec<T, F, M, N>>,
285>(
286 g: G,
287 (x, y): (&OVector<T, M>, &OVector<T, N>),
288) -> O::Output<(T, OVector<T, M>, OVector<T, N>, OMatrix<T, M, N>)>
289where
290 G: Fn(
291 (
292 OVector<HyperDualVec<T, F, M, N>, M>,
293 OVector<HyperDualVec<T, F, M, N>, N>,
294 ),
295 ) -> O,
296 DefaultAllocator: Allocator<N> + Allocator<M> + Allocator<M, N> + Allocator<U1, N>,
297{
298 let mut x = x.map(HyperDualVec::from_re);
299 let mut y = y.map(HyperDualVec::from_re);
300 let (m, _) = x.shape_generic();
301 for (i, xi) in x.iter_mut().enumerate() {
302 xi.eps1 = Derivative::derivative_generic(m, U1, i)
303 }
304 let (n, _) = y.shape_generic();
305 for (i, yi) in y.iter_mut().enumerate() {
306 yi.eps2 = Derivative::derivative_generic(U1, n, i)
307 }
308 g((x, y)).map_dual(|r| {
309 (
310 r.re,
311 r.eps1.unwrap_generic(m, U1),
312 r.eps2.unwrap_generic(U1, n).transpose(),
313 r.eps1eps2.unwrap_generic(m, n),
314 )
315 })
316}
317
318pub fn third_derivative<G, T: DualNum<F>, F, O: Mappable<Dual3<T, F>>>(
328 g: G,
329 x: T,
330) -> O::Output<(T, T, T, T)>
331where
332 G: Fn(Dual3<T, F>) -> O,
333{
334 let x = Dual3::from_re(x).derivative();
335 g(x).map_dual(|r| (r.re, r.v1, r.v2, r.v3))
336}
337
338#[expect(clippy::type_complexity)]
356pub fn third_partial_derivative<G, T: DualNum<F>, F, O: Mappable<HyperHyperDual<T, F>>>(
357 g: G,
358 (x, y, z): (T, T, T),
359) -> O::Output<(T, T, T, T, T, T, T, T)>
360where
361 G: Fn(
362 (
363 HyperHyperDual<T, F>,
364 HyperHyperDual<T, F>,
365 HyperHyperDual<T, F>,
366 ),
367 ) -> O,
368{
369 let x = HyperHyperDual::from_re(x).derivative1();
370 let y = HyperHyperDual::from_re(y).derivative2();
371 let z = HyperHyperDual::from_re(z).derivative3();
372 g((x, y, z)).map_dual(|r| {
373 (
374 r.re,
375 r.eps1,
376 r.eps2,
377 r.eps3,
378 r.eps1eps2,
379 r.eps1eps3,
380 r.eps2eps3,
381 r.eps1eps2eps3,
382 )
383 })
384}
385
386#[expect(clippy::type_complexity)]
405pub fn third_partial_derivative_vec<G, T: DualNum<F>, F, O: Mappable<HyperHyperDual<T, F>>>(
406 g: G,
407 x: &[T],
408 i: usize,
409 j: usize,
410 k: usize,
411) -> O::Output<(T, T, T, T, T, T, T, T)>
412where
413 G: Fn(&[HyperHyperDual<T, F>]) -> O,
414{
415 let mut x: Vec<_> = x
416 .iter()
417 .map(|x| HyperHyperDual::from_re(x.clone()))
418 .collect();
419 x[i].eps1 = T::one();
420 x[j].eps2 = T::one();
421 x[k].eps3 = T::one();
422 g(&x).map_dual(|r| {
423 (
424 r.re,
425 r.eps1,
426 r.eps2,
427 r.eps3,
428 r.eps1eps2,
429 r.eps1eps3,
430 r.eps2eps3,
431 r.eps1eps2eps3,
432 )
433 })
434}
435
436pub trait Gradients: Dim
439where
440 DefaultAllocator: Allocator<Self>,
441{
442 type Dual<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, Inner = T> + Copy;
443 type Dual2<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, Inner = T> + Copy;
444 type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat>: DualNum<F, Inner = T> + Copy;
445
446 fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
447 g: G,
448 x: &OVector<T, Self>,
449 args: &A::Inner,
450 ) -> (T, OVector<T, Self>)
451 where
452 G: Fn(OVector<Self::Dual<T, F>, Self>, &A) -> Self::Dual<T, F>;
453
454 fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual2<T, F>, F>>(
455 g: G,
456 x: &OVector<T, Self>,
457 args: &A::Inner,
458 ) -> (T, OVector<T, Self>, OMatrix<T, Self, Self>)
459 where
460 G: Fn(OVector<Self::Dual2<T, F>, Self>, &A) -> Self::Dual2<T, F>,
461 DefaultAllocator: Allocator<Self, Self>;
462
463 fn partial_hessian<
464 G,
465 T: DualNum<F> + Copy,
466 F: DualNumFloat,
467 A: DualStruct<Self::HyperDual<T, F>, F>,
468 >(
469 g: G,
470 x: &OVector<T, Self>,
471 y: T,
472 args: &A::Inner,
473 ) -> (T, OVector<T, Self>, T, OVector<T, Self>)
474 where
475 G: Fn(
476 OVector<Self::HyperDual<T, F>, Self>,
477 Self::HyperDual<T, F>,
478 &A,
479 ) -> Self::HyperDual<T, F>;
480
481 fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
482 g: G,
483 x: &OVector<T, Self>,
484 args: &A::Inner,
485 ) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
486 where
487 G: Fn(OVector<Self::Dual<T, F>, Self>, &A) -> OVector<Self::Dual<T, F>, Self>,
488 DefaultAllocator: Allocator<Self, Self>;
489}
490
491impl<const N: usize> Gradients for Const<N> {
492 type Dual<T: DualNum<F> + Copy, F: DualNumFloat> = DualSVec<T, F, N>;
493 type Dual2<T: DualNum<F> + Copy, F: DualNumFloat> = Dual2Vec<T, F, Const<N>>;
494 type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDualVec<T, F, Const<N>, U1>;
495
496 fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<DualSVec<T, F, N>, F>>(
497 g: G,
498 x: &SVector<T, N>,
499 args: &A::Inner,
500 ) -> (T, SVector<T, N>)
501 where
502 G: Fn(SVector<DualSVec<T, F, N>, N>, &A) -> DualSVec<T, F, N>,
503 {
504 gradient(partial(g, args), x)
505 }
506
507 fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual2<T, F>, F>>(
508 g: G,
509 x: &OVector<T, Self>,
510 args: &A::Inner,
511 ) -> (T, OVector<T, Self>, OMatrix<T, Self, Self>)
512 where
513 G: Fn(OVector<Self::Dual2<T, F>, Self>, &A) -> Self::Dual2<T, F>,
514 {
515 hessian(partial(g, args), x)
516 }
517
518 fn partial_hessian<
519 G,
520 T: DualNum<F> + Copy,
521 F: DualNumFloat,
522 A: DualStruct<Self::HyperDual<T, F>, F>,
523 >(
524 g: G,
525 x: &OVector<T, Self>,
526 y: T,
527 args: &A::Inner,
528 ) -> (T, OVector<T, Self>, T, OVector<T, Self>)
529 where
530 G: Fn(
531 OVector<Self::HyperDual<T, F>, Self>,
532 Self::HyperDual<T, F>,
533 &A,
534 ) -> Self::HyperDual<T, F>,
535 {
536 let (a, b, c, d) = partial_hessian(
537 |(x, y)| {
538 let [[y]] = y.data.0;
539 g(x, y, &A::from_inner(args))
540 },
541 (x, &SVector::from([y])),
542 );
543 let [[c]] = c.data.0;
544 (a, b, c, d)
545 }
546
547 fn jacobian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Self::Dual<T, F>, F>>(
548 g: G,
549 x: &OVector<T, Self>,
550 args: &A::Inner,
551 ) -> (OVector<T, Self>, OMatrix<T, Self, Self>)
552 where
553 G: Fn(OVector<DualVec<T, F, Self>, Self>, &A) -> OVector<DualVec<T, F, Self>, Self>,
554 {
555 jacobian(partial(g, args), x)
556 }
557}
558
559impl Gradients for Dyn {
560 type Dual<T: DualNum<F> + Copy, F: DualNumFloat> = Dual<T, F>;
561 type Dual2<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDual<T, F>;
562 type HyperDual<T: DualNum<F> + Copy, F: DualNumFloat> = HyperDual<T, F>;
563
564 fn gradient<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<Dual<T, F>, F>>(
565 g: G,
566 x: &DVector<T>,
567 args: &A::Inner,
568 ) -> (T, DVector<T>)
569 where
570 G: Fn(OVector<Dual<T, F>, Dyn>, &A) -> Dual<T, F>,
571 {
572 let mut re = T::zero();
573 let n = x.len();
574 let args = A::from_inner(args);
575 let grad = DVector::from_fn(n, |i, _| {
576 let mut x = x.map(Dual::from_re);
577 x[i].eps = T::one();
578 let res = g(x, &args);
579 re = res.re;
580 res.eps
581 });
582 (re, grad)
583 }
584
585 fn hessian<G, T: DualNum<F> + Copy, F: DualNumFloat, A: DualStruct<HyperDual<T, F>, F>>(
586 g: G,
587 x: &DVector<T>,
588 args: &A::Inner,
589 ) -> (T, DVector<T>, DMatrix<T>)
590 where
591 G: Fn(DVector<HyperDual<T, F>>, &A) -> HyperDual<T, F>,
592 {
593 let mut re = T::zero();
594 let n = x.len();
595 let args = A::from_inner(args);
596 let mut grad = DVector::zeros(n);
597 let hessian = DMatrix::from_fn(n, n, |i, j| {
598 let mut x = x.map(HyperDual::from_re);
599 x[i].eps1 = T::one();
600 x[j].eps2 = T::one();
601 let res = g(x, &args);
602 re = res.re;
603 grad[i] = res.eps1;
604 grad[j] = res.eps2;
605 res.eps1eps2
606 });
607 (re, grad, hessian)
608 }
609
610 fn partial_hessian<
611 G,
612 T: DualNum<F> + Copy,
613 F: DualNumFloat,
614 A: DualStruct<HyperDual<T, F>, F>,
615 >(
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<Self::Dual<T, F>, 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}