swarmkit 0.1.0

//! Global-best PSO: one swarm-wide attractor.

use crate::searcher_impl::{SearcherCore, evaluate_pbests, move_particles, reduce_best};
use crate::{Best, Contextful, FitCalc, Group, ParticleMover, ParticleRefMut, Searcher, Unit};
use rand::distr::StandardUniform;
use rand::{Rng, RngExt as _};
use std::marker::PhantomData;
use std::ops::{Add, Mul, Sub};

/// Coefficients of the standard PSO velocity update:
/// `v ← inertia·v + cognitive·r1·(pbest − pos) + social·r2·(gbest − pos)`.
#[derive(Copy, Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct PSOCoeffs {
    /// PSO `w`.
    pub inertia: f64,
    /// PSO `c1`.
    pub cognitive_coeff: f64,
    /// PSO `c2`.
    pub social_coeff: f64,
}

impl Default for PSOCoeffs {
    /// Textbook defaults — `#[derive(Default)]` would give an all-zero,
    /// inert searcher, which is worse than useless as a starting point.
    fn default() -> Self {
        Self::new(0.5, 1.5, 1.5)
    }
}

impl PSOCoeffs {
    /// Panics if any coefficient is non-finite or negative.
    ///
    /// ```
    /// use swarmkit::PSOCoeffs;
    ///
    /// let coeffs = PSOCoeffs::new(0.5, 1.5, 1.5);
    /// assert_eq!(coeffs.inertia, 0.5);
    /// ```
    ///
    /// Non-finite or negative values panic:
    ///
    /// ```should_panic
    /// # use swarmkit::PSOCoeffs;
    /// PSOCoeffs::new(0.5, f64::NAN, 1.5);
    /// ```
    #[must_use]
    #[track_caller]
    pub fn new(inertia: f64, cognitive_coeff: f64, social_coeff: f64) -> Self {
        assert!(
            inertia.is_finite() && inertia >= 0.0,
            "inertia must be finite and non-negative, got {inertia}",
        );
        assert!(
            cognitive_coeff.is_finite() && cognitive_coeff >= 0.0,
            "cognitive_coeff must be finite and non-negative, got {cognitive_coeff}",
        );
        assert!(
            social_coeff.is_finite() && social_coeff >= 0.0,
            "social_coeff must be finite and non-negative, got {social_coeff}",
        );
        Self {
            inertia,
            cognitive_coeff,
            social_coeff,
        }
    }
}

/// Global-best PSO searcher.
#[derive(Clone, Debug)]
pub struct GBestSearcher<TUnit, TFit, TMover>
where
    TUnit: Unit,
    TFit: FitCalc,
    TMover: ParticleMover,
{
    core: SearcherCore<TUnit, TFit, TMover>,
}

impl<TUnit, TFit, TMover> GBestSearcher<TUnit, TFit, TMover>
where
    TUnit: Unit,
    TFit: FitCalc<T = TUnit>,
    TMover: ParticleMover<TUnit = TUnit>,
{
    /// RNG is OS-seeded; call [`Searcher::reseed`] for determinism.
    #[must_use]
    pub fn new(fit_calc: TFit, mover: TMover) -> Self {
        Self {
            core: SearcherCore::new(fit_calc, mover),
        }
    }
}

/// Method-chain wrapper: `mover.into_gbest_searcher(fit_calc)`.
pub trait IntoGBestSearcher: ParticleMover + Sized {
    /// Wrap `self` and `fit_calc` into a [`GBestSearcher`].
    #[must_use]
    fn into_gbest_searcher<TFit>(self, fit_calc: TFit) -> GBestSearcher<Self::TUnit, TFit, Self>
    where
        TFit: FitCalc<T = Self::TUnit>,
    {
        GBestSearcher::new(fit_calc, self)
    }
}

impl<TUnit, M> IntoGBestSearcher for M
where
    TUnit: Unit,
    M: ParticleMover<TUnit = TUnit, TCommon = Best<TUnit>>,
{
}

