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
10pub 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#[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#[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#[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
161pub trait ImplicitFunction {
164 type Parameters<D>;
166
167 type Variable<D>;
169
170 fn residual<D: DualNum + Copy>(
172 x: Self::Variable<D>,
173 parameters: &Self::Parameters<D>,
174 ) -> Self::Variable<D>;
175}
176
177pub 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 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 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 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 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}