metalforge 0.1.0

//! Particle Swarm Optimization (Kennedy & Eberhart 1995; Clerc & Kennedy 2002).
//!
//! A swarm of particles flies through the search space, each pulled toward its
//! own best-seen position (*cognitive*) and the swarm's best-seen position
//! (*social*), with momentum (*inertia*). The default coefficients are the
//! constriction-equivalent inertia weights `w = 0.7298`, `c1 = c2 = 1.49618`
//! (Clerc & Kennedy 2002), which keep the swarm convergent without explicit
//! velocity damping.
//!
//! The update is *asynchronous*: each particle reads the current global best at
//! the moment it moves, so improvements propagate within an iteration. Velocity
//! is clamped to one full variable range per step, and a particle that hits a
//! bound has its position clamped and that velocity component zeroed, so it does
//! not stick to the wall at speed.

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

/// PSO configuration.
#[derive(Debug, Clone, Copy)]
pub struct Pso {
    /// Number of particles (≥ 2).
    pub swarm_size: usize,
    /// Inertia weight `w`: momentum carried from the previous velocity.
    pub inertia: f64,
    /// Cognitive coefficient `c1`: pull toward the particle's personal best.
    pub cognitive: f64,
    /// Social coefficient `c2`: pull toward the swarm's global best.
    pub social: f64,
    /// RNG seed; same seed + same problem + same budget ⇒ same result.
    pub seed: u64,
}

impl Default for Pso {
    fn default() -> Self {
        Pso {
            swarm_size: 30,
            inertia: 0.7298,
            cognitive: 1.49618,
            social: 1.49618,
            seed: 42,
        }
    }
}

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

    /// Minimizes `problem` within its bounds using PSO.
    fn optimize(&self, problem: &dyn Problem, term: &Termination) -> Report {
        let bounds = problem.bounds();
        let dim = bounds.len();
        let n = self.swarm_size.max(2);
        let mut rng = Rng::new(self.seed);

        // Per-dimension velocity clamp: one full range per step.
        let vmax: Vec<f64> = bounds.iter().map(|&(lo, hi)| hi - lo).collect();

        // Initialize positions and velocities, seeding the evaluator with the
        // first particle.
        let first_x = sample(bounds, &mut rng);
        let first_v = finite_or_worst(problem.objective(&first_x));
        let mut ev = Evaluator::new(
            problem,
            term,
            Solution {
                x: first_x.clone(),
                value: first_v,
            },
        );

        let mut pos: Vec<Vec<f64>> = Vec::with_capacity(n);
        let mut vel: Vec<Vec<f64>> = Vec::with_capacity(n);
        let mut pbest_x: Vec<Vec<f64>> = Vec::with_capacity(n);
        let mut pbest_f: Vec<f64> = Vec::with_capacity(n);

        // First particle (already evaluated).
        vel.push(init_velocity(&first_x, bounds, &mut rng));
        pbest_x.push(first_x.clone());
        pbest_f.push(first_v);
        pos.push(first_x);

        for _ in 1..n {
            let x = sample(bounds, &mut rng);
            let v = init_velocity(&x, bounds, &mut rng);
            let f = if ev.done() {
                f64::INFINITY
            } else {
                finite_or_worst(ev.eval(&x))
            };
            pbest_x.push(x.clone());
            pbest_f.push(f);
            pos.push(x);
            vel.push(v);
        }

        while !ev.done() {
            for i in 0..n {
                if ev.done() {
                    break;
                }
                // Asynchronous: read the current global best for this move.
                let gbest = ev.best.x.clone();
                for d in 0..dim {
                    let (lo, hi) = bounds[d];
                    let r1 = rng.uniform();
                    let r2 = rng.uniform();
                    let mut v = self.inertia * vel[i][d]
                        + self.cognitive * r1 * (pbest_x[i][d] - pos[i][d])
                        + self.social * r2 * (gbest[d] - pos[i][d]);
                    v = clamp(v, -vmax[d], vmax[d]);

                    let mut x = pos[i][d] + v;
                    if x < lo {
                        x = lo;
                        v = 0.0;
                    } else if x > hi {
                        x = hi;
                        v = 0.0;
                    }
                    vel[i][d] = v;
                    pos[i][d] = x;
                }

                let f = finite_or_worst(ev.eval(&pos[i]));
                if f < pbest_f[i] {
                    pbest_f[i] = f;
                    pbest_x[i].copy_from_slice(&pos[i]);
                }
            }
        }

        ev.finish()
    }
}

#[inline]
fn finite_or_worst(v: f64) -> f64 {
    if v.is_finite() {
        v
    } else {
        f64::INFINITY
    }
}

/// Initial velocity in the SPSO style: half the gap to a fresh random point,
/// `v[d] = ½(U(lo, hi) − x[d])`. Keeps early motion on the scale of the box.
fn init_velocity(x: &[f64], bounds: &[(f64, f64)], rng: &mut Rng) -> Vec<f64> {
    bounds
        .iter()
        .zip(x)
        .map(|(&(lo, hi), &xi)| 0.5 * (rng.uniform_in(lo, hi) - xi))
        .collect()
}