rlevo-evolution 0.3.0

//! Salp Swarm Algorithm.
//!
//! The population splits into a *leader* half and a *follower* chain:
//!
//! - Leaders (indices `0 .. N/2`) update toward the food source `F`
//!   (the best-so-far position), modulated by a time-decayed coefficient
//!   `c1 = 2·exp(−(4·t / T)²)`:
//!
//!   `X_i ← F ± c1 · ((ub − lb)·c2 + lb)` with `c2, c3 ∈ U[0, 1]`,
//!   where the sign follows `c3 ≥ 0.5 → +` and `c3 < 0.5 → −`.
//!
//! - Followers (indices `N/2 .. N`) track the chain:
//!
//!   `X_i ← (X_i + X_{i−1}) / 2`.
//!
//! The follower rule is realized as a parallel stencil — the shifted
//! copy is `concat([last_leader, followers[..−1]])` — rather than a
//! host-side Python-style loop. That keeps the update GPU-friendly and
//! deterministic.
//!
//! # Candor
//!
//! Legacy comparator. The leader update is a thinly-veiled biased
//! random walk toward the best; the follower chain is a stencil average
//! of old positions. Neither mechanism introduces a novel search
//! dynamic not already present in PSO or DE. Ship it for API coverage;
//! prefer CMA-ES or LSHADE when available.
//!
//! # References
//!
//! - Mirjalili, Gandomi, Mirjalili, Saremi, Faris & Mirjalili (2017),
//!   *Salp Swarm Algorithm*.

use std::marker::PhantomData;

use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
use rand::Rng;
use rand::RngExt;

use rlevo_core::bounds::Bounds;
use rlevo_core::config::{self, ConfigError, Validate};

use crate::ops::selection::argmax_host;
use crate::rng::{SeedPurpose, seed_stream};
use crate::strategy::{Strategy, StrategyMetrics};

/// Static configuration for [`SalpSwarm`].
#[derive(Debug, Clone)]
pub struct SalpConfig {
    /// Swarm size; must be `≥ 2` so there is at least one leader and
    /// one follower.
    pub pop_size: usize,
    /// Genome dimensionality.
    pub genome_dim: usize,
    /// Search-space bounds. The leader update pulls from a scaled
    /// uniform over this range.
    pub bounds: Bounds,
    /// Budget pacing the `c1` decay.
    pub max_generations: usize,
}

impl SalpConfig {
    /// Default configuration for a given population size and genome dimensionality.
    #[must_use]
    pub fn default_for(pop_size: usize, genome_dim: usize) -> Self {
        Self {
            pop_size,
            genome_dim,
            bounds: Bounds::new(-5.12, 5.12),
            max_generations: 500,
        }
    }
}

impl Validate for SalpConfig {
    fn validate(&self) -> Result<(), ConfigError> {
        const C: &str = "SalpConfig";
        config::at_least(C, "pop_size", self.pop_size, 2)?;
        config::nonzero(C, "genome_dim", self.genome_dim)?;
        config::at_least(C, "max_generations", self.max_generations, 1)?;
        Ok(())
    }
}

/// Generation state for [`SalpSwarm`].
#[derive(Debug, Clone)]
pub struct SalpState<B: Backend> {
    /// Current positions, shape `(pop_size, D)`.
    pub positions: Tensor<B, 2>,
    /// Host-side fitness cache.
    pub fitness: Vec<f32>,
    /// Food-source (best) genome.
    pub best_genome: Option<Tensor<B, 2>>,
    /// Food-source fitness.
    pub best_fitness: f32,
    /// Generation counter.
    pub generation: usize,
}

/// Salp Swarm Algorithm strategy.
///
/// # Panics
///
/// [`Strategy::init`] panics if `params.pop_size < 2`, since the
/// leader/follower split requires at least one of each.
///
/// [`Strategy::ask`] panics if called a second time without an intervening
/// [`Strategy::tell`] — `state.best_genome` is `None` until the first
/// `tell` populates it. Normal harness-driven usage prevents this: the
/// first `ask` returns initial positions unevaluated, `tell` sets
/// `best_genome`, and every subsequent `ask` finds it populated.
///
/// # Example
///
/// ```no_run
/// use burn::backend::Flex;
/// use rlevo_evolution::algorithms::metaheuristic::salp::{SalpConfig, SalpSwarm};
///
/// let strategy = SalpSwarm::<Flex>::new();
/// let params = SalpConfig::default_for(32, 10);
/// let _ = (strategy, params);
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct SalpSwarm<B: Backend> {
    _backend: PhantomData<fn() -> B>,
}