/// Standard PSO velocity update. `TContext` is pass-through only — the
/// mover ignores it, but it lets the type line up when chained with
/// contextful neighbours.
#[derive(Copy, Clone, Debug)]
pub struct GBestMover<TUnit, TContext = ()>
where
    TUnit: Unit,
    TContext: Copy,
{
    phantom: PhantomData<fn() -> (TUnit, TContext)>,
    config: PSOCoeffs,
}

impl<TUnit, TContext> GBestMover<TUnit, TContext>
where
    TUnit: Unit,
    TContext: Copy,
{
    /// Build with `config`.
    #[must_use]
    pub fn new(config: PSOCoeffs) -> Self {
        Self {
            phantom: PhantomData,
            config,
        }
    }
}

impl<TUnit, TContext> Contextful for GBestMover<TUnit, TContext>
where
    TUnit: Unit,
    TContext: Copy,
{
    type TContext = TContext;
}

impl<TUnit, TContext> ParticleMover for GBestMover<TUnit, TContext>
where
    TUnit: Unit + Add<Output = TUnit> + Sub<Output = TUnit> + Mul<f64, Output = TUnit>,
    TContext: Copy,
{
    type TUnit = TUnit;
    type TCommon = Best<TUnit>;

    fn update<R: Rng>(
        &self,
        common: &Self::TCommon,
        rng: &mut R,
        _idx: usize,
        p: &mut ParticleRefMut<'_, Self::TUnit>,
    ) {
        let p_diff = (*p.best_pos - *p.pos) * self.config.cognitive_coeff;
        let g_diff = (common.best_pos - *p.pos) * self.config.social_coeff;
        let a = *p.vel * self.config.inertia;
        let b = p_diff * rng.sample::<f64, StandardUniform>(StandardUniform);
        let c = g_diff * rng.sample::<f64, StandardUniform>(StandardUniform);
        *p.vel = a + b + c;
        *p.pos = *p.pos + *p.vel;
    }
}

impl<TUnit, TFit, TMover, TContext> Contextful for GBestSearcher<TUnit, TFit, TMover>
where
    TFit: FitCalc<T = TUnit, TContext = TContext>,
    TMover: ParticleMover<TUnit = TUnit, TCommon = Best<TUnit>, TContext = TContext>,
    TUnit: Unit,
    TContext: Copy,
{
    type TContext = TContext;

    fn set_context(&mut self, context: Self::TContext) {
        self.core.set_context(context);
    }

    fn set_iteration(&mut self, iteration: usize, max_iteration: usize) {
        self.core.set_iteration(iteration, max_iteration);
    }
}

impl<TUnit, TFit, TMover, TContext> Searcher for GBestSearcher<TUnit, TFit, TMover>
where
    TFit: FitCalc<T = TUnit, TContext = TContext>,
    TMover: ParticleMover<TUnit = TUnit, TCommon = Best<TUnit>, TContext = TContext>,
    TUnit: Unit,
    TContext: Copy,
{
    type TUnit = TUnit;

    fn init(&mut self, particles: &mut Group<Self::TUnit>) {
        self.core.swarm_best = Best::new();
        self.core.sync_rngs(particles.len());
        evaluate_pbests(&self.core.fit_calc, particles);
        reduce_best(particles, &mut self.core.swarm_best);
    }

    fn next(&mut self, particles: &mut Group<Self::TUnit>) -> &Best<TUnit> {
        self.core.sync_rngs(particles.len());
        move_particles(
            &self.core.mover,
            &self.core.swarm_best,
            particles,
            &mut self.core.particle_rngs,
        );
        evaluate_pbests(&self.core.fit_calc, particles);
        reduce_best(particles, &mut self.core.swarm_best);
        &self.core.swarm_best
    }

    fn reseed<R: Rng>(&mut self, rng: &mut R) {
        self.core.reseed(rng);
    }
}