conspire/math/optimize/nonlinear/gradient_descent/
mod.rs

1#[cfg(test)]
2mod test;
3
4use crate::math::{
5    Erase, Is, Jacobian, Matrix, Norm, Scalar, Solution, Tensor, Vector,
6    optimize::{
7        BacktrackingLineSearch, EqualityConstraint, EqualityConstraints, NoSearch, ObjectiveMerit,
8        Optimization, OptimizationError, RootFinding, StepSize, Tolerances, UNSUPPORTED,
9        no_objective,
10        nonlinear::step_size::{BarzilaiBorwein, Secant, StepEstimator},
11    },
12    sparse::SparseSolver,
13};
14use crate::units::{UnitDiv, UnitMul, UnitSum};
15use std::{
16    fmt::{self, Debug, Formatter},
17    ops::Mul,
18};
19
20const CUTBACK_FACTOR: Scalar = 0.8;
21const CUTBACK_FACTOR_MINUS_ONE: Scalar = 1.0 - CUTBACK_FACTOR;
22const INITIAL_STEP_SIZE: Scalar = 1e-2;
23
24/// The method of gradient descent.
25pub struct GradientDescent<S = NoSearch, Q = BarzilaiBorwein> {
26    /// Absolute error tolerances.
27    pub abs_tol: Tolerances,
28    /// Lagrangian dual.
29    pub dual: bool,
30    /// Norm type for error evaluation.
31    pub error_norm: Norm,
32    /// Line search algorithm.
33    pub line_search: S,
34    /// Maximum number of steps.
35    pub max_steps: usize,
36    /// Relative error tolerance.
37    pub rel_tol: Option<Scalar>,
38    /// Step size estimator.
39    pub step_estimator: Q,
40}
41
42impl<S, Q, J, X> BacktrackingLineSearch<J, X> for GradientDescent<S, Q>
43where
44    S: ObjectiveMerit,
45    Q: StepEstimator,
46{
47    type Search = S;
48    fn line_search(&self) -> &S {
49        &self.line_search
50    }
51}
52
53impl<S, Q> Debug for GradientDescent<S, Q>
54where
55    S: ObjectiveMerit,
56    Q: StepEstimator,
57{
58    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
59        write!(
60            f,
61            "GradientDescent {{ abs_tol: {:?}, dual: {:?}, line_search: {}, max_steps: {:?}, rel_tol: {:?}, step_estimator: {} }}",
62            self.abs_tol,
63            self.dual,
64            self.line_search,
65            self.max_steps,
66            self.rel_tol,
67            self.step_estimator
68        )
69    }
70}
71
72impl GradientDescent {
73    #[allow(clippy::should_implement_trait)]
74    pub fn default() -> Self {
75        Self::with(NoSearch)
76    }
77}
78
79impl<S> Default for GradientDescent<S>
80where
81    S: Default,
82{
83    fn default() -> Self {
84        Self::with(S::default())
85    }
86}
87
88impl<S, Q> GradientDescent<S, Q> {
89    /// The same solver, searching along each step with another line search.
90    pub fn with_line_search<T>(self, line_search: T) -> GradientDescent<T, Q> {
91        GradientDescent {
92            abs_tol: self.abs_tol,
93            dual: self.dual,
94            error_norm: self.error_norm,
95            line_search,
96            max_steps: self.max_steps,
97            rel_tol: self.rel_tol,
98            step_estimator: self.step_estimator,
99        }
100    }
101    /// The same solver, estimating each step size with another estimator.
102    pub fn with_step_estimator<T>(self, step_estimator: T) -> GradientDescent<S, T> {
103        GradientDescent {
104            abs_tol: self.abs_tol,
105            dual: self.dual,
106            error_norm: self.error_norm,
107            line_search: self.line_search,
108            max_steps: self.max_steps,
109            rel_tol: self.rel_tol,
110            step_estimator,
111        }
112    }
113}
114
115impl<S> GradientDescent<S> {
116    fn with(line_search: S) -> Self {
117        Self {
118            abs_tol: Tolerances::default(),
119            dual: false,
120            error_norm: Norm::Chebyshev,
121            line_search,
122            max_steps: 250,
123            rel_tol: None,
124            step_estimator: BarzilaiBorwein::default(),
125        }
126    }
127}
128
129impl<Q, F, X, E> RootFinding<F, (), X> for GradientDescent<NoSearch, Q>
130where
131    Q: StepEstimator,
132    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
133    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
134    X: Erase<Erased = E> + Jacobian + Solution,
135    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
136    E: Tensor,
137    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
138{
139    fn root(
140        &self,
141        function: impl FnMut(&X) -> Result<F, String>,
142        _jacobian: impl FnMut(&X) -> Result<(), String>,
143        mut initial_guess: X,
144        equality_constraint: EqualityConstraints,
145        _sparse: Option<SparseSolver>,
146    ) -> Result<X, OptimizationError> {
147        let Some(equality_constraint) =
148            EqualityConstraint::lower_single(equality_constraint, &mut initial_guess)
149        else {
150            return Err(OptimizationError::Intermediate(UNSUPPORTED.to_string()));
151        };
152        match equality_constraint {
153            EqualityConstraint::Fixed(indices) => {
154                constrained_fixed(self, no_objective, function, initial_guess, indices)
155            }
156            EqualityConstraint::Prescribed(_) => {
157                unreachable!("Prescribed constraints are lowered first.")
158            }
159            EqualityConstraint::Multipliers(_) => {
160                Err(OptimizationError::Intermediate(UNSUPPORTED.to_string()))
161            }
162            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => {
163                if self.dual {
164                    constrained_dual(
165                        self,
166                        function,
167                        initial_guess,
168                        constraint_matrix,
169                        constraint_rhs,
170                    )
171                } else {
172                    constrained(
173                        self,
174                        function,
175                        initial_guess,
176                        constraint_matrix,
177                        constraint_rhs,
178                    )
179                }
180            }
181            EqualityConstraint::None => {
182                unconstrained(self, no_objective, function, initial_guess, None)
183            }
184        }
185    }
186}
187
188impl<S, Q, F, J, X, E> Optimization<F, J, (), X> for GradientDescent<S, Q>
189where
190    S: ObjectiveMerit,
191    Q: StepEstimator,
192    F: Erase<Erased = Scalar> + Tensor,
193    <J as Tensor>::Unit: UnitMul<<X as Tensor>::Unit>,
194    <<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output: UnitSum,
195    <<<J as Tensor>::Unit as UnitMul<<X as Tensor>::Unit>>::Output as UnitSum>::Output:
196        Is<<F as Tensor>::Unit>,
197    J: Erase<Erased = E> + Jacobian + Mul<StepSize<J, X>, Output = X>,
198    for<'a> &'a J: Mul<StepSize<J, X>, Output = X>,
199    X: Erase<Erased = E> + Jacobian + Solution,
200    <X as Tensor>::Unit: UnitDiv<<J as Tensor>::Unit>,
201    E: Tensor,
202    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
203{
204    fn minimize(
205        &self,
206        mut function: impl FnMut(&X) -> Result<F, String>,
207        jacobian: impl FnMut(&X) -> Result<J, String>,
208        _hessian: impl FnMut(&X) -> Result<(), String>,
209        mut initial_guess: X,
210        equality_constraint: EqualityConstraints,
211        _sparse: Option<SparseSolver>,
212    ) -> Result<X, OptimizationError> {
213        let objective = move |argument: &X| function(argument).map(|value| *value.erase());
214        let Some(equality_constraint) =
215            EqualityConstraint::lower_single(equality_constraint, &mut initial_guess)
216        else {
217            return Err(OptimizationError::Intermediate(UNSUPPORTED.to_string()));
218        };
219        match equality_constraint {
220            EqualityConstraint::Fixed(indices) => {
221                constrained_fixed(self, objective, jacobian, initial_guess, indices)
222            }
223            EqualityConstraint::Prescribed(_) => {
224                unreachable!("Prescribed constraints are lowered first.")
225            }
226            EqualityConstraint::Multipliers(_) => {
227                Err(OptimizationError::Intermediate(UNSUPPORTED.to_string()))
228            }
229            EqualityConstraint::Linear(constraint_matrix, constraint_rhs) => {
230                if self.dual {
231                    constrained_dual(
232                        self,
233                        jacobian,
234                        initial_guess,
235                        constraint_matrix,
236                        constraint_rhs,
237                    )
238                } else {
239                    constrained(
240                        self,
241                        jacobian,
242                        initial_guess,
243                        constraint_matrix,
244                        constraint_rhs,
245                    )
246                }
247            }
248            EqualityConstraint::None => {
249                unconstrained(self, objective, jacobian, initial_guess, None)
250            }
251        }
252    }
253}
254
255fn unconstrained<S, Q, F, X, E>(
256    gradient_descent: &GradientDescent<S, Q>,
257    mut function: impl FnMut(&X) -> Result<Scalar, String>,
258    mut jacobian: impl FnMut(&X) -> Result<F, String>,
259    initial_guess: X,
260    linear_equality_constraint: Option<(&Matrix, &Vector)>,
261) -> Result<X, OptimizationError>
262where
263    S: ObjectiveMerit,
264    Q: StepEstimator,
265    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
266    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
267    X: Erase<Erased = E> + Jacobian + Solution,
268    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
269    E: Tensor,
270{
271    let constraint = if let Some((constraint_matrix, multipliers)) = linear_equality_constraint {
272        Some(multipliers * constraint_matrix)
273    } else {
274        None
275    };
276    let mut residual;
277    let mut secant = Secant::new(&initial_guess);
278    let mut solution = initial_guess;
279    let mut step_size = INITIAL_STEP_SIZE;
280    let mut steps = 0;
281    loop {
282        residual = if let Some(ref extra) = constraint {
283            jacobian(&solution)? - extra
284        } else {
285            jacobian(&solution)?
286        };
287        if gradient_descent.error_norm.apply(&residual) < gradient_descent.abs_tol.residual() {
288            return Ok(solution);
289        } else if steps == gradient_descent.max_steps {
290            return Err(OptimizationError::MaximumStepsReached(
291                gradient_descent.max_steps,
292                format!("{gradient_descent:?}"),
293            ));
294        } else {
295            steps += 1;
296            step_size = secant.step_size(
297                &gradient_descent.step_estimator,
298                &solution,
299                &residual,
300                step_size,
301            );
302            step_size = gradient_descent.backtracking_line_search::<F, E>(
303                |trial: &X, _: Scalar| function(trial),
304                &mut jacobian,
305                &solution,
306                &residual,
307                &residual,
308                step_size,
309            )?;
310            solution -= residual * StepSize::<F, X>::new(step_size);
311        }
312    }
313}
314
315fn constrained_fixed<S, Q, F, X, E>(
316    gradient_descent: &GradientDescent<S, Q>,
317    mut function: impl FnMut(&X) -> Result<Scalar, String>,
318    mut jacobian: impl FnMut(&X) -> Result<F, String>,
319    initial_guess: X,
320    indices: Vec<usize>,
321) -> Result<X, OptimizationError>
322where
323    S: ObjectiveMerit,
324    Q: StepEstimator,
325    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
326    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
327    X: Erase<Erased = E> + Jacobian + Solution,
328    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
329    E: Tensor,
330{
331    let mut relative_scale = 0.0;
332    let mut residual: F;
333    let mut residual_norm;
334    let mut secant = Secant::new(&initial_guess);
335    let mut solution = initial_guess;
336    let mut step_size = INITIAL_STEP_SIZE;
337    let mut steps = 0;
338    loop {
339        residual = jacobian(&solution)?;
340        residual.zero_out(&indices);
341        residual_norm = gradient_descent.error_norm.measure(&residual);
342        if gradient_descent.rel_tol.is_some() && steps == 0 {
343            relative_scale = gradient_descent.error_norm.measure(&residual)
344        }
345        if residual_norm < gradient_descent.abs_tol.residual {
346            return Ok(solution);
347        } else if let Some(rel_tol) = gradient_descent.rel_tol
348            && residual_norm / relative_scale < rel_tol
349        {
350            return Ok(solution);
351        } else if steps == gradient_descent.max_steps {
352            return Err(OptimizationError::MaximumStepsReached(
353                gradient_descent.max_steps,
354                format!("{gradient_descent:?}"),
355            ));
356        } else {
357            steps += 1;
358            step_size = secant.step_size(
359                &gradient_descent.step_estimator,
360                &solution,
361                &residual,
362                step_size,
363            );
364            step_size = gradient_descent.backtracking_line_search::<F, E>(
365                |trial: &X, _: Scalar| function(trial),
366                &mut jacobian,
367                &solution,
368                &residual,
369                &residual,
370                step_size,
371            )?;
372            solution -= residual * StepSize::<F, X>::new(step_size);
373        }
374    }
375}
376
377fn constrained<S, Q, F, X, E>(
378    gradient_descent: &GradientDescent<S, Q>,
379    mut jacobian: impl FnMut(&X) -> Result<F, String>,
380    initial_guess: X,
381    constraint_matrix: Matrix,
382    constraint_rhs: Vector,
383) -> Result<X, OptimizationError>
384where
385    S: ObjectiveMerit,
386    Q: StepEstimator,
387    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
388    X: Erase<Erased = E> + Jacobian,
389    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
390    E: Tensor,
391    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
392{
393    if S::SEARCHES {
394        panic!("Line search needs the exact penalty function in constrained optimization.")
395    }
396    let mut residual_solution;
397    let mut secant_solution = Secant::new(&initial_guess);
398    let mut solution = initial_guess;
399    let mut step_size_solution = INITIAL_STEP_SIZE;
400    let num_constraints = constraint_rhs.len();
401    let mut residual_multipliers;
402    let mut multipliers = Vector::zero(num_constraints);
403    let mut secant_multipliers = Secant::new(&multipliers);
404    let mut step_size_multipliers = INITIAL_STEP_SIZE;
405    let mut step_size;
406    let mut steps = 0;
407    loop {
408        residual_solution = jacobian(&solution)? - &multipliers * &constraint_matrix;
409        residual_multipliers = &constraint_rhs - &constraint_matrix * &solution;
410        if gradient_descent.error_norm.apply(&residual_solution)
411            < gradient_descent.abs_tol.residual()
412            && gradient_descent.error_norm.apply(&residual_multipliers)
413                < gradient_descent.abs_tol.constraint()
414        {
415            return Ok(solution);
416        } else if steps == gradient_descent.max_steps {
417            return Err(OptimizationError::MaximumStepsReached(
418                gradient_descent.max_steps,
419                format!("{gradient_descent:?}"),
420            ));
421        } else {
422            steps += 1;
423            step_size_solution = secant_solution.step_size(
424                &gradient_descent.step_estimator,
425                &solution,
426                &residual_solution,
427                step_size_solution,
428            );
429            step_size_multipliers = secant_multipliers.step_size(
430                &gradient_descent.step_estimator,
431                &multipliers,
432                &residual_multipliers,
433                step_size_multipliers,
434            );
435            step_size = step_size_solution.min(step_size_multipliers);
436            solution -= residual_solution * StepSize::<F, X>::new(step_size);
437            multipliers += residual_multipliers * step_size;
438        }
439    }
440}
441
442fn constrained_dual<S, Q, F, X, E>(
443    gradient_descent: &GradientDescent<S, Q>,
444    mut jacobian: impl FnMut(&X) -> Result<F, String>,
445    initial_guess: X,
446    constraint_matrix: Matrix,
447    constraint_rhs: Vector,
448) -> Result<X, OptimizationError>
449where
450    S: ObjectiveMerit,
451    Q: StepEstimator,
452    F: Erase<Erased = E> + Jacobian + Mul<StepSize<F, X>, Output = X>,
453    for<'a> &'a F: Mul<StepSize<F, X>, Output = X>,
454    X: Erase<Erased = E> + Jacobian + Solution,
455    <X as Tensor>::Unit: UnitDiv<<F as Tensor>::Unit>,
456    E: Tensor,
457    for<'a> &'a Matrix: Mul<&'a X, Output = Vector>,
458{
459    if S::SEARCHES {
460        panic!("Line search needs the exact penalty function in constrained optimization.")
461    }
462    let num_constraints = constraint_rhs.len();
463    let mut multipliers = Vector::zero(num_constraints);
464    let mut secant = Secant::new(&multipliers);
465    let mut residual;
466    let mut solution = initial_guess;
467    let mut step_size = INITIAL_STEP_SIZE;
468    for _ in 0..gradient_descent.max_steps {
469        if let Ok(result) = unconstrained(
470            gradient_descent,
471            |_: &X| {
472                panic!("Line search needs the exact penalty function in constrained optimization.")
473            },
474            &mut jacobian,
475            solution.clone(),
476            Some((&constraint_matrix, &multipliers)),
477        ) {
478            solution = result;
479            residual = &constraint_rhs - &constraint_matrix * &solution;
480            if gradient_descent.error_norm.apply(&residual) < gradient_descent.abs_tol.constraint()
481            {
482                return Ok(solution);
483            } else {
484                step_size = secant.step_size(
485                    &gradient_descent.step_estimator,
486                    &multipliers,
487                    &residual,
488                    step_size,
489                );
490                multipliers += residual * step_size;
491            }
492        } else {
493            //
494            // This sort of acts like LineSearch::Error, does it not?
495            //
496            multipliers -= (multipliers.clone() - secant.solution()) * CUTBACK_FACTOR_MINUS_ONE;
497            step_size *= CUTBACK_FACTOR;
498        }
499    }
500    Err(OptimizationError::MaximumStepsReached(
501        gradient_descent.max_steps,
502        format!("{gradient_descent:?}"),
503    ))
504}
conspire/math/optimize/nonlinear/gradient_descent/mod.rs

conspire/math/optimize/nonlinear/gradient_descent/
mod.rs