impl<B: Backend> SalpSwarm<B> {
    /// Builds a new (stateless) strategy object.
    #[must_use]
    pub fn new() -> Self {
        Self {
            _backend: PhantomData,
        }
    }
}

impl<B: Backend> Strategy<B> for SalpSwarm<B>
where
    B::Device: Clone,
{
    type Params = SalpConfig;
    type State = SalpState<B>;
    type Genome = Tensor<B, 2>;

    /// Samples the initial swarm uniformly within [`SalpConfig::bounds`]
    /// using the host-RNG convention and sets the generation counter to zero.
    ///
    /// The `pop_size >= 2` invariant is enforced by [`Validate::validate`] at
    /// the harness chokepoint.
    fn init(
        &self,
        params: &SalpConfig,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> SalpState<B> {
        debug_assert!(
            params.validate().is_ok(),
            "invalid SalpConfig reached init: {params:?}"
        );
        let (lo, hi): (f32, f32) = params.bounds.into();
        // Host-sample the initial swarm from a deterministic `seed_stream`
        // rather than the process-wide Flex RNG (`B::seed` + `Tensor::random`),
        // whose draws interleave with sibling tests under the parallel
        // runner and are therefore not reproducible across thread schedules.
        let pop = params.pop_size;
        let genome_dim = params.genome_dim;
        let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
        let mut position_rows = Vec::with_capacity(pop * genome_dim);
        for _ in 0..pop * genome_dim {
            position_rows.push(lo + (hi - lo) * stream.random::<f32>());
        }
        let positions =
            Tensor::<B, 2>::from_data(TensorData::new(position_rows, [pop, genome_dim]), device);
        SalpState {
            positions,
            fitness: Vec::new(),
            best_genome: None,
            best_fitness: f32::NEG_INFINITY,
            generation: 0,
        }
    }

    /// Proposes the next swarm positions.
    ///
    /// On the first call returns the initial positions unchanged. On
    /// subsequent calls, applies the two-phase update:
    ///
    /// - **Leaders** (`0..pop_size/2`): each leader moves toward the food
    ///   source with a signed, decaying step proportional to `c1` and a
    ///   host-sampled uniform offset over [`SalpConfig::bounds`].
    /// - **Followers** (`pop_size/2..pop_size`): each follower averages its
    ///   current position with the position of the salp directly ahead (a
    ///   parallel stencil over the updated leader block).
    ///
    /// Results are clamped to [`SalpConfig::bounds`].
    fn ask(
        &self,
        params: &SalpConfig,
        state: &SalpState<B>,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> (Tensor<B, 2>, SalpState<B>) {
        if state.fitness.is_empty() {
            return (state.positions.clone(), state.clone());
        }

        let pop_size = params.pop_size;
        let genome_dim = params.genome_dim;
        let n_leaders = pop_size / 2;
        let (lo, hi): (f32, f32) = params.bounds.into();

        // c1 decay: 2 * exp(-(4t/T)^2)
        #[allow(clippy::cast_precision_loss)]
        let t = state.generation as f32;
        #[allow(clippy::cast_precision_loss)]
        let max_t = params.max_generations.max(1) as f32;
        let frac = (4.0 * t / max_t).min(4.0);
        let c1 = 2.0 * (-(frac * frac)).exp();

        // Host-side sampling for c2, c3 so the leader step is
        // backend-agnostic and reproducible under the splitmix contract.
        let mut stream = seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Other);
        let mut leader_delta: Vec<f32> = Vec::with_capacity(n_leaders * genome_dim);
        for _ in 0..n_leaders {
            for _ in 0..genome_dim {
                let c2: f32 = stream.random::<f32>();
                let c3: f32 = stream.random::<f32>();
                let scaled = (hi - lo) * c2 + lo;
                let sign = if c3 >= 0.5 { 1.0 } else { -1.0 };
                leader_delta.push(sign * c1 * scaled);
            }
        }

        let best = state
            .best_genome
            .as_ref()
            .expect("best_genome populated after first tell")
            .clone()
            .expand([n_leaders, genome_dim]);
        let delta = Tensor::<B, 2>::from_data(
            TensorData::new(leader_delta, [n_leaders, genome_dim]),
            device,
        );
        let new_leaders = (best + delta).clamp(lo, hi);

        // Follower stencil: X_i ← (X_i + X_{i-1}) / 2.
        let followers = state
            .positions
            .clone()
            .slice([n_leaders..pop_size, 0..genome_dim]);
        // Shifted copy: follower i reads the position of the salp
        // directly ahead, which is the last leader for the first
        // follower and `followers[i-1]` otherwise. We build the shifted
        // stack by gathering indices `[n_leaders-1, n_leaders, …, pop_size-2]`
        // from the *updated* leader block + the old followers slice.
        let joined = Tensor::cat(vec![new_leaders.clone(), followers.clone()], 0); // (pop_size, D) — leaders are already the new ones.
        #[allow(clippy::cast_possible_wrap)]
        let shift_idx: Vec<i64> = (0..(pop_size - n_leaders))
            .map(|k| (n_leaders + k - 1) as i64)
            .collect();
        let idx = Tensor::<B, 1, Int>::from_data(
            TensorData::new(shift_idx, [pop_size - n_leaders]),
            device,
        );
        let previous = joined.clone().select(0, idx);
        let new_followers = (followers + previous).mul_scalar(0.5).clamp(lo, hi);

        let new_positions = Tensor::cat(vec![new_leaders, new_followers], 0);
        let mut next = state.clone();
        next.positions.clone_from(&new_positions);
        (new_positions, next)
    }

    /// Records evaluated fitness, updates the food-source (best-so-far), and
    /// increments the generation counter.
    ///
    /// Returns the updated [`SalpState`] and a [`StrategyMetrics`] snapshot
    /// for the completed generation.
    fn tell(
        &self,
        _params: &SalpConfig,
        population: Tensor<B, 2>,
        fitness: Tensor<B, 1>,
        mut state: SalpState<B>,
        _rng: &mut dyn Rng,
    ) -> (SalpState<B>, StrategyMetrics) {
        let fitness_host = fitness
            .into_data()
            .into_vec::<f32>()
            .expect("fitness tensor must be readable as f32");
        state.fitness.clone_from(&fitness_host);
        state.positions.clone_from(&population);
        let best_idx = argmax_host(&fitness_host);
        if fitness_host[best_idx] > state.best_fitness {
            state.best_fitness = fitness_host[best_idx];
            let device = population.device();
            #[allow(clippy::cast_possible_wrap)]
            let idx = Tensor::<B, 1, Int>::from_data(
                TensorData::new(vec![best_idx as i64], [1]),
                &device,
            );
            state.best_genome = Some(population.select(0, idx));
        }
        state.generation += 1;
        let m =
            StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
        state.best_fitness = m.best_fitness_ever();
        (state, m)
    }

    /// Returns the food-source genome and its fitness, or `None` before the
    /// first [`tell`](Strategy::tell) call.
    fn best(&self, state: &SalpState<B>) -> Option<(Tensor<B, 2>, f32)> {
        state
            .best_genome
            .as_ref()
            .map(|g| (g.clone(), state.best_fitness))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fitness::FromFitnessEvaluable;
    use crate::strategy::EvolutionaryHarness;
    use burn::backend::Flex;
    use rand::SeedableRng;
    use rand::rngs::StdRng;
    use rlevo_core::fitness::FitnessEvaluable;

    type TestBackend = Flex;

    /// Distinct finite fitness with the maximum at index 0, for direct
    /// (non-harness) `tell` calls. Finite, so it respects ADR 0034.
    #[allow(clippy::trivially_copy_pass_by_ref)] // mirror the by-ref device idiom
    fn finite_fitness(
        n: usize,
        device: &<TestBackend as burn::tensor::backend::BackendTypes>::Device,
    ) -> Tensor<TestBackend, 1> {
        #[allow(clippy::cast_precision_loss)]
        let vals: Vec<f32> = (0..n).map(|i| -(i as f32) - 1.0).collect();
        Tensor::<TestBackend, 1>::from_data(TensorData::new(vals, [n]), device)
    }

    #[test]
    fn default_config_validates() {
        assert!(SalpConfig::default_for(30, 10).validate().is_ok());
    }

    #[test]
    fn rejects_pop_size_below_two() {
        let mut cfg = SalpConfig::default_for(30, 10);
        cfg.pop_size = 1;
        assert_eq!(cfg.validate().unwrap_err().field, "pop_size");
    }

    struct Sphere;
    struct SphereFit;
    impl FitnessEvaluable for SphereFit {
        type Individual = Vec<f64>;
        type Landscape = Sphere;
        fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
            x.iter().map(|v| v * v).sum()
        }
    }

    #[test]
    fn ssa_converges_on_sphere_d10() {
        // SSA reduces strongly but, per module-level candor, is not
        // expected to reach machine precision. A 1e-2 target in 600
        // generations is the strong-reduction guarantee mirroring the
        // DE premature-convergence test pattern.
        let device = Default::default();
        let strategy = SalpSwarm::<TestBackend>::new();
        let params = SalpConfig::default_for(40, 10);
        let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
        let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
            strategy, params, fitness_fn, 3, device, 600,
        )
        .expect("valid params");
        harness.reset();
        while !harness.step(()).done {}
        let best = harness.latest_metrics().unwrap().best_fitness_ever();
        assert!(best < 1e-2, "SSA D10 best={best}");
    }

    #[test]
    fn best_is_none_until_first_tell() {
        let device = Default::default();
        let strategy = SalpSwarm::<TestBackend>::new();
        let params = SalpConfig::default_for(4, 3);
        let mut rng = StdRng::seed_from_u64(0);
        let state = strategy.init(&params, &mut rng, &device);
        assert!(strategy.best(&state).is_none());
        let (pop, state) = strategy.ask(&params, &state, &mut rng, &device);
        let fitness = finite_fitness(4, &device);
        let (state, _m) = strategy.tell(&params, pop, fitness, state, &mut rng);
        assert!(strategy.best(&state).is_some());
    }

    #[test]
    fn first_ask_returns_initial_positions_unchanged() {
        // Before any `tell`, `state.fitness` is empty, so `ask` short-circuits
        // and hands back the untouched initial swarm.
        let device = Default::default();
        let strategy = SalpSwarm::<TestBackend>::new();
        let params = SalpConfig::default_for(6, 4);
        let mut rng = StdRng::seed_from_u64(3);
        let state = strategy.init(&params, &mut rng, &device);
        let expected = state
            .positions
            .clone()
            .into_data()
            .into_vec::<f32>()
            .unwrap();
        let (pop, _state) = strategy.ask(&params, &state, &mut rng, &device);
        let got = pop.into_data().into_vec::<f32>().unwrap();
        assert_eq!(expected, got);
    }

    #[test]
    fn minimal_and_odd_pop_sizes_run() {
        // pop_size = 2 is the leader/follower minimum; pop_size = 3 exercises
        // the odd split (1 leader, 2 followers). Both must complete without a
        // panic in the follower stencil.
        for pop in [2usize, 3] {
            let device = Default::default();
            let strategy = SalpSwarm::<TestBackend>::new();
            let params = SalpConfig::default_for(pop, 3);
            let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
            let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
                strategy, params, fitness_fn, 0, device, 5,
            )
            .expect("valid params");
            harness.reset();
            while !harness.step(()).done {}
            assert!(
                harness
                    .latest_metrics()
                    .unwrap()
                    .best_fitness_ever()
                    .is_finite(),
                "pop_size {pop} produced a non-finite best"
            );
        }
    }

    #[test]
    fn ask_keeps_positions_in_bounds() {
        // Every position proposed by the two-phase update (leader step +
        // follower stencil) is clamped to bounds; verify across seeds.
        let device = Default::default();
        let strategy = SalpSwarm::<TestBackend>::new();
        let params = SalpConfig::default_for(6, 4);
        let (lo, hi): (f32, f32) = params.bounds.into();
        for seed in 0..32 {
            let mut rng = StdRng::seed_from_u64(seed);
            let state = strategy.init(&params, &mut rng, &device);
            let (pop1, state) = strategy.ask(&params, &state, &mut rng, &device);
            let fitness = finite_fitness(6, &device);
            let (state, _m) = strategy.tell(&params, pop1, fitness, state, &mut rng);
            let (pop2, _state) = strategy.ask(&params, &state, &mut rng, &device);
            let values = pop2.into_data().into_vec::<f32>().unwrap();
            for v in values {
                assert!(
                    v >= lo - 1e-4 && v <= hi + 1e-4,
                    "seed {seed}: position {v} out of bounds [{lo}, {hi}]"
                );
            }
        }
    }

    #[test]
    fn zero_budget_harness_produces_no_metrics() {
        // The harness *budget* (distinct from SalpConfig::max_generations) is
        // zero: a well-behaved driver takes no step, so `latest_metrics` stays
        // None after `reset`.
        let device = Default::default();
        let strategy = SalpSwarm::<TestBackend>::new();
        let params = SalpConfig::default_for(4, 3);
        let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
        let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
            strategy, params, fitness_fn, 0, device, 0,
        )
        .expect("valid params");
        harness.reset();
        assert!(harness.latest_metrics().is_none());
    }

    #[test]
    fn unit_budget_harness_runs_exactly_one_generation() {
        let device = Default::default();
        let strategy = SalpSwarm::<TestBackend>::new();
        let params = SalpConfig::default_for(4, 3);
        let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
        let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
            strategy, params, fitness_fn, 0, device, 1,
        )
        .expect("valid params");
        harness.reset();
        assert!(harness.step(()).done);
        assert_eq!(harness.generation(), 1);
        assert!(harness.latest_metrics().is_some());
    }
}