metalforge 0.3.0

//! 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 {
        crate::problem::validate(problem).unwrap_or_else(|e| panic!("CmaEs: invalid problem: {e}"));
        let run = cma_run(problem, term, self.seed, self.population, self.sigma0, 0);
        let stop = term.reason(run.evaluations, run.best.value).unwrap_or(
            // The loop only exits early when an internal criterion fired.
            StopReason::Converged,
        );
        Report {
            solution: run.best,
            stop,
            evaluations: run.evaluations,
        }
    }
}

/// Restart regime for [`RestartCmaEs`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Restart {
    /// IPOP (Auger & Hansen 2005): each restart doubles the population.
    Ipop,
    /// BIPOP (Hansen 2009): interlaces the IPOP regime with a small-population
    /// regime (random λ between the default and half the large λ, randomly
    /// reduced initial step size), giving both regimes similar evaluation
    /// budgets. The BBOB-2009 reference multimodal strategy.
    Bipop,
}

/// CMA-ES with restarts — the accepted, cheap way to make CMA-ES competitive
/// on multimodal functions (IPOP: Auger & Hansen, CEC 2005; BIPOP: Hansen,
/// GECCO/BBOB 2009). Each inner run stops on the standard internal criteria
/// (TolFunHist / TolX / ConditionCov / σ collapse) and the next run restarts
/// with a fresh sub-stream seed (`mix_seed(seed, run_index)`), so the whole
/// procedure stays deterministic for a given seed.
///
/// The global [`Termination`] budget/target bounds the *total* evaluations
/// across all runs; small-regime BIPOP runs are additionally capped at half
/// the evaluations of the latest large run, per the reference description.
#[derive(Debug, Clone, Copy)]
pub struct RestartCmaEs {
    /// Restart regime.
    pub strategy: Restart,
    /// Initial step size for large-regime runs (normalized [0, 1] space).
    pub sigma0: f64,
    /// Safety cap on the number of restarts (the budget usually binds first).
    pub max_restarts: usize,
    /// RNG seed; same seed + same problem + same budget ⇒ same result.
    pub seed: u64,
}

impl Default for RestartCmaEs {
    fn default() -> Self {
        RestartCmaEs {
            strategy: Restart::Bipop,
            sigma0: 0.3,
            max_restarts: 100,
            seed: 42,
        }
    }
}

impl RestartCmaEs {
    /// IPOP configuration with defaults.
    pub fn ipop() -> Self {
        RestartCmaEs {
            strategy: Restart::Ipop,
            ..RestartCmaEs::default()
        }
    }

