metalforge 0.1.1

//! Standard optimization test functions, as [`Problem`]s.
//!
//! These are the classic benchmarks used to validate global optimizers: the
//! convex unimodal **sphere**, the narrow-valley **Rosenbrock**, and the highly
//! multimodal **Rastrigin**. All have a known global minimum of `0`, which the
//! integration tests check each algorithm can reach.

use crate::problem::{Bound, MultiProblem, Problem};

/// Sphere: `f(x) = Σ xᵢ²`. Convex, unimodal, minimum `0` at the origin.
pub struct Sphere {
    bounds: Vec<Bound>,
}

impl Sphere {
    /// `dim`-dimensional sphere on `[-5.12, 5.12]^dim`.
    pub fn new(dim: usize) -> Self {
        Sphere {
            bounds: vec![(-5.12, 5.12); dim],
        }
    }
}

impl Problem for Sphere {
    fn dim(&self) -> usize {
        self.bounds.len()
    }
    fn bounds(&self) -> &[Bound] {
        &self.bounds
    }
    fn objective(&self, x: &[f64]) -> f64 {
        x.iter().map(|v| v * v).sum()
    }
}

/// Rosenbrock: `f(x) = Σ [100(x_{i+1} − xᵢ²)² + (1 − xᵢ)²]`. A long, curved,
/// nearly-flat valley; minimum `0` at the all-ones vector.
pub struct Rosenbrock {
    bounds: Vec<Bound>,
}

impl Rosenbrock {
    /// `dim`-dimensional Rosenbrock (`dim >= 2`) on `[-5, 10]^dim`.
    pub fn new(dim: usize) -> Self {
        assert!(dim >= 2, "Rosenbrock needs at least 2 dimensions");
        Rosenbrock {
            bounds: vec![(-5.0, 10.0); dim],
        }
    }
}

impl Problem for Rosenbrock {
    fn dim(&self) -> usize {
        self.bounds.len()
    }
    fn bounds(&self) -> &[Bound] {
        &self.bounds
    }
    fn objective(&self, x: &[f64]) -> f64 {
        x.windows(2)
            .map(|w| 100.0 * (w[1] - w[0] * w[0]).powi(2) + (1.0 - w[0]).powi(2))
            .sum()
    }
}

/// Rastrigin: `f(x) = 10n + Σ [xᵢ² − 10 cos(2π xᵢ)]`. Highly multimodal with a
/// regular lattice of local minima; global minimum `0` at the origin.
pub struct Rastrigin {
    bounds: Vec<Bound>,
}

impl Rastrigin {
    /// `dim`-dimensional Rastrigin on `[-5.12, 5.12]^dim`.
    pub fn new(dim: usize) -> Self {
        Rastrigin {
            bounds: vec![(-5.12, 5.12); dim],
        }
    }
}

impl Problem for Rastrigin {
    fn dim(&self) -> usize {
        self.bounds.len()
    }
    fn bounds(&self) -> &[Bound] {
        &self.bounds
    }
    fn objective(&self, x: &[f64]) -> f64 {
        let n = x.len() as f64;
        10.0 * n
            + x.iter()
                .map(|v| v * v - 10.0 * (2.0 * std::f64::consts::PI * v).cos())
                .sum::<f64>()
    }
}

/// ZDT1 (Zitzler, Deb & Thiele 2000): a two-objective benchmark with a known,
/// convex Pareto front. On `[0, 1]^n`,
///
/// ```text
/// f1 = x1
/// g  = 1 + 9 · (Σ_{i≥2} xi) / (n − 1)
/// f2 = g · (1 − √(f1 / g))
/// ```
///
/// The Pareto-optimal set has `x2 … xn = 0` (so `g = 1`), giving the analytical
/// front `f2 = 1 − √f1` for `f1 ∈ [0, 1]` — the ground truth the NSGA-II test
/// measures distance to.
pub struct Zdt1 {
    bounds: Vec<Bound>,
}

impl Zdt1 {
    /// `dim`-dimensional ZDT1 (`dim >= 2`) on `[0, 1]^dim`.
    pub fn new(dim: usize) -> Self {
        assert!(dim >= 2, "ZDT1 needs at least 2 dimensions");
        Zdt1 {
            bounds: vec![(0.0, 1.0); dim],
        }
    }

    /// The analytical Pareto front value `f2 = 1 − √f1` for a given `f1`.
    pub fn front_f2(f1: f64) -> f64 {
        1.0 - f1.sqrt()
    }
}

