geop_core_sketch/
bfgs.rs

1//! A dense BFGS minimizer with a backtracking (Armijo) line search.
2//!
3//! Sketches have at most a few hundred variables, so the inverse Hessian
4//! approximation is kept as a full `n x n` matrix.
5
6/// When [`minimize`] stops.
7#[derive(Clone, Copy, Debug)]
8pub struct BfgsOptions {
9    pub max_iterations: usize,
10    /// Stop once the objective is at or below this value.
11    pub f_tolerance: f64,
12    /// Stop once every gradient component is at or below this magnitude.
13    pub g_tolerance: f64,
14}
15
16#[derive(Clone, Debug)]
17pub struct BfgsResult {
18    pub x: Vec<f64>,
19    pub f: f64,
20    pub iterations: usize,
21}
22
23fn dot(a: &[f64], b: &[f64]) -> f64 {
24    a.iter().zip(b).map(|(x, y)| x * y).sum()
25}
26
27/// Minimize `f` from `x0`. `f` returns the objective and its gradient.
28///
29/// Always returns the best point found; whether that point is good enough is
30/// the caller's question (it knows what the objective means), which is why
31/// this reports no "converged" flag of its own.
32pub fn minimize(
33    f: impl Fn(&[f64]) -> (f64, Vec<f64>),
34    x0: Vec<f64>,
35    options: BfgsOptions,
36) -> BfgsResult {
37    let n = x0.len();
38    let mut x = x0;
39    let (mut fx, mut g) = f(&x);
40    // Inverse Hessian approximation, row-major. Starts as the identity and is
41    // rescaled after the first accepted step (Nocedal & Wright, eq. 6.20).
42    let mut h = vec![0.0; n * n];
43    for i in 0..n {
44        h[i * n + i] = 1.0;
45    }
46    let mut first_update = true;
47
48    for iteration in 0..options.max_iterations {
49        if fx <= options.f_tolerance || g.iter().all(|gi| gi.abs() <= options.g_tolerance) {
50            return BfgsResult {
51                x,
52                f: fx,
53                iterations: iteration,
54            };
55        }
56
57        // Search direction p = -H g; fall back to steepest descent if the
58        // approximation has stopped producing a descent direction.
59        let mut p: Vec<f64> = (0..n).map(|i| -dot(&h[i * n..(i + 1) * n], &g)).collect();
60        let mut slope = dot(&p, &g);
61        if slope >= 0.0 {
62            p = g.iter().map(|gi| -gi).collect();
63            slope = dot(&p, &g);
64            h.iter_mut().for_each(|v| *v = 0.0);
65            for i in 0..n {
66                h[i * n + i] = 1.0;
67            }
68            first_update = true;
69        }
70
71        // Backtracking line search on the Armijo condition.
72        let mut alpha = 1.0;
73        let mut accepted = None;
74        for _ in 0..60 {
75            let x_new: Vec<f64> = x.iter().zip(&p).map(|(xi, pi)| xi + alpha * pi).collect();
76            let (f_new, g_new) = f(&x_new);
77            if f_new.is_finite() && f_new <= fx + 1e-4 * alpha * slope {
78                accepted = Some((x_new, f_new, g_new));
79                break;
80            }
81            alpha *= 0.5;
82        }
83        let Some((x_new, f_new, g_new)) = accepted else {
84            // No decrease along a descent direction: we are at the limit of
85            // what floating point can resolve.
86            return BfgsResult {
87                x,
88                f: fx,
89                iterations: iteration,
90            };
91        };
92
93        let s: Vec<f64> = x_new.iter().zip(&x).map(|(a, b)| a - b).collect();
94        let y: Vec<f64> = g_new.iter().zip(&g).map(|(a, b)| a - b).collect();
95        let sy = dot(&s, &y);
96        // Skip the update unless it keeps H positive definite.
97        if sy > 1e-12 * dot(&s, &s).sqrt() * dot(&y, &y).sqrt() && sy > 0.0 {
98            if first_update {
99                let scale = sy / dot(&y, &y);
100                h.iter_mut().for_each(|v| *v = 0.0);
101                for i in 0..n {
102                    h[i * n + i] = scale;
103                }
104                first_update = false;
105            }
106            // H <- (I - rho s y^T) H (I - rho y s^T) + rho s s^T
107            let rho = 1.0 / sy;
108            let hy: Vec<f64> = (0..n).map(|i| dot(&h[i * n..(i + 1) * n], &y)).collect();
109            let yhy = dot(&y, &hy);
110            for i in 0..n {
111                for j in 0..n {
112                    h[i * n + j] += -rho * (hy[i] * s[j] + s[i] * hy[j])
113                        + (rho * rho * yhy + rho) * s[i] * s[j];
114                }
115            }
116        }
117
118        x = x_new;
119        fx = f_new;
120        g = g_new;
121    }
122
123    BfgsResult {
124        x,
125        f: fx,
126        iterations: options.max_iterations,
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn minimizes_rosenbrock() {
136        let rosenbrock = |x: &[f64]| {
137            let (a, b) = (x[0], x[1]);
138            let f = (1.0 - a).powi(2) + 100.0 * (b - a * a).powi(2);
139            let g = vec![
140                -2.0 * (1.0 - a) - 400.0 * a * (b - a * a),
141                200.0 * (b - a * a),
142            ];
143            (f, g)
144        };
145        let r = minimize(
146            rosenbrock,
147            vec![-1.2, 1.0],
148            BfgsOptions {
149                max_iterations: 500,
150                f_tolerance: 1e-20,
151                g_tolerance: 1e-12,
152            },
153        );
154        assert!(
155            (r.x[0] - 1.0).abs() < 1e-6 && (r.x[1] - 1.0).abs() < 1e-6,
156            "{r:?}"
157        );
158    }
159}
geop_core_sketch/bfgs.rs

geop_core_sketch/
bfgs.rs
