conspire/math/optimize/gradient_descent/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::{
5    super::{Jacobian, Matrix, Scalar, Solution, Tensor, Vector},
6    BacktrackingLineSearch, EqualityConstraint, FirstOrderOptimization, LineSearch,
7    OptimizationError, ZerothOrderRootFinding,
8};
9use crate::ABS_TOL;
10use crate::math::Norm;
11use std::{
12    fmt::{self, Debug, Formatter},
13    ops::Mul,
14};
15
16const CUTBACK_FACTOR: Scalar = 0.8;
17const CUTBACK_FACTOR_MINUS_ONE: Scalar = 1.0 - CUTBACK_FACTOR;
18const INITIAL_STEP_SIZE: Scalar = 1e-2;
19
20/// The method of gradient descent.
21pub struct GradientDescent {
22    /// Absolute error tolerance.
23    pub abs_tol: Scalar,
24    /// Lagrangian dual.
25    pub dual: bool,
26    /// Line search algorithm.
27    pub line_search: LineSearch,
28    /// Maximum number of steps.
29    pub max_steps: usize,
30    /// Norm type for error evaluation.
31    pub norm: Norm,
32    /// Relative error tolerance.
33    pub rel_tol: Option<Scalar>,
34}
35
36impl<J, X> BacktrackingLineSearch<J, X> for GradientDescent {
37    fn get_line_search(&self) -> &LineSearch {
38        &self.line_search
39    }
40}
41
42impl Debug for GradientDescent {
43    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
44        write!(
45            f,
46            "GradientDescent {{ abs_tol: {:?}, dual: {:?}, line_search: {}, max_steps: {:?}, rel_tol: {:?} }}",
47            self.abs_tol, self.dual, self.line_search, self.max_steps, self.rel_tol
48        )
49    }
50}
51
52impl Default for GradientDescent {
53    fn default() -> Self {
54        Self {
55            abs_tol: ABS_TOL,
56            dual: false,
57            line_search: LineSearch::None,
58            max_steps: 250,
59            norm: Norm::Chebyshev,
60            rel_tol: None,
61        }
62    }
63}
64
65impl<X> ZerothOrderRootFinding<X> for GradientDescent
66where
67    X: Jacobian + Solution,
68    for<'a> &'a X: Mul<Scalar, Output = X>,
69    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
70{
71    fn root(
72        &self,
73        function: impl FnMut(&X) -> Result<X, String>,
74        initial_guess: X,
75        equality_constraint: EqualityConstraint,
76    ) -> Result<X, OptimizationError> {
77        match equality_constraint {
78            EqualityConstraint::Fixed(indices) => constrained_fixed(
79                self,
80                |_: &X| panic!("No line search in root finding."),
81                function,
82                initial_guess,
83                indices,
84            ),
85            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => {
86                if self.dual {
87                    constrained_dual(
88                        self,
89                        function,
90                        initial_guess,
91                        constraint_matrix,
92                        constraint_rhs,
93                    )
94                } else {
95                    constrained(
96                        self,
97                        function,
98                        initial_guess,
99                        constraint_matrix,
100                        constraint_rhs,
101                    )
102                }
103            }
104            EqualityConstraint::None => unconstrained(
105                self,
106                |_: &X| panic!("No line search in root finding."),
107                function,
108                initial_guess,
109                None,
110            ),
111        }
112    }
113}
114
115impl<X> FirstOrderOptimization<Scalar, X> for GradientDescent
116where
117    X: Jacobian + Solution,
118    for<'a> &'a X: Mul<Scalar, Output = X>,
119    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
120{
121    fn minimize(
122        &self,
123        function: impl FnMut(&X) -> Result<Scalar, String>,
124        jacobian: impl FnMut(&X) -> Result<X, String>,
125        initial_guess: X,
126        equality_constraint: EqualityConstraint,
127    ) -> Result<X, OptimizationError> {
128        match equality_constraint {
129            EqualityConstraint::Fixed(indices) => {
130                constrained_fixed(self, function, jacobian, initial_guess, indices)
131            }
132            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => {
133                if self.dual {
134                    constrained_dual(
135                        self,
136                        jacobian,
137                        initial_guess,
138                        constraint_matrix,
139                        constraint_rhs,
140                    )
141                } else {
142                    constrained(
143                        self,
144                        jacobian,
145                        initial_guess,
146                        constraint_matrix,
147                        constraint_rhs,
148                    )
149                }
150            }
151            EqualityConstraint::None => {
152                unconstrained(self, function, jacobian, initial_guess, None)
153            }
154        }
155    }
156}
157
158fn unconstrained<X>(
159    gradient_descent: &GradientDescent,
160    mut function: impl FnMut(&X) -> Result<Scalar, String>,
161    mut jacobian: impl FnMut(&X) -> Result<X, String>,
162    initial_guess: X,
163    linear_equality_constraint: Option<(&Matrix, &Vector)>,
164) -> Result<X, OptimizationError>
165where
166    X: Jacobian + Solution,
167    for<'a> &'a X: Mul<Scalar, Output = X>,
168{
169    let constraint = if let Some((constraint_matrix, multipliers)) = linear_equality_constraint {
170        Some(multipliers * constraint_matrix)
171    } else {
172        None
173    };
174    let mut residual;
175    let mut residual_change = initial_guess.clone() * 0.0;
176    let mut solution = initial_guess.clone();
177    let mut solution_change = solution.clone();
178    let mut step_size = INITIAL_STEP_SIZE;
179    let mut step_trial;
180    for _ in 0..=gradient_descent.max_steps {
181        residual = if let Some(ref extra) = constraint {
182            jacobian(&solution)? - extra
183        } else {
184            jacobian(&solution)?
185        };
186        if gradient_descent.norm.apply(&residual) < gradient_descent.abs_tol {
187            return Ok(solution);
188        } else {
189            solution_change -= &solution;
190            residual_change -= &residual;
191            step_trial =
192                residual_change.full_contraction(&solution_change) / residual_change.norm_squared();
193            if step_trial.abs() > 0.0 && !step_trial.is_nan() {
194                step_size = step_trial.abs()
195            }
196            step_size = gradient_descent.backtracking_line_search(
197                &mut function,
198                &mut jacobian,
199                &solution,
200                &residual,
201                &residual,
202                step_size,
203            )?;
204            residual_change = residual.clone();
205            solution_change = solution.clone();
206            solution -= residual * step_size;
207        }
208    }
209    Err(OptimizationError::MaximumStepsReached(
210        gradient_descent.max_steps,
211        format!("{gradient_descent:?}"),
212    ))
213}
214
215fn constrained_fixed<X>(
216    gradient_descent: &GradientDescent,
217    mut function: impl FnMut(&X) -> Result<Scalar, String>,
218    mut jacobian: impl FnMut(&X) -> Result<X, String>,
219    initial_guess: X,
220    indices: Vec<usize>,
221) -> Result<X, OptimizationError>
222where
223    X: Jacobian + Solution,
224    for<'a> &'a X: Mul<Scalar, Output = X>,
225{
226    let mut relative_scale = 0.0;
227    let mut residual: X;
228    let mut residual_change = initial_guess.clone() * 0.0;
229    let mut residual_norm;
230    let mut solution = initial_guess.clone();
231    let mut solution_change = solution.clone();
232    let mut step_size = INITIAL_STEP_SIZE;
233    let mut step_trial;
234    for iteration in 0..=gradient_descent.max_steps {
235        residual = jacobian(&solution)?;
236        residual.zero_out(&indices);
237        residual_norm = gradient_descent.norm.apply(&residual);
238        if gradient_descent.rel_tol.is_some() && iteration == 0 {
239            relative_scale = gradient_descent.norm.apply(&residual)
240        }
241        if residual_norm < gradient_descent.abs_tol {
242            return Ok(solution);
243        } else if let Some(rel_tol) = gradient_descent.rel_tol
244            && residual_norm / relative_scale < rel_tol
245        {
246            return Ok(solution);
247        } else {
248            solution_change -= &solution;
249            residual_change -= &residual;
250            step_trial =
251                residual_change.full_contraction(&solution_change) / residual_change.norm_squared();
252            if step_trial.abs() > 0.0 && !step_trial.is_nan() {
253                step_size = step_trial.abs()
254            }
255            step_size = gradient_descent.backtracking_line_search(
256                &mut function,
257                &mut jacobian,
258                &solution,
259                &residual,
260                &residual,
261                step_size,
262            )?;
263            residual_change = residual.clone();
264            solution_change = solution.clone();
265            solution -= residual * step_size;
266        }
267    }
268    Err(OptimizationError::MaximumStepsReached(
269        gradient_descent.max_steps,
270        format!("{gradient_descent:?}"),
271    ))
272}
273
274fn constrained<X>(
275    gradient_descent: &GradientDescent,
276    mut jacobian: impl FnMut(&X) -> Result<X, String>,
277    initial_guess: X,
278    constraint_matrix: Matrix,
279    constraint_rhs: Vector,
280) -> Result<X, OptimizationError>
281where
282    X: Jacobian,
283    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
284{
285    if !matches!(gradient_descent.line_search, LineSearch::None) {
286        panic!("Line search needs the exact penalty function in constrained optimization.")
287    }
288    let mut residual_solution;
289    let mut residual_solution_change = initial_guess.clone() * 0.0;
290    let mut solution = initial_guess.clone();
291    let mut solution_change = solution.clone();
292    let mut step_size_solution = INITIAL_STEP_SIZE;
293    let mut step_trial_solution;
294    let num_constraints = constraint_rhs.len();
295    let mut residual_multipliers;
296    let mut residual_multipliers_change = Vector::zero(num_constraints);
297    let mut multipliers = Vector::zero(num_constraints);
298    let mut multipliers_change = Vector::zero(num_constraints);
299    let mut step_size_multipliers = INITIAL_STEP_SIZE;
300    let mut step_trial_multipliers;
301    let mut step_size;
302    for _ in 0..=gradient_descent.max_steps {
303        residual_solution = jacobian(&solution)? - &multipliers * &constraint_matrix;
304        residual_multipliers = &constraint_rhs - &constraint_matrix * &solution;
305        if gradient_descent.norm.apply(&residual_solution) < gradient_descent.abs_tol
306            && gradient_descent.norm.apply(&residual_multipliers) < gradient_descent.abs_tol
307        {
308            return Ok(solution);
309        } else {
310            solution_change -= &solution;
311            residual_solution_change -= &residual_solution;
312            step_trial_solution = residual_solution_change.full_contraction(&solution_change)
313                / residual_solution_change.norm_squared();
314            if step_trial_solution.abs() > 0.0 && !step_trial_solution.is_nan() {
315                step_size_solution = step_trial_solution.abs()
316            }
317            residual_solution_change = residual_solution.clone();
318            solution_change = solution.clone();
319            multipliers_change -= &multipliers;
320            residual_multipliers_change -= &residual_multipliers;
321            step_trial_multipliers = residual_multipliers_change
322                .full_contraction(&multipliers_change)
323                / residual_multipliers_change.norm_squared();
324            if step_trial_multipliers.abs() > 0.0 && !step_trial_multipliers.is_nan() {
325                step_size_multipliers = step_trial_multipliers.abs()
326            }
327            residual_multipliers_change = residual_multipliers.clone();
328            multipliers_change = multipliers.clone();
329            step_size = step_size_solution.min(step_size_multipliers);
330            solution -= residual_solution * step_size;
331            multipliers += residual_multipliers * step_size;
332        }
333    }
334    Err(OptimizationError::MaximumStepsReached(
335        gradient_descent.max_steps,
336        format!("{gradient_descent:?}"),
337    ))
338}
339
340fn constrained_dual<X>(
341    gradient_descent: &GradientDescent,
342    mut jacobian: impl FnMut(&X) -> Result<X, String>,
343    initial_guess: X,
344    constraint_matrix: Matrix,
345    constraint_rhs: Vector,
346) -> Result<X, OptimizationError>
347where
348    X: Jacobian + Solution,
349    for<'a> &'a X: Mul<Scalar, Output = X>,
350    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
351{
352    if !matches!(gradient_descent.line_search, LineSearch::None) {
353        panic!("Line search needs the exact penalty function in constrained optimization.")
354    }
355    let num_constraints = constraint_rhs.len();
356    let mut multipliers = Vector::zero(num_constraints);
357    let mut multipliers_change = multipliers.clone();
358    let mut residual;
359    let mut residual_change = Vector::zero(num_constraints);
360    let mut solution = initial_guess;
361    let mut step_size = INITIAL_STEP_SIZE;
362    let mut step_trial;
363    for _ in 0..=gradient_descent.max_steps {
364        if let Ok(result) = unconstrained(
365            gradient_descent,
366            |_: &X| {
367                panic!("Line search needs the exact penalty function in constrained optimization.")
368            },
369            &mut jacobian,
370            solution.clone(),
371            Some((&constraint_matrix, &multipliers)),
372        ) {
373            solution = result;
374            residual = &constraint_rhs - &constraint_matrix * &solution;
375            if gradient_descent.norm.apply(&residual) < gradient_descent.abs_tol {
376                return Ok(solution);
377            } else {
378                multipliers_change -= &multipliers;
379                residual_change -= &residual;
380                step_trial = residual_change.full_contraction(&multipliers_change)
381                    / residual_change.norm_squared();
382                if step_trial.abs() > 0.0 && !step_trial.is_nan() {
383                    step_size = step_trial.abs()
384                }
385                residual_change = residual.clone();
386                multipliers_change = multipliers.clone();
387                multipliers += residual * step_size;
388            }
389        } else {
390            //
391            // This sort of acts like LineSearch::Error, does it not?
392            //
393            multipliers -= (multipliers.clone() - &multipliers_change) * CUTBACK_FACTOR_MINUS_ONE;
394            step_size *= CUTBACK_FACTOR;
395        }
396    }
397    Err(OptimizationError::MaximumStepsReached(
398        gradient_descent.max_steps,
399        format!("{gradient_descent:?}"),
400    ))
401}
conspire/math/optimize/gradient_descent/mod.rs

conspire/math/optimize/gradient_descent/
mod.rs