impl MultiProblem for Zdt1 {
    fn dim(&self) -> usize {
        self.bounds.len()
    }
    fn bounds(&self) -> &[Bound] {
        &self.bounds
    }
    fn n_objectives(&self) -> usize {
        2
    }
    fn objectives(&self, x: &[f64]) -> Vec<f64> {
        let n = x.len();
        let f1 = x[0];
        let g = 1.0 + 9.0 * x[1..].iter().sum::<f64>() / (n - 1) as f64;
        let f2 = g * (1.0 - (f1 / g).sqrt());
        vec![f1, f2]
    }
}

/// DTLZ2 (Deb, Thiele, Laumanns & Zitzler 2002): a scalable many-objective
/// benchmark whose Pareto front is the unit hypersphere in the positive orthant
/// (`Σ fᵢ² = 1`). On `[0, 1]^n` with `M` objectives and `k = n − M + 1` distance
/// variables,
///
/// ```text
/// g    = Σ_{i ≥ M} (xᵢ − 0.5)²
/// f_1  = (1+g) · cos(x₁ π/2) ··· cos(x_{M−1} π/2)
/// f_j  = (1+g) · cos(x₁ π/2) ··· sin(x_{M−j+1} π/2)
/// f_M  = (1+g) · sin(x₁ π/2)
/// ```
///
/// The front is reached when `g = 0` (all distance variables at 0.5); there the
/// objective vector lies exactly on the unit sphere — the ground truth the
/// NSGA-III test measures distance to.
pub struct Dtlz2 {
    bounds: Vec<Bound>,
    m: usize,
}

impl Dtlz2 {
    /// `m`-objective DTLZ2 in `dim` variables (`dim > m`) on `[0, 1]^dim`. The
    /// recommended `dim` is `m + 9` (so `k = 10` distance variables).
    pub fn new(m: usize, dim: usize) -> Self {
        assert!(m >= 2 && dim > m, "DTLZ2 needs m >= 2 and dim > m");
        Dtlz2 {
            bounds: vec![(0.0, 1.0); dim],
            m,
        }
    }
}

impl MultiProblem for Dtlz2 {
    fn dim(&self) -> usize {
        self.bounds.len()
    }
    fn bounds(&self) -> &[Bound] {
        &self.bounds
    }
    fn n_objectives(&self) -> usize {
        self.m
    }
    fn objectives(&self, x: &[f64]) -> Vec<f64> {
        let m = self.m;
        let g: f64 = x[m - 1..].iter().map(|&xi| (xi - 0.5).powi(2)).sum();
        let half_pi = std::f64::consts::FRAC_PI_2;
        let mut f = vec![1.0 + g; m];
        for (i, fi) in f.iter_mut().enumerate() {
            // Product of cosines for the first (M-1-i) angles...
            for &xj in x.iter().take(m - 1 - i) {
                *fi *= (xj * half_pi).cos();
            }
            // ...and one sine, except for the first objective.
            if i > 0 {
                *fi *= (x[m - 1 - i] * half_pi).sin();
            }
        }
        f
    }
}

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

    #[test]
    fn known_minima_are_zero() {
        assert_eq!(Sphere::new(3).objective(&[0.0, 0.0, 0.0]), 0.0);
        assert!(Rosenbrock::new(3).objective(&[1.0, 1.0, 1.0]).abs() < 1e-12);
        assert!(Rastrigin::new(4).objective(&[0.0; 4]).abs() < 1e-12);
    }

    #[test]
    fn dtlz2_front_lies_on_the_unit_sphere() {
        // Distance variables at 0.5 ⇒ g = 0 ⇒ Σ fᵢ² = 1 for any angle variables.
        let p = Dtlz2::new(3, 12);
        for angles in [[0.0, 0.0], [0.5, 0.5], [1.0, 0.3]] {
            let mut x = vec![0.5; 12];
            x[0] = angles[0];
            x[1] = angles[1];
            let f = p.objectives(&x);
            let sumsq: f64 = f.iter().map(|v| v * v).sum();
            assert!((sumsq - 1.0).abs() < 1e-9, "Σf² = {sumsq}");
        }
    }

    #[test]
    fn zdt1_pareto_set_lies_on_the_analytical_front() {
        // x2..xn = 0 ⇒ g = 1 ⇒ f2 = 1 − √f1 for any f1 = x1.
        let p = Zdt1::new(5);
        for &f1 in &[0.0, 0.25, 0.5, 1.0] {
            let mut x = vec![0.0; 5];
            x[0] = f1;
            let o = p.objectives(&x);
            assert!((o[0] - f1).abs() < 1e-12);
            assert!((o[1] - Zdt1::front_f2(f1)).abs() < 1e-12);
        }
    }
}