Skip to main content

num_dual/
implicit.rs

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