conspire/math/optimize/newton_raphson/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::{
5    super::{
6        Hessian, Jacobian, Matrix, Scalar, Solution, SquareMatrix, Tensor, Vector,
7        sparse::SparseSolver,
8    },
9    BacktrackingLineSearch, EqualityConstraint, FirstOrderRootFinding, LineSearch,
10    OptimizationError, SecondOrderOptimization,
11};
12use crate::ABS_TOL;
13use crate::math::Norm;
14use std::{
15    fmt::{self, Debug, Formatter},
16    ops::{Div, Mul},
17};
18
19/// The Newton-Raphson method.
20pub struct NewtonRaphson {
21    /// Absolute error tolerance.
22    pub abs_tol: Scalar,
23    /// Line search algorithm.
24    pub line_search: LineSearch,
25    /// Maximum number of steps.
26    pub max_steps: usize,
27    /// Norm type for error evaluation.
28    pub norm: Norm,
29}
30
31impl<J, X> BacktrackingLineSearch<J, X> for NewtonRaphson {
32    fn get_line_search(&self) -> &LineSearch {
33        &self.line_search
34    }
35}
36
37impl Debug for NewtonRaphson {
38    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
39        write!(
40            f,
41            "NewtonRaphson {{ abs_tol: {:?}, line_search: {}, max_steps: {:?} }}",
42            self.abs_tol, self.line_search, self.max_steps
43        )
44    }
45}
46
47impl Default for NewtonRaphson {
48    fn default() -> Self {
49        Self {
50            abs_tol: ABS_TOL,
51            line_search: LineSearch::None,
52            max_steps: 25,
53            norm: Norm::Chebyshev,
54        }
55    }
56}
57
58impl<F, J, X> FirstOrderRootFinding<F, J, X> for NewtonRaphson
59where
60    F: Jacobian,
61    for<'a> &'a F: Div<J, Output = X> + From<&'a X>,
62    J: Hessian,
63    X: Solution,
64    for<'a> &'a X: Mul<Scalar, Output = X>,
65    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
66{
67    fn root(
68        &self,
69        function: impl FnMut(&X) -> Result<F, String>,
70        jacobian: impl FnMut(&X) -> Result<J, String>,
71        initial_guess: X,
72        equality_constraint: EqualityConstraint,
73        sparse: Option<SparseSolver>,
74    ) -> Result<X, OptimizationError> {
75        match equality_constraint {
76            EqualityConstraint::Fixed(indices) => constrained_fixed(
77                self,
78                |_: &X| panic!("No line search in root finding"),
79                function,
80                jacobian,
81                initial_guess,
82                sparse,
83                indices,
84            ),
85            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => constrained(
86                self,
87                |_: &X| panic!("No line search in root finding"),
88                function,
89                jacobian,
90                initial_guess,
91                sparse,
92                constraint_matrix,
93                constraint_rhs,
94            ),
95            EqualityConstraint::None => unconstrained(
96                self,
97                |_: &X| panic!("No line search in root finding"),
98                function,
99                jacobian,
100                initial_guess,
101            ),
102        }
103    }
104}
105
106impl<J, H, X> SecondOrderOptimization<Scalar, J, H, X> for NewtonRaphson
107where
108    H: Hessian,
109    J: Jacobian,
110    for<'a> &'a J: Div<H, Output = X> + From<&'a X>,
111    X: Solution,
112    for<'a> &'a X: Mul<Scalar, Output = X>,
113    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
114{
115    fn minimize(
116        &self,
117        function: impl FnMut(&X) -> Result<Scalar, String>,
118        jacobian: impl FnMut(&X) -> Result<J, String>,
119        hessian: impl FnMut(&X) -> Result<H, String>,
120        initial_guess: X,
121        equality_constraint: EqualityConstraint,
122        sparse: Option<SparseSolver>,
123    ) -> Result<X, OptimizationError> {
124        match match equality_constraint {
125            EqualityConstraint::Fixed(indices) => constrained_fixed(
126                self,
127                function,
128                jacobian,
129                hessian,
130                initial_guess,
131                sparse,
132                indices,
133            ),
134            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => constrained(
135                self,
136                function,
137                jacobian,
138                hessian,
139                initial_guess,
140                sparse,
141                constraint_matrix,
142                constraint_rhs,
143            ),
144            EqualityConstraint::None => {
145                unconstrained(self, function, jacobian, hessian, initial_guess)
146            }
147        } {
148            Ok(solution) => Ok(solution),
149            Err(error) => Err(OptimizationError::Upstream(
150                format!("{error}"),
151                format!("{self:?}"),
152            )),
153        }
154    }
155}
156
157fn unconstrained<J, H, X>(
158    newton_raphson: &NewtonRaphson,
159    mut function: impl FnMut(&X) -> Result<Scalar, String>,
160    mut jacobian: impl FnMut(&X) -> Result<J, String>,
161    mut hessian: impl FnMut(&X) -> Result<H, String>,
162    initial_guess: X,
163) -> Result<X, OptimizationError>
164where
165    H: Hessian,
166    J: Jacobian,
167    for<'a> &'a J: Div<H, Output = X> + From<&'a X>,
168    X: Solution,
169    for<'a> &'a X: Mul<Scalar, Output = X>,
170{
171    let mut decrement;
172    let mut residual;
173    let mut solution = initial_guess;
174    let mut step_size;
175    let mut tangent;
176    for _ in 0..=newton_raphson.max_steps {
177        residual = jacobian(&solution)?;
178        if newton_raphson.norm.apply(&residual) < newton_raphson.abs_tol {
179            return Ok(solution);
180        } else {
181            tangent = hessian(&solution)?;
182            decrement = &residual / tangent;
183            step_size = newton_raphson.backtracking_line_search(
184                &mut function,
185                &mut jacobian,
186                &solution,
187                &residual,
188                &decrement,
189                1.0,
190            )?;
191            if step_size != 1.0 {
192                decrement *= step_size
193            }
194            solution -= decrement
195        }
196    }
197    Err(OptimizationError::MaximumStepsReached(
198        newton_raphson.max_steps,
199        format!("{:?}", newton_raphson),
200    ))
201}
202
203#[allow(clippy::too_many_arguments)]
204fn constrained_fixed<J, H, X>(
205    newton_raphson: &NewtonRaphson,
206    mut function: impl FnMut(&X) -> Result<Scalar, String>,
207    mut jacobian: impl FnMut(&X) -> Result<J, String>,
208    mut hessian: impl FnMut(&X) -> Result<H, String>,
209    initial_guess: X,
210    sparse: Option<SparseSolver>,
211    indices: Vec<usize>,
212) -> Result<X, OptimizationError>
213where
214    H: Hessian,
215    J: Jacobian,
216    for<'a> &'a J: From<&'a X>,
217    X: Solution,
218    for<'a> &'a X: Mul<Scalar, Output = X>,
219{
220    let mut retained = vec![true; initial_guess.size()];
221    indices.iter().for_each(|&index| retained[index] = false);
222    let unmap: Vec<usize> = retained
223        .iter()
224        .enumerate()
225        .filter_map(|(index, &keep)| keep.then_some(index))
226        .collect();
227    let mut decrement;
228    let mut residual;
229    let mut solution = initial_guess;
230    let mut step_size;
231    for _ in 0..=newton_raphson.max_steps {
232        residual = jacobian(&solution)?.retain_from(&retained);
233        if newton_raphson.norm.apply(&residual) < newton_raphson.abs_tol {
234            return Ok(solution);
235        } else if let Some(ref solver) = sparse {
236            let hess = hessian(&solution)?;
237            decrement = solver.solve(|i, j| hess.entry(unmap[i], unmap[j]), &residual)?
238        } else {
239            decrement = hessian(&solution)?
240                .retain_from(&retained)
241                .solve_lu(&residual)?
242        }
243        if !matches!(newton_raphson.line_search, LineSearch::None) {
244            let jac = jacobian(&solution)?;
245            let mut decrement_full = &solution * 0.0;
246            decrement_full.decrement_from_retained(&retained, &decrement);
247            decrement_full *= -1.0;
248            step_size = newton_raphson.backtracking_line_search(
249                &mut function,
250                &mut jacobian,
251                &solution,
252                &jac,
253                &decrement_full,
254                1.0,
255            )?;
256            if step_size != 1.0 {
257                decrement *= step_size
258            }
259        }
260        solution.decrement_from_retained(&retained, &decrement)
261    }
262    Err(OptimizationError::MaximumStepsReached(
263        newton_raphson.max_steps,
264        format!("{:?}", newton_raphson),
265    ))
266}
267
268#[allow(clippy::too_many_arguments)]
269fn constrained<J, H, X>(
270    newton_raphson: &NewtonRaphson,
271    _function: impl FnMut(&X) -> Result<Scalar, String>,
272    mut jacobian: impl FnMut(&X) -> Result<J, String>,
273    mut hessian: impl FnMut(&X) -> Result<H, String>,
274    initial_guess: X,
275    sparse: Option<SparseSolver>,
276    constraint_matrix: Matrix,
277    constraint_rhs: Vector,
278) -> Result<X, OptimizationError>
279where
280    H: Hessian,
281    J: Jacobian,
282    X: Solution,
283    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
284{
285    if !matches!(newton_raphson.line_search, LineSearch::None) {
286        panic!("Line search needs the exact penalty function in constrained optimization.")
287    }
288    let mut decrement;
289    let num_variables = initial_guess.size();
290    let num_constraints = constraint_rhs.len();
291    let num_total = num_variables + num_constraints;
292    let mut multipliers = Vector::zero(num_constraints);
293    let mut residual = Vector::zero(num_total);
294    let mut solution = initial_guess;
295    let mut tangent = SquareMatrix::zero(if sparse.is_none() { num_total } else { 0 });
296    if sparse.is_none() {
297        constraint_matrix
298            .iter()
299            .enumerate()
300            .for_each(|(i, constraint_matrix_i)| {
301                constraint_matrix_i
302                    .iter()
303                    .enumerate()
304                    .for_each(|(j, constraint_matrix_ij)| {
305                        tangent[i + num_variables][j] = -constraint_matrix_ij;
306                        tangent[j][i + num_variables] = -constraint_matrix_ij;
307                    })
308            });
309    }
310    for _ in 0..=newton_raphson.max_steps {
311        (jacobian(&solution)? - &multipliers * &constraint_matrix).fill_into_chained(
312            &constraint_rhs - &constraint_matrix * &solution,
313            &mut residual,
314        );
315        if newton_raphson.norm.apply(&residual) < newton_raphson.abs_tol {
316            return Ok(solution);
317        } else if let Some(ref solver) = sparse {
318            let hess = hessian(&solution)?;
319            decrement = solver.solve(
320                |i, j| {
321                    if i >= num_variables {
322                        -constraint_matrix[i - num_variables][j]
323                    } else if j >= num_variables {
324                        -constraint_matrix[j - num_variables][i]
325                    } else {
326                        hess.entry(i, j)
327                    }
328                },
329                &residual,
330            )?;
331        } else {
332            hessian(&solution)?.fill_into(&mut tangent);
333            decrement = tangent.solve_lu(&residual)?
334        }
335        solution.decrement_from_chained(&mut multipliers, decrement)
336    }
337    Err(OptimizationError::MaximumStepsReached(
338        newton_raphson.max_steps,
339        format!("{:?}", newton_raphson),
340    ))
341}
conspire/math/optimize/newton_raphson/mod.rs

conspire/math/optimize/newton_raphson/
mod.rs
