Skip to main content

conspire/math/optimize/nonlinear/newton_raphson/solve/
mod.rs

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}