metalforge 0.1.4

//! CMA-ES — Covariance Matrix Adaptation Evolution Strategy
//! (Hansen & Ostermeier 2001; parameters per Hansen's tutorial, 2016).
//!
//! The reference black-box optimizer for hard continuous problems: it samples a
//! multivariate normal, then adapts its mean, step size, and full covariance
//! matrix to the local landscape — learning variable scales and correlations,
//! which makes it excel on ill-conditioned, non-separable surfaces (e.g.
//! Rosenbrock) where coordinate-wise methods crawl.
//!
//! The search runs in a **normalized `[0, 1]^n` space** (each variable mapped
//! from its bounds), so a single scalar `sigma0` is meaningful regardless of how
//! the box is scaled per dimension; candidates are clamped into the box for
//! evaluation. The covariance is diagonalized each generation with a Jacobi
//! eigensolver (no external linear-algebra dependency), keeping the run
//! deterministic for a given seed.

// Index-based loops are clearer than iterator chains for the matrix algebra and
// Jacobi rotations in this module.
#![allow(clippy::needless_range_loop)]

use super::Optimizer;
use crate::problem::Problem;
use crate::rng::Rng;
use crate::solution::{Report, Solution, StopReason};
use crate::termination::Termination;

/// CMA-ES configuration.
#[derive(Debug, Clone, Copy)]
pub struct CmaEs {
    /// Population size `λ`; `None` uses the default `4 + ⌊3 ln n⌋`.
    pub population: Option<usize>,
    /// Initial step size in the normalized `[0, 1]` space (the box is rescaled
    /// per dimension, so `0.3` ≈ covering a third of each variable's range).
    pub sigma0: f64,
    /// RNG seed; same seed + same problem + same budget ⇒ same result.
    pub seed: u64,
}

impl Default for CmaEs {
    fn default() -> Self {
        CmaEs {
            population: None,
            sigma0: 0.3,
            seed: 42,
        }
    }
}

impl Optimizer for CmaEs {
    fn with_seed(&self, seed: u64) -> Self {
        CmaEs { seed, ..*self }
    }

