num_dual/
implicit.rs

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

num_dual/
implicit.rs
