conspire/math/optimize/newton_raphson/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::{
5    super::{
6        Banded, Hessian, Jacobian, Matrix, Scalar, Solution, SquareMatrix, Tensor, TensorVec,
7        Vector,
8    },
9    EqualityConstraint, FirstOrderRootFinding, LineSearch, OptimizeError, SecondOrderOptimization,
10};
11use crate::ABS_TOL;
12use std::ops::{Div, Mul};
13
14/// The Newton-Raphson method.
15#[derive(Debug)]
16pub struct NewtonRaphson {
17    /// Absolute error tolerance.
18    pub abs_tol: Scalar,
19    /// Line search algorithm.
20    pub line_search: Option<LineSearch>,
21    /// Maximum number of steps.
22    pub max_steps: usize,
23}
24
25impl Default for NewtonRaphson {
26    fn default() -> Self {
27        Self {
28            abs_tol: ABS_TOL,
29            line_search: None,
30            max_steps: 25,
31        }
32    }
33}
34
35impl<F, J, X> FirstOrderRootFinding<F, J, X> for NewtonRaphson
36where
37    F: Jacobian,
38    for<'a> &'a F: Div<J, Output = X> + From<&'a X>,
39    J: Hessian,
40    X: Solution,
41    for<'a> &'a X: Mul<Scalar, Output = X>,
42    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
43{
44    fn root(
45        &self,
46        function: impl Fn(&X) -> Result<F, OptimizeError>,
47        jacobian: impl Fn(&X) -> Result<J, OptimizeError>,
48        initial_guess: X,
49        equality_constraint: EqualityConstraint,
50    ) -> Result<X, OptimizeError> {
51        match equality_constraint {
52            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => constrained(
53                self,
54                |_: &X| panic!("No line search in root finding"),
55                function,
56                jacobian,
57                initial_guess,
58                None,
59                constraint_matrix,
60                constraint_rhs,
61            ),
62            EqualityConstraint::None => unconstrained(
63                self,
64                |_: &X| panic!("No line search in root finding"),
65                function,
66                jacobian,
67                initial_guess,
68            ),
69        }
70    }
71}
72
73impl<J, H, X> SecondOrderOptimization<Scalar, J, H, X> for NewtonRaphson
74where
75    H: Hessian,
76    J: Jacobian,
77    for<'a> &'a J: Div<H, Output = X> + From<&'a X>,
78    X: Solution,
79    for<'a> &'a X: Mul<Scalar, Output = X>,
80    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
81{
82    fn minimize(
83        &self,
84        function: impl Fn(&X) -> Result<Scalar, OptimizeError>,
85        jacobian: impl Fn(&X) -> Result<J, OptimizeError>,
86        hessian: impl Fn(&X) -> Result<H, OptimizeError>,
87        initial_guess: X,
88        equality_constraint: EqualityConstraint,
89        banded: Option<Banded>,
90    ) -> Result<X, OptimizeError> {
91        match equality_constraint {
92            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => constrained(
93                self,
94                function,
95                jacobian,
96                hessian,
97                initial_guess,
98                banded,
99                constraint_matrix,
100                constraint_rhs,
101            ),
102            EqualityConstraint::None => {
103                unconstrained(self, function, jacobian, hessian, initial_guess)
104            }
105        }
106    }
107}
108
109fn unconstrained<J, H, X>(
110    newton_raphson: &NewtonRaphson,
111    function: impl Fn(&X) -> Result<Scalar, OptimizeError>,
112    jacobian: impl Fn(&X) -> Result<J, OptimizeError>,
113    hessian: impl Fn(&X) -> Result<H, OptimizeError>,
114    initial_guess: X,
115) -> Result<X, OptimizeError>
116where
117    H: Hessian,
118    J: Jacobian,
119    for<'a> &'a J: Div<H, Output = X> + From<&'a X>,
120    X: Solution,
121    for<'a> &'a X: Mul<Scalar, Output = X>,
122    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
123{
124    let mut decrement;
125    let mut residual;
126    let mut solution = initial_guess;
127    let mut tangent;
128    for _ in 0..newton_raphson.max_steps {
129        residual = jacobian(&solution)?;
130        tangent = hessian(&solution)?;
131        if residual.norm_inf() < newton_raphson.abs_tol {
132            return Ok(solution);
133        } else {
134            decrement = &residual / tangent;
135            if let Some(line_search) = &newton_raphson.line_search {
136                decrement *=
137                    line_search.backtrack(&function, &jacobian, &solution, &decrement, &1.0)?
138            }
139            solution -= decrement
140        }
141    }
142    Err(OptimizeError::MaximumStepsReached(
143        newton_raphson.max_steps,
144        format!("{:?}", &newton_raphson),
145    ))
146}
147
148#[allow(clippy::too_many_arguments)]
149fn constrained<J, H, X>(
150    newton_raphson: &NewtonRaphson,
151    _function: impl Fn(&X) -> Result<Scalar, OptimizeError>,
152    jacobian: impl Fn(&X) -> Result<J, OptimizeError>,
153    hessian: impl Fn(&X) -> Result<H, OptimizeError>,
154    initial_guess: X,
155    banded: Option<Banded>,
156    constraint_matrix: Matrix,
157    constraint_rhs: Vector,
158) -> Result<X, OptimizeError>
159where
160    H: Hessian,
161    J: Jacobian,
162    X: Solution,
163    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
164{
165    if newton_raphson.line_search.is_some() {
166        panic!("Line search needs the exact penalty function in constrained optimization.");
167    }
168    let mut decrement;
169    let num_variables = initial_guess.num_entries();
170    let num_constraints = constraint_rhs.len();
171    let num_total = num_variables + num_constraints;
172    let mut multipliers = Vector::zero(num_constraints);
173    let mut residual = Vector::zero(num_total);
174    let mut solution = initial_guess;
175    let mut tangent = SquareMatrix::zero(num_total);
176    constraint_matrix
177        .iter()
178        .enumerate()
179        .for_each(|(i, constraint_matrix_i)| {
180            constraint_matrix_i
181                .iter()
182                .enumerate()
183                .for_each(|(j, constraint_matrix_ij)| {
184                    tangent[i + num_variables][j] = -constraint_matrix_ij;
185                    tangent[j][i + num_variables] = -constraint_matrix_ij;
186                })
187        });
188    for _ in 0..newton_raphson.max_steps {
189        (jacobian(&solution)? - &multipliers * &constraint_matrix).fill_into_chained(
190            &constraint_rhs - &constraint_matrix * &solution,
191            &mut residual,
192        );
193        hessian(&solution)?.fill_into(&mut tangent);
194        if residual.norm_inf() < newton_raphson.abs_tol {
195            return Ok(solution);
196        } else if let Some(ref band) = banded {
197            decrement = tangent.solve_lu_banded(&residual, band)?
198        } else {
199            decrement = tangent.solve_lu(&residual)?
200        }
201        solution.decrement_from_chained(&mut multipliers, decrement)
202        // The convexity of every step of the solves can be verified (with LDL, LL, etc.).
203        // Also, consider revisiting null-space method to drastically reduce solve size.
204    }
205    Err(OptimizeError::MaximumStepsReached(
206        newton_raphson.max_steps,
207        format!("{:?}", &newton_raphson),
208    ))
209}
conspire/math/optimize/newton_raphson/mod.rs

conspire/math/optimize/newton_raphson/
mod.rs
