1use crate::linalg::LU;
2use crate::{
3 Dual, Dual2Vec, DualNum, DualNumFloat, DualSVec, DualStruct, DualVec, first_derivative,
4 hessian, 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<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
37pub 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
84pub 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
125pub fn implicit_derivative_sp<
143 G,
144 D: DualNum<F> + Copy,
145 F: DualNumFloat,
146 A: DualStruct<Dual2Vec<D, F, N>, F>,
147 N: Dim,
148>(
149 g: G,
150 x: OVector<F, N>,
151 args: &A::Inner,
152) -> OVector<D, N>
153where
154 DefaultAllocator: Allocator<N> + Allocator<N, N> + Allocator<U1, N>,
155 G: Fn(OVector<Dual2Vec<D, F, N>, N>, &A) -> Dual2Vec<D, F, N>,
156{
157 let mut x = x.map(D::from);
158 let args = A::from_inner(args);
159 for _ in 0..D::NDERIV {
160 let (_, grad, hess) = hessian(|x| g(x, &args), &x);
161 x -= LU::new(hess).unwrap().solve(&grad);
162 }
163 x
164}
165
166pub trait ImplicitFunction<F> {
169 type Parameters<D>;
171
172 type Variable<D>;
174
175 fn residual<D: DualNum<F> + Copy>(
177 x: Self::Variable<D>,
178 parameters: &Self::Parameters<D>,
179 ) -> Self::Variable<D>;
180}
181
182pub struct ImplicitDerivative<G: ImplicitFunction<F>, D: DualNum<F> + Copy, F: DualNumFloat, V> {
186 base: G::Parameters<D::Real>,
187 derivative: G::Parameters<D>,
188 phantom: PhantomData<V>,
189}
190
191impl<G: ImplicitFunction<F>, D: DualNum<F> + Copy, F: DualNum<F> + DualNumFloat>
192 ImplicitDerivative<G, D, F, G::Variable<f64>>
193where
194 G::Parameters<D>: DualStruct<D, F, Real = G::Parameters<F>>,
195{
196 pub fn new(_: G, parameters: G::Parameters<D>) -> Self {
197 Self {
198 base: parameters.re(),
199 derivative: parameters,
200 phantom: PhantomData,
201 }
202 }
203
204 pub fn residual(&self, x: G::Variable<F>) -> G::Variable<F> {
206 G::residual(x, &self.base)
207 }
208}
209
210impl<G: ImplicitFunction<F>, D: DualNum<F> + Copy, F: DualNum<F> + DualNumFloat>
211 ImplicitDerivative<G, D, F, F>
212where
213 G::Parameters<D>: DualStruct<D, F, Real = G::Parameters<F>>,
214{
215 pub fn implicit_derivative<A: DualStruct<Dual<D, F>, F, Inner = G::Parameters<D>>>(
217 &self,
218 x: F,
219 ) -> D
220 where
221 G: ImplicitFunction<F, Variable<Dual<D, F>> = Dual<D, F>, Parameters<Dual<D, F>> = A>,
222 {
223 implicit_derivative(G::residual::<Dual<D, F>>, x, &self.derivative)
224 }
225}
226
227impl<G: ImplicitFunction<F>, D: DualNum<F> + Copy, F: DualNum<F> + DualNumFloat>
228 ImplicitDerivative<G, D, F, [F; 2]>
229where
230 G::Parameters<D>: DualStruct<D, F, Real = G::Parameters<F>>,
231{
232 pub fn implicit_derivative<A: DualStruct<DualVec<D, F, U2>, F, Inner = G::Parameters<D>>>(
234 &self,
235 x: F,
236 y: F,
237 ) -> [D; 2]
238 where
239 G: ImplicitFunction<
240 F,
241 Variable<DualVec<D, F, U2>> = [DualVec<D, F, U2>; 2],
242 Parameters<DualVec<D, F, U2>> = A,
243 >,
244 {
245 implicit_derivative_binary(
246 |x, y, args: &A| G::residual::<DualVec<D, F, U2>>([x, y], args),
247 x,
248 y,
249 &self.derivative,
250 )
251 }
252}
253
254impl<G: ImplicitFunction<F>, D: DualNum<F> + Copy, F: DualNum<F> + DualNumFloat, const N: usize>
255 ImplicitDerivative<G, D, F, SVector<F, N>>
256where
257 G::Parameters<D>: DualStruct<D, F, Real = G::Parameters<F>>,
258{
259 pub fn implicit_derivative<A: DualStruct<DualSVec<D, F, N>, F, Inner = G::Parameters<D>>>(
261 &self,
262 x: SVector<F, N>,
263 ) -> SVector<D, N>
264 where
265 G: ImplicitFunction<
266 F,
267 Variable<DualSVec<D, F, N>> = SVector<DualSVec<D, F, N>, N>,
268 Parameters<DualSVec<D, F, N>> = A,
269 >,
270 {
271 implicit_derivative_vec(G::residual::<DualSVec<D, F, N>>, x, &self.derivative)
272 }
273}
274
275#[cfg(test)]
276mod test {
277 use super::*;
278 use nalgebra::SVector;
279
280 struct TestFunction;
281 impl ImplicitFunction<f64> for TestFunction {
282 type Parameters<D> = D;
283 type Variable<D> = D;
284
285 fn residual<D: DualNum<f64> + Copy>(x: D, square: &D) -> D {
286 *square - x * x
287 }
288 }
289
290 struct TestFunction2;
291 impl ImplicitFunction<f64> for TestFunction2 {
292 type Parameters<D> = (D, D);
293 type Variable<D> = [D; 2];
294
295 fn residual<D: DualNum<f64> + Copy>([x, y]: [D; 2], (square_sum, sum): &(D, D)) -> [D; 2] {
296 [*square_sum - x * x - y * y, *sum - x - y]
297 }
298 }
299
300 struct TestFunction3<const N: usize>;
301 impl<const N: usize> ImplicitFunction<f64> for TestFunction3<N> {
302 type Parameters<D> = D;
303 type Variable<D> = SVector<D, N>;
304
305 fn residual<D: DualNum<f64> + Copy>(x: SVector<D, N>, &square_sum: &D) -> SVector<D, N> {
306 let mut res = x;
307 for i in 1..N {
308 res[i] = x[i] - x[i - 1] - D::from(1.0);
309 }
310 res[0] = square_sum - x.dot(&x);
311 res
312 }
313 }
314
315 #[test]
316 fn test() {
317 let f: crate::Dual64 = Dual::from(25.0).derivative();
318 let func = ImplicitDerivative::new(TestFunction, f);
319 println!("{}", func.residual(5.0));
320 println!("{}", func.implicit_derivative(5.0));
321 println!("{}", f.sqrt());
322 assert_eq!(f.sqrt(), func.implicit_derivative(5.0));
323
324 let a: crate::Dual64 = Dual::from(25.0).derivative();
325 let b: crate::Dual64 = Dual::from(7.0);
326 let func = ImplicitDerivative::new(TestFunction2, (a, b));
327 println!("\n{:?}", func.residual([4.0, 3.0]));
328 let [x, y] = func.implicit_derivative(4.0, 3.0);
329 let xa = (b + (a * 2.0 - b * b).sqrt()) * 0.5;
330 let ya = (b - (a * 2.0 - b * b).sqrt()) * 0.5;
331 println!("{x}, {y}");
332 println!("{xa}, {ya}");
333 assert_eq!(x, xa);
334 assert_eq!(y, ya);
335
336 let s: crate::Dual64 = Dual::from(30.0).derivative();
337 let func = ImplicitDerivative::new(TestFunction3, s);
338 println!("\n{:?}", func.residual(SVector::from([1.0, 2.0, 3.0, 4.0])));
339 let x = func.implicit_derivative(SVector::from([1.0, 2.0, 3.0, 4.0]));
340 let x0 = ((s - 5.0).sqrt() - 5.0) * 0.5;
341 println!("{}, {}, {}, {}", x[0], x[1], x[2], x[3]);
342 println!("{}, {}, {}, {}", x0 + 1.0, x0 + 2.0, x0 + 3.0, x0 + 4.0);
343 assert_eq!(x0 + 1.0, x[0]);
344 assert_eq!(x0 + 2.0, x[1]);
345 assert_eq!(x0 + 3.0, x[2]);
346 assert_eq!(x0 + 4.0, x[3]);
347 }
348}