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, StepSize, Tolerances, ZerothOrderRootFinding,
8};
9use crate::math::{Erase, Is, Norm};
10use crate::units::{UnitDiv, UnitMul, UnitSum};
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 tolerances.
23    pub abs_tol: Tolerances,
24    /// Lagrangian dual.
25    pub dual: bool,
26    /// Norm type for error evaluation.
27    pub error_norm: Norm,
28    /// Line search algorithm.
29    pub line_search: LineSearch,
30    /// Maximum number of steps.
31    pub max_steps: usize,
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: Tolerances::default(),
56            dual: false,
57            error_norm: Norm::Chebyshev,
58            line_search: LineSearch::None,
59            max_steps: 250,
60            rel_tol: None,
61        }
62    }
63}
64
65impl<F, X, E> ZerothOrderRootFinding<F, X> for GradientDescent
66where
67    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
68    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
69    X: Erase<Erased = E> + Jacobian + Solution,
70    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
71    E: Tensor,
72    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
73{
74    fn root(
75        &self,
76        function: impl FnMut(&X) -> Result<F, String>,
77        initial_guess: X,
78        equality_constraint: EqualityConstraint,
79    ) -> Result<X, OptimizationError> {
80        match equality_constraint {
81            EqualityConstraint::Fixed(indices) => constrained_fixed(
82                self,
83                |_: &X| panic!("No line search in root finding."),
84                function,
85                initial_guess,
86                indices,
87            ),
88            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => {
89                if self.dual {
90                    constrained_dual(
91                        self,
92                        function,
93                        initial_guess,
94                        constraint_matrix,
95                        constraint_rhs,
96                    )
97                } else {
98                    constrained(
99                        self,
100                        function,
101                        initial_guess,
102                        constraint_matrix,
103                        constraint_rhs,
104                    )
105                }
106            }
107            EqualityConstraint::None => unconstrained(
108                self,
109                |_: &X| panic!("No line search in root finding."),
110                function,
111                initial_guess,
112                None,
113            ),
114        }
115    }
116}
117
118impl<F, J, X, E> FirstOrderOptimization<F, J, X> for GradientDescent
119where
120    F: Erase<Erased = Scalar> + Tensor,
121    <J as Tensor>::Unit: UnitMul<<X as Tensor>::Unit>,
122    <<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output: UnitSum,
123    <<<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output as UnitSum>::Output:
124        Is<<F as Tensor>::Unit>,
125    J: Erase<Erased = E> + Jacobian + Mul<StepSize<J, X>, Output = X>,
126    for<'a> &'a J: Mul<StepSize<J, X>, Output = X>,
127    X: Erase<Erased = E> + Jacobian + Solution,
128    <X as Tensor>::Unit: UnitDiv<<J as Tensor>::Unit>,
129    E: Tensor,
130    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
131{
132    fn minimize(
133        &self,
134        mut function: impl FnMut(&X) -> Result<F, String>,
135        jacobian: impl FnMut(&X) -> Result<J, String>,
136        initial_guess: X,
137        equality_constraint: EqualityConstraint,
138    ) -> Result<X, OptimizationError> {
139        let objective = move |argument: &X| function(argument).map(|value| *value.erase());
140        match equality_constraint {
141            EqualityConstraint::Fixed(indices) => {
142                constrained_fixed(self, objective, jacobian, initial_guess, indices)
143            }
144            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => {
145                if self.dual {
146                    constrained_dual(
147                        self,
148                        jacobian,
149                        initial_guess,
150                        constraint_matrix,
151                        constraint_rhs,
152                    )
153                } else {
154                    constrained(
155                        self,
156                        jacobian,
157                        initial_guess,
158                        constraint_matrix,
159                        constraint_rhs,
160                    )
161                }
162            }
163            EqualityConstraint::None => {
164                unconstrained(self, objective, jacobian, initial_guess, None)
165            }
166        }
167    }
168}
169
170fn unconstrained<F, X, E>(
171    gradient_descent: &GradientDescent,
172    mut function: impl FnMut(&X) -> Result<Scalar, String>,
173    mut jacobian: impl FnMut(&X) -> Result<F, String>,
174    initial_guess: X,
175    linear_equality_constraint: Option<(&Matrix, &Vector)>,
176) -> Result<X, OptimizationError>
177where
178    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
179    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
180    X: Erase<Erased = E> + Jacobian + Solution,
181    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
182    E: Tensor,
183{
184    let constraint = if let Some((constraint_matrix, multipliers)) = linear_equality_constraint {
185        Some(multipliers * constraint_matrix)
186    } else {
187        None
188    };
189    let mut residual;
190    let mut residual_change = None;
191    let mut solution = initial_guess.clone();
192    let mut solution_change = solution.clone();
193    let mut step_size = INITIAL_STEP_SIZE;
194    let mut step_trial;
195    let mut steps = 0;
196    loop {
197        residual = if let Some(ref extra) = constraint {
198            jacobian(&solution)? - extra
199        } else {
200            jacobian(&solution)?
201        };
202        if gradient_descent.error_norm.apply(&residual) < gradient_descent.abs_tol.residual() {
203            return Ok(solution);
204        } else if steps == gradient_descent.max_steps {
205            return Err(OptimizationError::MaximumStepsReached(
206                gradient_descent.max_steps,
207                format!("{gradient_descent:?}"),
208            ));
209        } else {
210            steps += 1;
211            solution_change -= &solution;
212            let change = residual_change.get_or_insert_with(|| zeroed(&residual));
213            *change -= &residual;
214            step_trial = change.erase().full_contraction(solution_change.erase())
215                / change.erase().full_contraction(change.erase());
216            if step_trial.abs() > 0.0 && !step_trial.is_nan() {
217                step_size = step_trial.abs()
218            }
219            step_size = gradient_descent.backtracking_line_search::<F, E>(
220                |trial: &X, _: Scalar| function(trial),
221                &mut jacobian,
222                &solution,
223                &residual,
224                &residual,
225                step_size,
226            )?;
227            *change = residual.clone();
228            solution_change = solution.clone();
229            solution -= residual * StepSize::<F, X>::new(step_size);
230        }
231    }
232}
233
234fn constrained_fixed<F, X, E>(
235    gradient_descent: &GradientDescent,
236    mut function: impl FnMut(&X) -> Result<Scalar, String>,
237    mut jacobian: impl FnMut(&X) -> Result<F, String>,
238    initial_guess: X,
239    indices: Vec<usize>,
240) -> Result<X, OptimizationError>
241where
242    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
243    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
244    X: Erase<Erased = E> + Jacobian + Solution,
245    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
246    E: Tensor,
247{
248    let mut relative_scale = 0.0;
249    let mut residual: F;
250    let mut residual_change = None;
251    let mut residual_norm;
252    let mut solution = initial_guess.clone();
253    let mut solution_change = solution.clone();
254    let mut step_size = INITIAL_STEP_SIZE;
255    let mut step_trial;
256    let mut steps = 0;
257    loop {
258        residual = jacobian(&solution)?;
259        residual.zero_out(&indices);
260        residual_norm = gradient_descent.error_norm.measure(&residual);
261        if gradient_descent.rel_tol.is_some() && steps == 0 {
262            relative_scale = gradient_descent.error_norm.measure(&residual)
263        }
264        if residual_norm < gradient_descent.abs_tol.residual {
265            return Ok(solution);
266        } else if let Some(rel_tol) = gradient_descent.rel_tol
267            && residual_norm / relative_scale < rel_tol
268        {
269            return Ok(solution);
270        } else if steps == gradient_descent.max_steps {
271            return Err(OptimizationError::MaximumStepsReached(
272                gradient_descent.max_steps,
273                format!("{gradient_descent:?}"),
274            ));
275        } else {
276            steps += 1;
277            solution_change -= &solution;
278            let change = residual_change.get_or_insert_with(|| zeroed(&residual));
279            *change -= &residual;
280            step_trial = change.erase().full_contraction(solution_change.erase())
281                / change.erase().full_contraction(change.erase());
282            if step_trial.abs() > 0.0 && !step_trial.is_nan() {
283                step_size = step_trial.abs()
284            }
285            step_size = gradient_descent.backtracking_line_search::<F, E>(
286                |trial: &X, _: Scalar| function(trial),
287                &mut jacobian,
288                &solution,
289                &residual,
290                &residual,
291                step_size,
292            )?;
293            *change = residual.clone();
294            solution_change = solution.clone();
295            solution -= residual * StepSize::<F, X>::new(step_size);
296        }
297    }
298}
299
300fn constrained<F, X, E>(
301    gradient_descent: &GradientDescent,
302    mut jacobian: impl FnMut(&X) -> Result<F, String>,
303    initial_guess: X,
304    constraint_matrix: Matrix,
305    constraint_rhs: Vector,
306) -> Result<X, OptimizationError>
307where
308    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
309    X: Erase<Erased = E> + Jacobian,
310    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
311    E: Tensor,
312    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
313{
314    if !matches!(gradient_descent.line_search, LineSearch::None) {
315        panic!("Line search needs the exact penalty function in constrained optimization.")
316    }
317    let mut residual_solution;
318    let mut residual_solution_change = None;
319    let mut solution = initial_guess.clone();
320    let mut solution_change = solution.clone();
321    let mut step_size_solution = INITIAL_STEP_SIZE;
322    let mut step_trial_solution;
323    let num_constraints = constraint_rhs.len();
324    let mut residual_multipliers;
325    let mut residual_multipliers_change = Vector::zero(num_constraints);
326    let mut multipliers = Vector::zero(num_constraints);
327    let mut multipliers_change = Vector::zero(num_constraints);
328    let mut step_size_multipliers = INITIAL_STEP_SIZE;
329    let mut step_trial_multipliers;
330    let mut step_size;
331    let mut steps = 0;
332    loop {
333        residual_solution = jacobian(&solution)? - &multipliers * &constraint_matrix;
334        residual_multipliers = &constraint_rhs - &constraint_matrix * &solution;
335        if gradient_descent.error_norm.apply(&residual_solution)
336            < gradient_descent.abs_tol.residual()
337            && gradient_descent.error_norm.apply(&residual_multipliers)
338                < gradient_descent.abs_tol.constraint()
339        {
340            return Ok(solution);
341        } else if steps == gradient_descent.max_steps {
342            return Err(OptimizationError::MaximumStepsReached(
343                gradient_descent.max_steps,
344                format!("{gradient_descent:?}"),
345            ));
346        } else {
347            steps += 1;
348            solution_change -= &solution;
349            let change = residual_solution_change.get_or_insert_with(|| zeroed(&residual_solution));
350            *change -= &residual_solution;
351            step_trial_solution = change.erase().full_contraction(solution_change.erase())
352                / change.erase().full_contraction(change.erase());
353            if step_trial_solution.abs() > 0.0 && !step_trial_solution.is_nan() {
354                step_size_solution = step_trial_solution.abs()
355            }
356            *change = residual_solution.clone();
357            solution_change = solution.clone();
358            multipliers_change -= &multipliers;
359            residual_multipliers_change -= &residual_multipliers;
360            step_trial_multipliers = residual_multipliers_change
361                .full_contraction(&multipliers_change)
362                / residual_multipliers_change.full_contraction(&residual_multipliers_change);
363            if step_trial_multipliers.abs() > 0.0 && !step_trial_multipliers.is_nan() {
364                step_size_multipliers = step_trial_multipliers.abs()
365            }
366            residual_multipliers_change = residual_multipliers.clone();
367            multipliers_change = multipliers.clone();
368            step_size = step_size_solution.min(step_size_multipliers);
369            solution -= residual_solution * StepSize::<F, X>::new(step_size);
370            multipliers += residual_multipliers * step_size;
371        }
372    }
373}
374
375fn constrained_dual<F, X, E>(
376    gradient_descent: &GradientDescent,
377    mut jacobian: impl FnMut(&X) -> Result<F, String>,
378    initial_guess: X,
379    constraint_matrix: Matrix,
380    constraint_rhs: Vector,
381) -> Result<X, OptimizationError>
382where
383    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
384    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
385    X: Erase<Erased = E> + Jacobian + Solution,
386    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
387    E: Tensor,
388    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
389{
390    if !matches!(gradient_descent.line_search, LineSearch::None) {
391        panic!("Line search needs the exact penalty function in constrained optimization.")
392    }
393    let num_constraints = constraint_rhs.len();
394    let mut multipliers = Vector::zero(num_constraints);
395    let mut multipliers_change = multipliers.clone();
396    let mut residual;
397    let mut residual_change = Vector::zero(num_constraints);
398    let mut solution = initial_guess;
399    let mut step_size = INITIAL_STEP_SIZE;
400    let mut step_trial;
401    for _ in 0..gradient_descent.max_steps {
402        if let Ok(result) = unconstrained(
403            gradient_descent,
404            |_: &X| {
405                panic!("Line search needs the exact penalty function in constrained optimization.")
406            },
407            &mut jacobian,
408            solution.clone(),
409            Some((&constraint_matrix, &multipliers)),
410        ) {
411            solution = result;
412            residual = &constraint_rhs - &constraint_matrix * &solution;
413            if gradient_descent.error_norm.apply(&residual) < gradient_descent.abs_tol.constraint()
414            {
415                return Ok(solution);
416            } else {
417                multipliers_change -= &multipliers;
418                residual_change -= &residual;
419                step_trial = residual_change.full_contraction(&multipliers_change)
420                    / residual_change.full_contraction(&residual_change);
421                if step_trial.abs() > 0.0 && !step_trial.is_nan() {
422                    step_size = step_trial.abs()
423                }
424                residual_change = residual.clone();
425                multipliers_change = multipliers.clone();
426                multipliers += residual * step_size;
427            }
428        } else {
429            //
430            // This sort of acts like LineSearch::Error, does it not?
431            //
432            multipliers -= (multipliers.clone() - &multipliers_change) * CUTBACK_FACTOR_MINUS_ONE;
433            step_size *= CUTBACK_FACTOR;
434        }
435    }
436    Err(OptimizationError::MaximumStepsReached(
437        gradient_descent.max_steps,
438        format!("{gradient_descent:?}"),
439    ))
440}
441
442/// A zeroed copy, to start the difference the step size is estimated from.
443fn zeroed<F>(residual: &F) -> F
444where
445    F: Tensor,
446{
447    let mut zero = residual.clone();
448    zero *= 0.0;
449    zero
450}
conspire/math/optimize/gradient_descent/mod.rs

conspire/math/optimize/gradient_descent/
mod.rs