    fn optimize(&self, problem: &dyn Problem, term: &Termination) -> Report {
        let bounds = problem.bounds();
        let n = bounds.len();
        let mut rng = Rng::new(self.seed);

        // Strategy parameters (Hansen tutorial).
        let lambda = self
            .population
            .unwrap_or(4 + (3.0 * (n as f64).ln()) as usize)
            .max(4);
        let mu = lambda / 2;
        // Recombination weights (positive, log-decreasing), normalized to 1.
        let raw: Vec<f64> = (0..mu)
            .map(|i| ((mu as f64) + 0.5).ln() - ((i + 1) as f64).ln())
            .collect();
        let wsum: f64 = raw.iter().sum();
        let w: Vec<f64> = raw.iter().map(|&v| v / wsum).collect();
        let mu_eff = 1.0 / w.iter().map(|&v| v * v).sum::<f64>();

        let nf = n as f64;
        let c_sigma = (mu_eff + 2.0) / (nf + mu_eff + 5.0);
        let d_sigma = 1.0 + 2.0 * (((mu_eff - 1.0) / (nf + 1.0)).sqrt() - 1.0).max(0.0) + c_sigma;
        let c_c = (4.0 + mu_eff / nf) / (nf + 4.0 + 2.0 * mu_eff / nf);
        let c_1 = 2.0 / ((nf + 1.3).powi(2) + mu_eff);
        let c_mu =
            (1.0 - c_1).min(2.0 * (mu_eff - 2.0 + 1.0 / mu_eff) / ((nf + 2.0).powi(2) + mu_eff));
        // Expected length of an N(0, I) vector.
        let e_n = nf.sqrt() * (1.0 - 1.0 / (4.0 * nf) + 1.0 / (21.0 * nf * nf));

        // State, in normalized [0, 1] space.
        let mut mean = vec![0.5; n];
        let mut sigma = self.sigma0;
        let mut cov = identity(n);
        let mut p_sigma = vec![0.0; n];
        let mut p_c = vec![0.0; n];
        let mut generation = 0i32;

        let mut best = Solution {
            x: denormalize(&mean, bounds),
            value: f64::INFINITY,
        };
        let mut evaluations = 0usize;

        while term.reason(evaluations, best.value).is_none() {
            // Diagonalize C = B diag(d²) Bᵀ; D = sqrt of eigenvalues.
            let (eigvals, b) = jacobi_eigen(&cov);
            let d: Vec<f64> = eigvals.iter().map(|&v| v.max(1e-20).sqrt()).collect();

            // Sample λ offspring; store their normalized step y = B D z.
            let mut pop: Vec<(f64, Vec<f64>)> = Vec::with_capacity(lambda);
            for _ in 0..lambda {
                let z: Vec<f64> = (0..n).map(|_| rng.normal()).collect();
                let dz: Vec<f64> = (0..n).map(|j| d[j] * z[j]).collect();
                let y = matvec(&b, &dz); // B (D z)
                let u: Vec<f64> = (0..n).map(|i| mean[i] + sigma * y[i]).collect();
                let x = denormalize_clamped(&u, bounds);
                let f = problem.objective(&x);
                evaluations += 1;
                let f = if f.is_finite() { f } else { f64::INFINITY };
                if f < best.value {
                    best = Solution { x, value: f };
                }
                pop.push((f, y));
            }

            // Rank by fitness (ascending) and recombine the μ best steps.
            pop.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
            let mut y_w = vec![0.0; n];
            for (i, wi) in w.iter().enumerate() {
                for k in 0..n {
                    y_w[k] += wi * pop[i].1[k];
                }
            }

            // Move the mean: m ← m + σ y_w (kept inside the box).
            for k in 0..n {
                mean[k] = (mean[k] + sigma * y_w[k]).clamp(0.0, 1.0);
            }

            // Step-size evolution path: p_σ ← (1−c_σ) p_σ + √(c_σ(2−c_σ)μ_eff) C^{-1/2} y_w.
            let c_inv_sqrt_yw = c_inv_sqrt_mul(&b, &d, &y_w);
            let cs_factor = (c_sigma * (2.0 - c_sigma) * mu_eff).sqrt();
            for k in 0..n {
                p_sigma[k] = (1.0 - c_sigma) * p_sigma[k] + cs_factor * c_inv_sqrt_yw[k];
            }
            let ps_norm = norm(&p_sigma);

            generation += 1;
            // Heaviside step stalls the rank-one update when σ would jump.
            let hsig = if ps_norm / (1.0 - (1.0 - c_sigma).powi(2 * generation)).sqrt()
                < (1.4 + 2.0 / (nf + 1.0)) * e_n
            {
                1.0
            } else {
                0.0
            };

            // Covariance evolution path.
            let cc_factor = (c_c * (2.0 - c_c) * mu_eff).sqrt();
            for k in 0..n {
                p_c[k] = (1.0 - c_c) * p_c[k] + hsig * cc_factor * y_w[k];
            }

            // Rank-one + rank-μ covariance update.
            let delta_hsig = (1.0 - hsig) * c_c * (2.0 - c_c);
            for a in 0..n {
                for bcol in a..n {
                    let mut rank_mu = 0.0;
                    for (i, wi) in w.iter().enumerate() {
                        rank_mu += wi * pop[i].1[a] * pop[i].1[bcol];
                    }
                    let rank_one = p_c[a] * p_c[bcol];
                    let val = (1.0 - c_1 - c_mu) * cov[a][bcol]
                        + c_1 * (rank_one + delta_hsig * cov[a][bcol])
                        + c_mu * rank_mu;
                    cov[a][bcol] = val;
                    cov[bcol][a] = val; // keep symmetric
                }
            }

            // Step-size update.
            sigma *= ((c_sigma / d_sigma) * (ps_norm / e_n - 1.0)).exp();
            if !sigma.is_finite() || sigma < 1e-300 {
                break; // converged / degenerate
            }
        }

        let stop = term
            .reason(evaluations, best.value)
            .unwrap_or(StopReason::BudgetExhausted);
        Report {
            solution: best,
            stop,
            evaluations,
        }
    }
}

/// Maps a normalized point in `[0, 1]^n` to the box (no clamping).
fn denormalize(u: &[f64], bounds: &[(f64, f64)]) -> Vec<f64> {
    u.iter()
        .zip(bounds)
        .map(|(&ui, &(lo, hi))| lo + ui * (hi - lo))
        .collect()
}

