Skip to main content

num_dual/
implicit.rs

1#[cfg(feature = "nalgebra")]
2use crate::linalg::LU;
3use crate::{Dual, DualNum, DualNumFloat, DualStruct, first_derivative, partial};
4#[cfg(feature = "nalgebra")]
5use crate::{DualSVec, DualVec, Gradients, jacobian};
6#[cfg(feature = "nalgebra")]
7use nalgebra::{DefaultAllocator, Dim, OVector, SVector, U1, U2, allocator::Allocator};
8use std::marker::PhantomData;
9
10/// Calculate the derivative of the unary implicit function
11///         g(x, args) = 0
12/// ```
13/// # use num_dual::{implicit_derivative, DualNum, Dual2_64};
14/// # use approx::assert_relative_eq;
15/// let y = Dual2_64::from(25.0).derivative();
16/// let x = implicit_derivative(|x,y| x.powi(2)-y, 5.0f64, &y);
17/// assert_relative_eq!(x.re, y.sqrt().re, max_relative=1e-16);
18/// assert_relative_eq!(x.v1, y.sqrt().v1, max_relative=1e-16);
19/// assert_relative_eq!(x.v2, y.sqrt().v2, max_relative=1e-16);
20/// ```
21pub fn implicit_derivative<G, D: DualNum, A: DualStruct>(
22    g: G,
23    x: D::Primitive,
24    args: &A::Inner,
25) -> D
26where
27    G: Fn(Dual<D>, &A) -> Dual<D>,
28{
29    let mut x = D::from(x);
30    for _ in 0..D::NDERIV {
31        let (f, df) = first_derivative(partial(&g, args), x.clone());
32        x -= f / df;
33    }
34    x
35}
36
37/// Calculate the derivative of the binary implicit function
38///         g(x, y, args) = 0
39/// ```
40/// # use num_dual::{implicit_derivative_binary, Dual64};
41/// # use approx::assert_relative_eq;
42/// let a = Dual64::from(4.0).derivative();
43/// let [x, y] =
44///     implicit_derivative_binary(|x, y, a| [x * y - a, x + y - a - 1.0], 1.0f64, 4.0f64, &a);
45/// assert_relative_eq!(x.re, 1.0, max_relative = 1e-16);
46/// assert_relative_eq!(x.eps, 0.0, max_relative = 1e-16);
47/// assert_relative_eq!(y.re, a.re, max_relative = 1e-16);
48/// assert_relative_eq!(y.eps, a.eps, max_relative = 1e-16);
49/// ```
50#[cfg(feature = "nalgebra")]
51pub fn implicit_derivative_binary<G, D: DualNum, A: DualStruct>(
52    g: G,
53    x: D::Primitive,
54    y: D::Primitive,
55    args: &A::Inner,
56) -> [D; 2]
57where
58    G: Fn(DualVec<D, U2>, DualVec<D, U2>, &A) -> [DualVec<D, U2>; 2],
59{
60    let mut x = D::from(x);
61    let mut y = D::from(y);
62    let args = A::from_inner(args);
63    for _ in 0..D::NDERIV {
64        let (f, jac) = jacobian(
65            |x| {
66                let [[x, y]] = x.data.0;
67                SVector::from(g(x, y, &args))
68            },
69            &SVector::from([x.clone(), y.clone()]),
70        );
71        let [[f0, f1]] = f.data.0;
72        let [[j00, j10], [j01, j11]] = jac.data.0;
73        let det = (j00.clone() * &j11 - j01.clone() * &j10).recip();
74        x -= (j11 * &f0 - j01 * &f1) * &det;
75        y -= (j00 * &f1 - j10 * &f0) * &det;
76    }
77    [x, y]
78}
79
80/// Calculate the derivative of the multivariate implicit function
81///         g(x, args) = 0
82/// ```
83/// # use num_dual::{implicit_derivative_vec, Dual64};
84/// # use approx::assert_relative_eq;
85/// # use nalgebra::SVector;
86/// let a = Dual64::from(4.0).derivative();
87/// let x = implicit_derivative_vec(
88///     |x, a| SVector::from([x[0] * x[1] - a, x[0] + x[1] - a - 1.0]),
89///     SVector::from([1.0f64, 4.0f64]),
90///     &a,
91///     );
92/// assert_relative_eq!(x[0].re, 1.0, max_relative = 1e-16);
93/// assert_relative_eq!(x[0].eps, 0.0, max_relative = 1e-16);
94/// assert_relative_eq!(x[1].re, a.re, max_relative = 1e-16);
95/// assert_relative_eq!(x[1].eps, a.eps, max_relative = 1e-16);
96/// ```
97#[cfg(feature = "nalgebra")]
98pub fn implicit_derivative_vec<G, D: DualNum + Copy, A: DualStruct, N: Dim>(
99    g: G,
100    x: OVector<D::Primitive, N>,
101    args: &A::Inner,
102) -> OVector<D, N>
103where
104    DefaultAllocator: Allocator<N> + Allocator<N, N> + Allocator<U1, N>,
105    G: Fn(OVector<DualVec<D, N>, N>, &A) -> OVector<DualVec<D, N>, N>,
106{
107    let mut x = x.map(D::from);
108    let args = A::from_inner(args);
109    for _ in 0..D::NDERIV {
110        let (f, jac) = jacobian(|x| g(x, &args), &x);
111        x -= LU::new(jac).unwrap().solve(&f);
112    }
113    x
114}
115
116/// Calculate the derivative of stationary points of the scalar potential
117///         g(x, args)
118/// ```
119/// # use num_dual::{implicit_derivative_sp, Dual64, DualNum, Dual2Vec, HyperDual};
120/// # use approx::assert_relative_eq;
121/// # use nalgebra::{vector, dvector};
122/// let a = Dual64::from(2.0).derivative();
123/// let x = implicit_derivative_sp(
124///     |x, a: &Dual2Vec<_, _>| (a - x[0]).powi(2) + (x[1] - x[0]*x[0]).powi(2)*100.0,
125///     vector![2.0f64, 4.0f64],
126///     &a,
127///     );
128/// assert_relative_eq!(x[0].re, a.re, max_relative = 1e-13);
129/// assert_relative_eq!(x[0].eps, a.eps, max_relative = 1e-13);
130/// assert_relative_eq!(x[1].re, (a*a).re, max_relative = 1e-13);
131/// assert_relative_eq!(x[1].eps, (a*a).eps, max_relative = 1e-13);
132///
133/// let x = implicit_derivative_sp(
134///     |x, a: &HyperDual<_>| (a - x[0]).powi(2) + (x[1] - x[0]*x[0]).powi(2)*100.0,
135///     dvector![2.0f64, 4.0f64],
136///     &a,
137///     );
138/// assert_relative_eq!(x[0].re, a.re, max_relative = 1e-13);
139/// assert_relative_eq!(x[0].eps, a.eps, max_relative = 1e-13);
140/// assert_relative_eq!(x[1].re, (a*a).re, max_relative = 1e-13);
141/// assert_relative_eq!(x[1].eps, (a*a).eps, max_relative = 1e-13);
142/// ```
143#[cfg(feature = "nalgebra")]
144pub fn implicit_derivative_sp<G, D: DualNum + Copy, A: DualStruct, N: Gradients>(
145    g: G,
146    x: OVector<D::Primitive, N>,
147    args: &A::Inner,
148) -> OVector<D, N>
149where
150    DefaultAllocator: Allocator<N> + Allocator<N, N> + Allocator<U1, N>,
151    G: Fn(OVector<N::Dual2<D>, N>, &A) -> N::Dual2<D>,
152{
153    let mut x = x.map(D::from);
154    for _ in 0..D::NDERIV {
155        let (_, grad, hess) = N::hessian(|x, args| g(x, args), &x, args);
156        x -= LU::new(hess).unwrap().solve(&grad);
157    }
158    x
159}
160
161/// An implicit function g(x, args) = 0 for which derivatives of x can be
162/// calculated with the [ImplicitDerivative] struct.
163pub trait ImplicitFunction {
164    /// data type of the parameter struct, needs to implement [DualStruct<F>].
165    type Parameters<D>;
166
167    /// data type of the variable `x`, needs to be either `D`, `[D; 2]`, or `SVector<D, N>`.
168    type Variable<D>;
169
170    /// implementation of the residual function g(x, args) = 0.
171    fn residual<D: DualNum + Copy>(
172        x: Self::Variable<D>,
173        parameters: &Self::Parameters<D>,
174    ) -> Self::Variable<D>;
175}
176
177/// Helper struct that stores parameters in dual and real form and provides functions
178/// for evaluating real residuals (for external solvers) and implicit derivatives for
179/// arbitrary dual numbers.
180pub struct ImplicitDerivative<G: ImplicitFunction, D: DualNum + Copy, V> {
181    base: G::Parameters<D::Real>,
182    derivative: G::Parameters<D>,
183    phantom: PhantomData<V>,
184}
185
186impl<G: ImplicitFunction, D: DualNum + Copy> ImplicitDerivative<G, D, G::Variable<f64>>
187where
188    G::Parameters<D>: DualStruct<Real = G::Parameters<D::Primitive>>,
189{
190    pub fn new(_: G, parameters: G::Parameters<D>) -> Self {
191        Self {
192            base: parameters.re(),
193            derivative: parameters,
194            phantom: PhantomData,
195        }
196    }
197
198    /// Evaluate the (real) residual for a scalar function.
199    pub fn residual(&self, x: G::Variable<D::Primitive>) -> G::Variable<D::Primitive> {
200        G::residual(x, &self.base)
201    }
202}
203
204impl<G: ImplicitFunction, D: DualNum<Primitive = F> + Copy, F: DualNumFloat>
205    ImplicitDerivative<G, D, F>
206where
207    G::Parameters<D>: DualStruct<Real = G::Parameters<D::Primitive>>,
208{
209    /// Evaluate the implicit derivative for a scalar function.
210    pub fn implicit_derivative<A: DualStruct<Inner = G::Parameters<D>>>(&self, x: F) -> D
211    where
212        G: ImplicitFunction<Variable<Dual<D>> = Dual<D>, Parameters<Dual<D>> = A>,
213    {
214        implicit_derivative(G::residual::<Dual<D>>, x, &self.derivative)
215    }
216}
217
218#[cfg(feature = "nalgebra")]
219impl<G: ImplicitFunction, D: DualNum<Primitive = F> + Copy, F: DualNumFloat>
220    ImplicitDerivative<G, D, [F; 2]>
221where
222    G::Parameters<D>: DualStruct<Real = G::Parameters<D::Primitive>>,
223{
224    /// Evaluate the implicit derivative for a bivariate function.
225    pub fn implicit_derivative<A: DualStruct<Inner = G::Parameters<D>>>(&self, x: F, y: F) -> [D; 2]
226    where
227        G: ImplicitFunction<
228                Variable<DualVec<D, U2>> = [DualVec<D, U2>; 2],
229                Parameters<DualVec<D, U2>> = A,
230            >,
231    {
232        implicit_derivative_binary(
233            |x, y, args: &A| G::residual::<DualVec<D, U2>>([x, y], args),
234            x,
235            y,
236            &self.derivative,
237        )
238    }
239}
240
241#[cfg(feature = "nalgebra")]
242impl<G: ImplicitFunction, D: DualNum<Primitive = F> + Copy, F: DualNumFloat, const N: usize>
243    ImplicitDerivative<G, D, SVector<F, N>>
244where
245    G::Parameters<D>: DualStruct<Real = G::Parameters<D::Primitive>>,
246{
247    /// Evaluate the implicit derivative for a multivariate function.
248    pub fn implicit_derivative<A: DualStruct<Inner = G::Parameters<D>>>(
249        &self,
250        x: SVector<F, N>,
251    ) -> SVector<D, N>
252    where
253        G: ImplicitFunction<
254                Variable<DualSVec<D, N>> = SVector<DualSVec<D, N>, N>,
255                Parameters<DualSVec<D, N>> = A,
256            >,
257    {
258        implicit_derivative_vec(G::residual::<DualSVec<D, N>>, x, &self.derivative)
259    }
260}
261
262#[cfg(test)]
263mod test {
264    use super::*;
265
266    struct TestFunction;
267    impl ImplicitFunction for TestFunction {
268        type Parameters<D> = D;
269        type Variable<D> = D;
270
271        fn residual<D: DualNum + Copy>(x: D, square: &D) -> D {
272            *square - x * x
273        }
274    }
275
276    #[test]
277    fn test() {
278        let f: crate::Dual64 = Dual::from(25.0).derivative();
279        let func = ImplicitDerivative::new(TestFunction, f);
280        println!("{}", func.residual(5.0));
281        println!("{}", func.implicit_derivative(5.0));
282        println!("{}", f.sqrt());
283        assert_eq!(f.sqrt(), func.implicit_derivative(5.0));
284    }
285}
286
287#[cfg(test)]
288#[cfg(feature = "nalgebra")]
289mod test_nalgebra {
290    use super::*;
291    use nalgebra::SVector;
292
293    struct TestFunction2;
294    impl ImplicitFunction for TestFunction2 {
295        type Parameters<D> = (D, D);
296        type Variable<D> = [D; 2];
297
298        fn residual<D: DualNum + Copy>([x, y]: [D; 2], (square_sum, sum): &(D, D)) -> [D; 2] {
299            [*square_sum - x * x - y * y, *sum - x - y]
300        }
301    }
302
303    struct TestFunction3<const N: usize>;
304    impl<const N: usize> ImplicitFunction for TestFunction3<N> {
305        type Parameters<D> = D;
306        type Variable<D> = SVector<D, N>;
307
308        fn residual<D: DualNum + Copy>(x: SVector<D, N>, &square_sum: &D) -> SVector<D, N> {
309            let mut res = x;
310            for i in 1..N {
311                res[i] = x[i] - x[i - 1] - D::one();
312            }
313            res[0] = square_sum - x.dot(&x);
314            res
315        }
316    }
317
318    #[test]
319    fn test_nalgebra() {
320        let a: crate::Dual64 = Dual::from(25.0).derivative();
321        let b: crate::Dual64 = Dual::from(7.0);
322        let func = ImplicitDerivative::new(TestFunction2, (a, b));
323        println!("\n{:?}", func.residual([4.0, 3.0]));
324        let [x, y] = func.implicit_derivative(4.0, 3.0);
325        let xa = (b + (a * 2.0 - b * b).sqrt()) * 0.5;
326        let ya = (b - (a * 2.0 - b * b).sqrt()) * 0.5;
327        println!("{x}, {y}");
328        println!("{xa}, {ya}");
329        assert_eq!(x, xa);
330        assert_eq!(y, ya);
331
332        let s: crate::Dual64 = Dual::from(30.0).derivative();
333        let func = ImplicitDerivative::new(TestFunction3, s);
334        println!("\n{:?}", func.residual(SVector::from([1.0, 2.0, 3.0, 4.0])));
335        let x = func.implicit_derivative(SVector::from([1.0, 2.0, 3.0, 4.0]));
336        let x0 = ((s - 5.0).sqrt() - 5.0) * 0.5;
337        println!("{}, {}, {}, {}", x[0], x[1], x[2], x[3]);
338        println!("{}, {}, {}, {}", x0 + 1.0, x0 + 2.0, x0 + 3.0, x0 + 4.0);
339        assert_eq!(x0 + 1.0, x[0]);
340        assert_eq!(x0 + 2.0, x[1]);
341        assert_eq!(x0 + 3.0, x[2]);
342        assert_eq!(x0 + 4.0, x[3]);
343    }
344}