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
10pub 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
37pub 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
79pub 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
114pub 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
164pub trait ImplicitFunction<F> {
167 type Parameters<D>;
169
170 type Variable<D>;
172
173 fn residual<D: DualNum<F> + Copy>(
175 x: Self::Variable<D>,
176 parameters: &Self::Parameters<D>,
177 ) -> Self::Variable<D>;
178}
179
180pub 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 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 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 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 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}