/// Maps a normalized point to the box, clamping into `[lo, hi]`.
fn denormalize_clamped(u: &[f64], bounds: &[(f64, f64)]) -> Vec<f64> {
    u.iter()
        .zip(bounds)
        .map(|(&ui, &(lo, hi))| lo + ui.clamp(0.0, 1.0) * (hi - lo))
        .collect()
}

fn identity(n: usize) -> Vec<Vec<f64>> {
    let mut m = vec![vec![0.0; n]; n];
    for (i, row) in m.iter_mut().enumerate() {
        row[i] = 1.0;
    }
    m
}

/// Matrix-vector product `M v`.
fn matvec(m: &[Vec<f64>], v: &[f64]) -> Vec<f64> {
    m.iter()
        .map(|row| row.iter().zip(v).map(|(a, b)| a * b).sum())
        .collect()
}

fn norm(v: &[f64]) -> f64 {
    v.iter().map(|x| x * x).sum::<f64>().sqrt()
}

/// Computes `C^{-1/2} v = B D^{-1} Bᵀ v`, where columns of `B` are eigenvectors
/// and `D` holds the square roots of the eigenvalues.
fn c_inv_sqrt_mul(b: &[Vec<f64>], d: &[f64], v: &[f64]) -> Vec<f64> {
    let n = v.len();
    // a = Bᵀ v
    let mut a = vec![0.0; n];
    for (j, aj) in a.iter_mut().enumerate() {
        for i in 0..n {
            *aj += b[i][j] * v[i];
        }
        *aj /= d[j];
    }
    // result = B a
    matvec(b, &a)
}

/// Symmetric eigendecomposition by the cyclic Jacobi method. Returns the
/// eigenvalues and an orthonormal matrix whose **columns** are the matching
/// eigenvectors. Deterministic; intended for the small-to-moderate `n` typical
/// of optimization problems.
fn jacobi_eigen(input: &[Vec<f64>]) -> (Vec<f64>, Vec<Vec<f64>>) {
    let n = input.len();
    let mut a: Vec<Vec<f64>> = input.to_vec();
    let mut v = identity(n);
    if n == 1 {
        return (vec![a[0][0]], v);
    }

    for _ in 0..100 {
        // Off-diagonal magnitude; stop once negligible.
        let mut off = 0.0;
        for p in 0..n {
            for q in p + 1..n {
                off += a[p][q] * a[p][q];
            }
        }
        if off.sqrt() < 1e-14 {
            break;
        }

        for p in 0..n {
            for q in p + 1..n {
                if a[p][q].abs() < 1e-300 {
                    continue;
                }
                // Jacobi rotation angle that zeros a[p][q].
                let theta = (a[q][q] - a[p][p]) / (2.0 * a[p][q]);
                let t = theta.signum() / (theta.abs() + (theta * theta + 1.0).sqrt());
                let c = 1.0 / (t * t + 1.0).sqrt();
                let s = t * c;

                // Rotate rows/columns p and q.
                for k in 0..n {
                    let akp = a[k][p];
                    let akq = a[k][q];
                    a[k][p] = c * akp - s * akq;
                    a[k][q] = s * akp + c * akq;
                }
                for k in 0..n {
                    let apk = a[p][k];
                    let aqk = a[q][k];
                    a[p][k] = c * apk - s * aqk;
                    a[q][k] = s * apk + c * aqk;
                }
                // Accumulate the rotation into the eigenvector matrix.
                for k in 0..n {
                    let vkp = v[k][p];
                    let vkq = v[k][q];
                    v[k][p] = c * vkp - s * vkq;
                    v[k][q] = s * vkp + c * vkq;
                }
            }
        }
    }

    let eigvals = (0..n).map(|i| a[i][i]).collect();
    (eigvals, v)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn jacobi_recovers_known_eigenpairs() {
        // [[2,1],[1,2]] has eigenvalues 1 and 3.
        let (vals, vecs) = jacobi_eigen(&[vec![2.0, 1.0], vec![1.0, 2.0]]);
        let mut sorted = vals.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        assert!((sorted[0] - 1.0).abs() < 1e-9 && (sorted[1] - 3.0).abs() < 1e-9);
        // Eigenvectors are orthonormal: columns have unit norm.
        for j in 0..2 {
            let col_norm = (0..2).map(|i| vecs[i][j] * vecs[i][j]).sum::<f64>().sqrt();
            assert!((col_norm - 1.0).abs() < 1e-9);
        }
    }
}