1use crate::math::{
2 Erase, Hessian, HessianBlock, Is, Jacobian, LuDecomposition, Matrix, Scalar, Solution, Tensor,
3 Vector,
4 optimize::{
5 Direct, EqualityConstraint, EqualityConstraints, LinearSolver, ObjectiveMerit,
6 Optimization, OptimizationBlock, OptimizationError, OptimizationIncremental, ResidualMerit,
7 RootFinding, RootFindingBlock, RootFindingIncremental, Rows, SolveStrategy, StepSize,
8 no_objective, no_objective_blocks,
9 },
10 sparse::{CscMatrix, SparseSolver},
11};
12use crate::units::{UnitMul, UnitRatio, UnitSum};
13use std::{
14 fmt::Display,
15 ops::{Div, Mul},
16};
17
18use super::NewtonRaphson;
19use super::backtrack::{Backtrack, Objective, Residually};
20use super::blocked::blocked;
21use super::constrained::constrained;
22use super::fixed::constrained_fixed;
23use super::unconstrained::unconstrained;
24use super::verify::verified;
25
26impl<S, F, J, X, E> RootFinding<F, J, X> for NewtonRaphson<Direct, S>
27where
28 S: ResidualMerit,
29 F: Jacobian,
30 for<'a> &'a F: Div<J, Output = X>,
31 J: Hessian,
32 F: Erase<Erased = E>,
33 X: Erase<Erased = E> + Solution,
34 E: Tensor,
35 <X as Tensor>::Unit: UnitRatio,
36 for<'a> &'a X: Mul<StepSize<X, X>, Output = X> + Mul<Scalar, Output = X>,
37 for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
38 for<'a> &'a CscMatrix: Mul<&'a X, Output = Vector>,
39{
40 fn root(
41 &self,
42 function: impl FnMut(&X) -> Result<F, String>,
43 jacobian: impl FnMut(&X) -> Result<J, String>,
44 initial_guess: X,
45 equality_constraint: EqualityConstraints,
46 sparse: Option<SparseSolver>,
47 ) -> Result<X, OptimizationError> {
48 iterate(
49 self,
50 Residually(&self.line_search),
51 no_objective,
52 function,
53 jacobian,
54 no_update,
55 initial_guess,
56 equality_constraint,
57 sparse,
58 false,
59 )
60 }
61}
62
63impl<L, S, F, J, X, E> RootFinding<F, J, X> for NewtonRaphson<L, S>
64where
65 S: ResidualMerit,
66 L: LinearSolver<Tangent = J>,
67 F: Jacobian,
68 F: Erase<Erased = E>,
69 X: Erase<Erased = E> + Solution,
70 E: Tensor,
71 <X as Tensor>::Unit: UnitRatio,
72 for<'a> &'a X: Mul<StepSize<X, X>, Output = X>,
73{
74 fn root(
75 &self,
76 function: impl FnMut(&X) -> Result<F, String>,
77 jacobian: impl FnMut(&X) -> Result<J, String>,
78 initial_guess: X,
79 equality_constraint: EqualityConstraints,
80 _sparse: Option<SparseSolver>,
81 ) -> Result<X, OptimizationError> {
82 iterate_linear(
83 self,
84 Residually(&self.line_search),
85 no_objective,
86 function,
87 jacobian,
88 no_update,
89 initial_guess,
90 equality_constraint,
91 )
92 }
93}
94
95impl<S, F, J, X, E> RootFindingIncremental<F, J, X> for NewtonRaphson<Direct, S>
96where
97 S: ResidualMerit,
98 F: Jacobian,
99 for<'a> &'a F: Div<J, Output = X>,
100 J: Hessian,
101 F: Erase<Erased = E>,
102 X: Erase<Erased = E> + Solution,
103 E: Tensor,
104 <X as Tensor>::Unit: UnitRatio,
105 for<'a> &'a X: Mul<StepSize<X, X>, Output = X> + Mul<Scalar, Output = X>,
106 for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
107 for<'a> &'a CscMatrix: Mul<&'a X, Output = Vector>,
108{
109 fn root_incremental(
110 &self,
111 function: impl FnMut(&X) -> Result<F, String>,
112 jacobian: impl FnMut(&X) -> Result<J, String>,
113 update: impl FnMut(&X, &Vector, Scalar, bool) -> Result<(), String>,
114 initial_guess: X,
115 equality_constraint: EqualityConstraints,
116 sparse: Option<SparseSolver>,
117 ) -> Result<X, OptimizationError> {
118 iterate(
119 self,
120 Residually(&self.line_search),
121 no_objective,
122 function,
123 jacobian,
124 update,
125 initial_guess,
126 equality_constraint,
127 sparse,
128 true,
129 )
130 }
131}
132
133impl<S, F, J, H, X, E> Optimization<F, J, H, X> for NewtonRaphson<Direct, S>
134where
135 S: ObjectiveMerit,
136 F: Erase<Erased = Scalar> + Tensor,
137 <J as Tensor>::Unit: UnitMul<<X as Tensor>::Unit>,
138 <<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output: UnitSum,
139 <<<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output as UnitSum>::Output:
140 Is<<F as Tensor>::Unit>,
141 H: Hessian,
142 J: Jacobian,
143 for<'a> &'a J: Div<H, Output = X>,
144 J: Erase<Erased = E>,
145 X: Erase<Erased = E> + Solution,
146 E: Tensor,
147 <X as Tensor>::Unit: UnitRatio,
148 for<'a> &'a X: Mul<StepSize<X, X>, Output = X> + Mul<Scalar, Output = X>,
149 for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
150 for<'a> &'a CscMatrix: Mul<&'a X, Output = Vector>,
151{
152 fn minimize(
153 &self,
154 mut function: impl FnMut(&X) -> Result<F, String>,
155 jacobian: impl FnMut(&X) -> Result<J, String>,
156 hessian: impl FnMut(&X) -> Result<H, String>,
157 initial_guess: X,
158 equality_constraint: EqualityConstraints,
159 sparse: Option<SparseSolver>,
160 ) -> Result<X, OptimizationError> {
161 let function = move |argument: &X| function(argument).map(|value| *value.erase());
162 iterate(
163 self,
164 Objective(&self.line_search),
165 function,
166 jacobian,
167 hessian,
168 no_update,
169 initial_guess,
170 equality_constraint,
171 sparse,
172 false,
173 )
174 }
175}
176
177impl<L, S, F, J, H, X, E> Optimization<F, J, H, X> for NewtonRaphson<L, S>
178where
179 S: ObjectiveMerit,
180 L: LinearSolver<Tangent = H>,
181 F: Erase<Erased = Scalar> + Tensor,
182 <J as Tensor>::Unit: UnitMul<<X as Tensor>::Unit>,
183 <<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output: UnitSum,
184 <<<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output as UnitSum>::Output:
185 Is<<F as Tensor>::Unit>,
186 J: Jacobian,
187 J: Erase<Erased = E>,
188 X: Erase<Erased = E> + Solution,
189 E: Tensor,
190 <X as Tensor>::Unit: UnitRatio,
191 for<'a> &'a X: Mul<StepSize<X, X>, Output = X>,
192{
193 fn minimize(
194 &self,
195 mut function: impl FnMut(&X) -> Result<F, String>,
196 jacobian: impl FnMut(&X) -> Result<J, String>,
197 hessian: impl FnMut(&X) -> Result<H, String>,
198 initial_guess: X,
199 equality_constraint: EqualityConstraints,
200 _sparse: Option<SparseSolver>,
201 ) -> Result<X, OptimizationError> {
202 let function = move |argument: &X| function(argument).map(|value| *value.erase());
203 iterate_linear(
204 self,
205 Objective(&self.line_search),
206 function,
207 jacobian,
208 hessian,
209 no_update,
210 initial_guess,
211 equality_constraint,
212 )
213 }
214}
215
216impl<S, F, J, H, X, E> OptimizationIncremental<F, J, H, X> for NewtonRaphson<Direct, S>
217where
218 S: ObjectiveMerit,
219 F: Erase<Erased = Scalar> + Tensor,
220 <J as Tensor>::Unit: UnitMul<<X as Tensor>::Unit>,
221 <<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output: UnitSum,
222 <<<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output as UnitSum>::Output:
223 Is<<F as Tensor>::Unit>,
224 H: Hessian,
225 J: Jacobian,
226 for<'a> &'a J: Div<H, Output = X>,
227 J: Erase<Erased = E>,
228 X: Erase<Erased = E> + Solution,
229 E: Tensor,
230 <X as Tensor>::Unit: UnitRatio,
231 for<'a> &'a X: Mul<StepSize<X, X>, Output = X> + Mul<Scalar, Output = X>,
232 for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
233 for<'a> &'a CscMatrix: Mul<&'a X, Output = Vector>,
234{
235 fn minimize_incremental(
236 &self,
237 mut function: impl FnMut(&X) -> Result<F, String>,
238 jacobian: impl FnMut(&X) -> Result<J, String>,
239 hessian: impl FnMut(&X) -> Result<H, String>,
240 update: impl FnMut(&X, &Vector, Scalar, bool) -> Result<(), String>,
241 initial_guess: X,
242 equality_constraint: EqualityConstraints,
243 sparse: Option<SparseSolver>,
244 ) -> Result<X, OptimizationError> {
245 let function = move |argument: &X| function(argument).map(|value| *value.erase());
246 iterate(
247 self,
248 Objective(&self.line_search),
249 function,
250 jacobian,
251 hessian,
252 update,
253 initial_guess,
254 equality_constraint,
255 sparse,
256 true,
257 )
258 }
259}
260
261impl<S, U, V, Ru, Rv, Kuu, Kvu, Kuv, Kvv> RootFindingBlock<U, V, Ru, Rv, Kuu, Kvu, Kuv, Kvv>
262 for NewtonRaphson<Direct, S>
263where
264 S: ResidualMerit,
265 U: Solution,
266 V: Solution,
267 Ru: Jacobian,
268 Rv: Jacobian,
269 Kuu: HessianBlock,
270 Kvu: HessianBlock,
271 Kuv: HessianBlock,
272 Kvv: HessianBlock,
273 for<'a> &'a CscMatrix: Mul<&'a U, Output = Vector> + Mul<&'a V, Output = Vector>,
274{
275 fn root_block(
276 &self,
277 residual_global: impl FnMut(&U, &V) -> Result<Ru, String>,
278 residual_local: impl FnMut(&U, &V) -> Result<Rv, String>,
279 tangents: impl FnMut(&U, &V) -> Result<(Kuu, Kvu, Kuv, Kvv), String>,
280 initial_guess: (U, V),
281 constraint_global: (CscMatrix, Vector),
282 constraint_local: (CscMatrix, Vector),
283 sparse: Option<SparseSolver>,
284 strategy: SolveStrategy,
285 ) -> Result<(U, V), OptimizationError> {
286 blocked(
287 self,
288 Residually(&self.line_search),
289 no_objective_blocks,
290 residual_global,
291 residual_local,
292 tangents,
293 initial_guess,
294 constraint_global,
295 constraint_local,
296 sparse,
297 strategy,
298 )
299 .map_err(|error| OptimizationError::upstream(error, self))
300 }
301}
302
303impl<S, F, U, V, Ru, Rv, Kuu, Kvu, Kuv, Kvv> OptimizationBlock<F, U, V, Ru, Rv, Kuu, Kvu, Kuv, Kvv>
304 for NewtonRaphson<Direct, S>
305where
306 S: ObjectiveMerit,
307 F: Erase<Erased = Scalar> + Tensor,
308 <Ru as Tensor>::Unit: UnitMul<<U as Tensor>::Unit>,
309 <<Ru as Tensor>::Unit as UnitMul<<U as Tensor>::Unit>>::Output: UnitSum,
310 <<<Ru as Tensor>::Unit as UnitMul<<U as Tensor>::Unit>>::Output as UnitSum>::Output:
311 Is<<F as Tensor>::Unit>,
312 <Rv as Tensor>::Unit: UnitMul<<V as Tensor>::Unit>,
313 <<Rv as Tensor>::Unit as UnitMul<<V as Tensor>::Unit>>::Output: UnitSum,
314 <<<Rv as Tensor>::Unit as UnitMul<<V as Tensor>::Unit>>::Output as UnitSum>::Output:
315 Is<<F as Tensor>::Unit>,
316 U: Solution,
317 V: Solution,
318 Ru: Jacobian,
319 Rv: Jacobian,
320 Kuu: HessianBlock,
321 Kvu: HessianBlock,
322 Kuv: HessianBlock,
323 Kvv: HessianBlock,
324 for<'a> &'a CscMatrix: Mul<&'a U, Output = Vector> + Mul<&'a V, Output = Vector>,
325{
326 fn minimize_block(
327 &self,
328 mut function: impl FnMut(&U, &V) -> Result<F, String>,
329 residual_global: impl FnMut(&U, &V) -> Result<Ru, String>,
330 residual_local: impl FnMut(&U, &V) -> Result<Rv, String>,
331 tangents: impl FnMut(&U, &V) -> Result<(Kuu, Kvu, Kuv, Kvv), String>,
332 initial_guess: (U, V),
333 constraint_global: (CscMatrix, Vector),
334 constraint_local: (CscMatrix, Vector),
335 sparse: Option<SparseSolver>,
336 strategy: SolveStrategy,
337 ) -> Result<(U, V), OptimizationError> {
338 let function = move |global: &U, local: &V| function(global, local).map(|v| *v.erase());
339 blocked(
340 self,
341 Objective(&self.line_search),
342 function,
343 residual_global,
344 residual_local,
345 tangents,
346 initial_guess,
347 constraint_global,
348 constraint_local,
349 sparse,
350 strategy,
351 )
352 .map_err(|error| OptimizationError::upstream(error, self))
353 }
354}
355
356pub(super) fn no_update<X>(_: &X, _: &Vector, _: Scalar, _: bool) -> Result<(), String> {
357 Ok(())
358}
359
360#[expect(clippy::too_many_arguments)]
361pub(super) fn iterate<B, S, J, H, X, E>(
362 newton_raphson: &NewtonRaphson<Direct, S>,
363 search: B,
364 objective: impl FnMut(&X) -> Result<Scalar, String>,
365 jacobian: impl FnMut(&X) -> Result<J, String>,
366 hessian: impl FnMut(&X) -> Result<H, String>,
367 update: impl FnMut(&X, &Vector, Scalar, bool) -> Result<(), String>,
368 mut initial_guess: X,
369 equality_constraint: EqualityConstraints,
370 sparse: Option<SparseSolver>,
371 lent: bool,
372) -> Result<X, OptimizationError>
373where
374 B: Backtrack,
375 S: Display,
376 H: Hessian,
377 J: Jacobian,
378 for<'a> &'a J: Div<H, Output = X>,
379 J: Erase<Erased = E>,
380 X: Erase<Erased = E> + Solution,
381 E: Tensor,
382 <X as Tensor>::Unit: UnitRatio,
383 for<'a> &'a X: Mul<StepSize<X, X>, Output = X> + Mul<Scalar, Output = X>,
384 for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
385 for<'a> &'a CscMatrix: Mul<&'a X, Output = Vector>,
386{
387 match EqualityConstraint::split(equality_constraint, &mut initial_guess) {
388 (Some(indices), Rows::None) => constrained_fixed(
389 newton_raphson,
390 search,
391 objective,
392 jacobian,
393 hessian,
394 update,
395 initial_guess,
396 direct(sparse.clone()),
397 B::OBJECTIVE.then(|| verified(sparse)),
398 indices,
399 ),
400 (indices, Rows::Dense(constraint_matrix, constraint_rhs)) => constrained(
401 newton_raphson,
402 search,
403 objective,
404 jacobian,
405 hessian,
406 update,
407 initial_guess,
408 sparse,
409 &indices.unwrap_or_default(),
410 constraint_matrix,
411 constraint_rhs,
412 ),
413 (indices, Rows::Sparse(constraint_matrix, constraint_rhs)) => constrained(
414 newton_raphson,
415 search,
416 objective,
417 jacobian,
418 hessian,
419 update,
420 initial_guess,
421 sparse,
422 &indices.unwrap_or_default(),
423 constraint_matrix,
424 constraint_rhs,
425 ),
426 (None, Rows::None) => {
427 if lent {
428 unimplemented!(
429 "An unconstrained solution has no chained vector to lend the increment through."
430 )
431 }
432 unconstrained(
433 newton_raphson,
434 search,
435 objective,
436 jacobian,
437 hessian,
438 initial_guess,
439 sparse,
440 )
441 }
442 }
443 .map_err(|error| OptimizationError::upstream(error, newton_raphson))
444}
445
446#[expect(clippy::too_many_arguments)]
447pub(super) fn iterate_linear<B, L, S, J, H, X, E>(
448 newton_raphson: &NewtonRaphson<L, S>,
449 search: B,
450 objective: impl FnMut(&X) -> Result<Scalar, String>,
451 jacobian: impl FnMut(&X) -> Result<J, String>,
452 hessian: impl FnMut(&X) -> Result<H, String>,
453 update: impl FnMut(&X, &Vector, Scalar, bool) -> Result<(), String>,
454 mut initial_guess: X,
455 equality_constraint: EqualityConstraints,
456) -> Result<X, OptimizationError>
457where
458 B: Backtrack,
459 L: LinearSolver<Tangent = H>,
460 S: Display,
461 J: Jacobian,
462 J: Erase<Erased = E>,
463 X: Erase<Erased = E> + Solution,
464 E: Tensor,
465 <X as Tensor>::Unit: UnitRatio,
466 for<'a> &'a X: Mul<StepSize<X, X>, Output = X>,
467{
468 match EqualityConstraint::split(equality_constraint, &mut initial_guess) {
469 (Some(indices), Rows::None) => constrained_fixed(
470 newton_raphson,
471 search,
472 objective,
473 jacobian,
474 hessian,
475 update,
476 initial_guess,
477 |tangent, _, retained, residual, decrement| {
478 *decrement = newton_raphson
479 .linear_solver
480 .solve(tangent, retained, residual)?;
481 Ok(())
482 },
483 None::<fn(H, &[bool], &[usize]) -> Option<bool>>,
484 indices,
485 ),
486 _ => Err(OptimizationError::Intermediate(
487 "This linear solver requires a fixed equality constraint.".to_string(),
488 )),
489 }
490 .map_err(|error| OptimizationError::upstream(error, newton_raphson))
491}
492
493pub(super) fn direct<H>(
494 sparse: Option<SparseSolver>,
495) -> impl FnMut(H, &[bool], &[usize], &Vector, &mut Vector) -> Result<(), OptimizationError>
496where
497 H: Hessian,
498{
499 let mut factorization = None;
500 move |hessian, retained, unmap, residual, decrement| {
501 if let Some(ref solver) = sparse {
502 *decrement = solver.solve(|i, j| hessian.entry(unmap[i], unmap[j]), residual)?
503 } else {
504 let factorization =
505 factorization.get_or_insert_with(|| LuDecomposition::zero(unmap.len()));
506 hessian
507 .retain_from(retained)
508 .factorize_lu_into(factorization)?;
509 factorization.solve_into(residual, decrement)
510 }
511 Ok(())
512 }
513}