    /// BIPOP configuration with defaults.
    pub fn bipop() -> Self {
        RestartCmaEs::default()
    }
}

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

    fn optimize(&self, problem: &dyn Problem, term: &Termination) -> Report {
        crate::problem::validate(problem)
            .unwrap_or_else(|e| panic!("RestartCmaEs: invalid problem: {e}"));
        let n = problem.bounds().len();
        let lambda_def = (4 + (3.0 * (n as f64).ln()) as usize).max(4);

        // Control stream for BIPOP's random small-regime draws, disjoint from
        // every inner run's stream (which use mix_seed(seed, run_index)).
        let mut ctl = Rng::split(self.seed, u64::MAX);

        let mut best = Solution {
            x: vec![0.0; n],
            value: f64::INFINITY,
        };
        let mut evaluations = 0usize;
        let mut lambda_large = lambda_def;
        let mut budget_large = 0usize;
        let mut budget_small = 0usize;
        let mut last_large_run = 0usize;
        let mut any_converged = false;

        for run_idx in 0..self.max_restarts.max(1) as u64 {
            if term.reason(evaluations, best.value).is_some() {
                break;
            }
            // Pick the regime. The first run is a large-regime run at the
            // default λ; afterwards IPOP always doubles, while BIPOP gives
            // the turn to whichever regime has spent less budget.
            let (lambda, sigma0, cap, is_large) = if run_idx == 0 {
                (lambda_def, self.sigma0, None, true)
            } else {
                match self.strategy {
                    Restart::Ipop => {
                        lambda_large *= 2;
                        (lambda_large, self.sigma0, None, true)
                    }
                    Restart::Bipop => {
                        if budget_large <= budget_small {
                            lambda_large *= 2;
                            (lambda_large, self.sigma0, None, true)
                        } else {
                            // λ_s = ⌊λ_def · (λ_large / (2 λ_def))^(u₁²)⌋,
                            // σ_s = σ0 · 10^(−2 u₂)  (Hansen 2009).
                            let u1 = ctl.uniform();
                            let u2 = ctl.uniform();
                            let ratio = lambda_large as f64 / (2.0 * lambda_def as f64);
                            let lam = ((lambda_def as f64) * ratio.max(1.0).powf(u1 * u1)) as usize;
                            let sig = self.sigma0 * 10f64.powf(-2.0 * u2);
                            // Small runs get at most half the latest large
                            // run's evaluations.
                            (
                                lam.max(4),
                                sig,
                                Some((last_large_run / 2).max(lam.max(4))),
                                false,
                            )
                        }
                    }
                }
            };

            let inner_term = Termination {
                max_evaluations: match cap {
                    Some(c) => (evaluations + c).min(term.max_evaluations),
                    None => term.max_evaluations,
                },
                target: term.target,
            };
            let out = cma_run(
                problem,
                &inner_term,
                crate::rng::mix_seed(self.seed, run_idx),
                Some(lambda),
                sigma0,
                evaluations,
            );
            let used = out.evaluations - evaluations;
            evaluations = out.evaluations;
            if is_large {
                budget_large += used;
                last_large_run = used;
            } else {
                budget_small += used;
            }
            if out.best.value < best.value {
                best = out.best;
            }
            any_converged |= out.converged;
            // A run stopped by the small-regime cap counts as restartable;
            // only a global budget/target stop ends the procedure.
            if !out.converged && term.reason(evaluations, best.value).is_some() {
                break;
            }
        }

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

/// Outcome of one CMA-ES run (restart wrappers merge several of these).
pub(crate) struct CmaRunOutcome {
    pub best: Solution,
    /// Global evaluation count after the run (continues from `start_evals`).
    pub evaluations: usize,
    /// True when an internal convergence criterion stopped the run (budget
    /// remained) — the restart trigger.
    pub converged: bool,
}

/// One CMA-ES run counting evaluations from `start_evals` against `term`.
///
/// Stops on the global budget/target, or on the standard internal criteria
/// (Hansen's defaults, simplified): **TolFunHist** (range of the
/// best-of-generation history over `10 + ⌈30n/λ⌉` generations below `1e-12`),
/// **TolX** (`σ·p_c` and `σ·√C_kk` all below `1e-12·σ0`), **ConditionCov**
/// (condition number of `C` above `1e14`), and σ collapse/divergence.
pub(crate) fn cma_run(
    problem: &dyn Problem,
    term: &Termination,
    seed: u64,
    population: Option<usize>,
    sigma0: f64,
    start_evals: usize,
) -> CmaRunOutcome {
    let bounds = problem.bounds();
    let n = bounds.len();
    let mut rng = Rng::new(seed);

    // Strategy parameters (Hansen tutorial).
    let lambda = 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.
    // Canonical form: wᵢ' = ln((λ+1)/2) − ln(i+1)  (Hansen tutorial eq. 49;
    // identical to ln(μ+½) only for even λ).
    let raw: Vec<f64> = (0..mu)
        .map(|i| ((lambda as f64 + 1.0) / 2.0).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 = 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;

    // Internal convergence bookkeeping.
    let hist_len = 10 + (30.0 * nf / lambda as f64).ceil() as usize;
    let mut hist: Vec<f64> = Vec::with_capacity(hist_len + 1);
    let tolx = 1e-12 * sigma0;
    let mut converged = false;

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

    'outer: while term.reason(evaluations, best.value).is_none() {
        // Diagonalize C = B diag(d²) Bᵀ; D = sqrt of eigenvalues. The
        // eigenvalue floor is relative to the largest eigenvalue: an
        // absolute floor could inject huge 1/d factors into C^{-1/2}.
        let (eigvals, b) = jacobi_eigen(&cov);
        let max_eig = eigvals.iter().cloned().fold(f64::MIN_POSITIVE, f64::max);
        let d: Vec<f64> = eigvals
            .iter()
            .map(|&v| v.max(max_eig * 1e-14).sqrt())
            .collect();

        // Sample λ offspring; store their normalized step y. Candidates are
        // clamped into the box, and the *clamped* step (u_c − m)/σ is what
        // feeds every adaptation update, so the fitness of the repaired
        // point is never attributed to a step the search did not take.
        let mut pop: Vec<(f64, Vec<f64>)> = Vec::with_capacity(lambda);
        for _ in 0..lambda {
            if term.reason(evaluations, best.value).is_some() {
                break 'outer; // budget/target hit mid-generation
            }
            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_c: Vec<f64> = (0..n)
                .map(|i| (mean[i] + sigma * y[i]).clamp(0.0, 1.0))
                .collect();
            let x = denormalize(&u_c, 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 };
            }
            let y_eff: Vec<f64> = (0..n).map(|i| (u_c[i] - mean[i]) / sigma).collect();
            pop.push((f, y_eff));
        }

        // 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. The recombined steps are built from
        // in-box points, so the mean stays in [0, 1] without clamping.
        for k in 0..n {
            mean[k] += sigma * y_w[k];
        }

        // 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, with collapse *and* divergence guards (the box
        // is [0, 1]^n, so a step size of 10^6 ranges is degenerate).
        sigma *= ((c_sigma / d_sigma) * (ps_norm / e_n - 1.0)).exp();
        if !(1e-300..=1e6).contains(&sigma) {
            converged = true; // collapsed / degenerate (NaN also fails this)
            break;
        }

        // Internal convergence criteria (restart triggers).
        hist.push(pop[0].0);
        if hist.len() > hist_len {
            hist.remove(0);
            let (lo, hi) = hist
                .iter()
                .fold((f64::INFINITY, f64::NEG_INFINITY), |(l, h), &v| {
                    (l.min(v), h.max(v))
                });
            if hi - lo < 1e-12 {
                converged = true; // TolFunHist
                break;
            }
        }
        let tolx_met = (0..n)
            .all(|k| (sigma * p_c[k]).abs() < tolx && sigma * cov[k][k].max(0.0).sqrt() < tolx);
        if tolx_met {
            converged = true; // TolX
            break;
        }
        let min_eig = eigvals.iter().cloned().fold(f64::INFINITY, f64::min);
        if max_eig / min_eig.max(f64::MIN_POSITIVE) > 1e14 {
            converged = true; // ConditionCov
            break;
        }
    }

    CmaRunOutcome {
        best,
        evaluations,
        converged,
    }
}

/// 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()
}

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 *relative to the
        // diagonal scale* (an absolute tolerance would be too strict for large
        // entries and vacuous for tiny ones).
        let mut off = 0.0;
        for p in 0..n {
            for q in p + 1..n {
                off += a[p][q] * a[p][q];
            }
        }
        let scale = (0..n)
            .map(|i| a[i][i].abs())
            .fold(f64::MIN_POSITIVE, f64::max);
        if off.sqrt() < 1e-14 * scale {
            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);
        }
    }
}