Skip to main content

slain_num/
optim.rs

1pub fn golden(mut f: impl FnMut(f64) -> f64, mut a: f64, mut b: f64, tol: f64) -> (f64, f64) {
2    let g = (5f64.sqrt() - 1.0) / 2.0;
3    let mut c = b - g * (b - a);
4    let mut d = a + g * (b - a);
5    let (mut fc, mut fd) = (f(c), f(d));
6    while (b - a).abs() > tol {
7        if fc < fd {
8            b = d;
9            (d, fd) = (c, fc);
10            c = b - g * (b - a);
11            fc = f(c);
12        } else {
13            a = c;
14            (c, fc) = (d, fd);
15            d = a + g * (b - a);
16            fd = f(d);
17        }
18    }
19    let x = (a + b) / 2.0;
20    (x, f(x))
21}
22
23pub fn nelder_mead(mut f: impl FnMut(&[f64]) -> f64, start: &[f64], step: f64, iters: usize) -> (Vec<f64>, f64) {
24    let n = start.len();
25    let mut pts: Vec<Vec<f64>> = vec![start.to_vec()];
26    for i in 0..n {
27        let mut p = start.to_vec();
28        p[i] += step;
29        pts.push(p);
30    }
31    let mut vals: Vec<f64> = pts.iter().map(|p| f(p)).collect();
32    let lerp = |a: &[f64], b: &[f64], t: f64| -> Vec<f64> { a.iter().zip(b).map(|(x, y)| x + t * (y - x)).collect() };
33    for _ in 0..iters {
34        let mut order: Vec<usize> = (0..=n).collect();
35        order.sort_by(|&i, &j| vals[i].total_cmp(&vals[j]));
36        pts = order.iter().map(|&i| pts[i].clone()).collect();
37        vals = order.iter().map(|&i| vals[i]).collect();
38        let centroid: Vec<f64> = (0..n).map(|k| pts[..n].iter().map(|p| p[k]).sum::<f64>() / n as f64).collect();
39        let worst = pts[n].clone();
40        let refl = lerp(&worst, &centroid, 2.0);
41        let fr = f(&refl);
42        if fr < vals[0] {
43            let exp = lerp(&worst, &centroid, 3.0);
44            let fe = f(&exp);
45            (pts[n], vals[n]) = if fe < fr { (exp, fe) } else { (refl, fr) };
46        } else if fr < vals[n - 1] {
47            (pts[n], vals[n]) = (refl, fr);
48        } else {
49            let con = lerp(&worst, &centroid, 0.5);
50            let fcn = f(&con);
51            if fcn < vals[n] {
52                (pts[n], vals[n]) = (con, fcn);
53            } else {
54                for i in 1..=n {
55                    pts[i] = lerp(&pts[0], &pts[i], 0.5);
56                    vals[i] = f(&pts[i]);
57                }
58            }
59        }
60    }
61    let best = (0..=n).min_by(|&i, &j| vals[i].total_cmp(&vals[j])).unwrap_or(0);
62    (pts[best].clone(), vals[best])
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn finds_minima() {
71        let (x, _) = golden(|x| (x - 1.7).powi(2) + 3.0, -5.0, 5.0, 1e-8);
72        assert!((x - 1.7).abs() < 1e-6);
73        let (p, v) = nelder_mead(|p| (p[0] - 1.0).powi(2) + 10.0 * (p[1] + 2.0).powi(2), &[0.0, 0.0], 0.5, 400);
74        assert!((p[0] - 1.0).abs() < 1e-4 && (p[1] + 2.0).abs() < 1e-4 && v < 1e-7);
75    }
76}