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

conspire/math/optimize/gradient_descent/
mod.rs
