swarmkit 0.1.0

//! [`ParticleMover`] trait and the standard combinators:
//! [`BoundedMover`], [`ChainedMover`], [`AdapterMover`], [`NestedMover`].

use crate::{
    Best, Boundary, Contextful, ParticleInitDependent, ParticleRefFrom, ParticleRefMut, Searcher,
    SearcherIter, SetTo, Unit,
};
use rand::Rng;
use std::marker::PhantomData;

/// Advances one particle per step.
///
/// `TCommon` is the per-iteration swarm-level value the mover reads
/// (`Best<TUnit>` for gbest, a per-particle lbest for lbest, a per-niche
/// best for niched). Movers don't touch `fit` / `best_fit` — the
/// searcher's post-move `evaluate_pbests` does.
pub trait ParticleMover: Contextful + Sync {
    /// The unit this mover operates on.
    type TUnit: Unit;
    /// Swarm-level value the mover reads each step.
    type TCommon;

    /// Parallel mover pass stops splitting at this size. Default
    /// `usize::MAX` runs serially; heavy movers (e.g. [`NestedMover`])
    /// override to a small value.
    const PAR_LEAF_SIZE: usize = usize::MAX;

    /// `idx` is the absolute swarm index; topology-aware movers use it
    /// to look up per-particle context.
    fn update<R: Rng>(
        &self,
        common: &Self::TCommon,
        rng: &mut R,
        idx: usize,
        p: &mut ParticleRefMut<'_, Self::TUnit>,
    );

    /// Sequence into a [`ChainedMover`].
    #[must_use]
    fn chain<Other>(self, other: Other) -> ChainedMover<Self, Other>
    where
        Self: Sized,
        Other:
            ParticleMover<TUnit = Self::TUnit, TCommon = Self::TCommon, TContext = Self::TContext>,
    {
        ChainedMover::new(self, other)
    }

    /// Wrap with a [`BoundedMover`] so positions are clamped after each
    /// update.
    #[must_use]
    fn bounded_by<TBoundary>(self, boundary: TBoundary) -> BoundedMover<Self, TBoundary>
    where
        Self: Sized,
        TBoundary: Boundary<T = Self::TUnit, TContext = Self::TContext>,
    {
        BoundedMover::new(self, boundary)
    }

    /// Lift a sub-unit mover into a parent-unit [`AdapterMover`].
    #[must_use]
    fn adapt<TParentCommon>(self) -> AdapterMover<Self, TParentCommon>
    where
        Self: Sized,
        Self::TUnit: ParticleRefFrom,
        Self::TCommon: for<'a> From<&'a TParentCommon> + Copy,
    {
        AdapterMover::new(self)
    }
}

/// Wraps a mover so [`Boundary::handle`] runs after each update.
#[derive(Clone, Debug)]
pub struct BoundedMover<TMover, TBoundary> {
    mover: TMover,
    boundary: TBoundary,
}

impl<TMover, TBoundary> BoundedMover<TMover, TBoundary> {
    /// Wrap `mover` with `boundary`.
    #[must_use]
    pub fn new(mover: TMover, boundary: TBoundary) -> Self {
        Self { mover, boundary }
    }
}

impl<TMover, TBoundary> ParticleMover for BoundedMover<TMover, TBoundary>
where
    TBoundary: Boundary,
    TMover: ParticleMover<TUnit = TBoundary::T, TContext = TBoundary::TContext>,
{
    type TUnit = TBoundary::T;
    type TCommon = TMover::TCommon;

    const PAR_LEAF_SIZE: usize = TMover::PAR_LEAF_SIZE;

    fn update<R: Rng>(
        &self,
        common: &Self::TCommon,
        rng: &mut R,
        idx: usize,
        p: &mut ParticleRefMut<'_, Self::TUnit>,
    ) {
        self.mover.update(common, rng, idx, p);
        *p.pos = self.boundary.handle(*p.pos);
    }
}

impl<TMover, TBoundary> Contextful for BoundedMover<TMover, TBoundary>
where
    TBoundary: Boundary,
    TMover: ParticleMover<TUnit = TBoundary::T, TContext = TBoundary::TContext>,
{
    type TContext = TBoundary::TContext;

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

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

/// Runs two movers in sequence on the same particle, sharing the RNG.
#[derive(Clone, Debug)]
pub struct ChainedMover<TU1, TU2> {
    tu1: TU1,
    tu2: TU2,
}

impl<TU1, TU2> ChainedMover<TU1, TU2> {
    /// `tu1` runs first, then `tu2`.
    #[must_use]
    pub fn new(tu1: TU1, tu2: TU2) -> Self {
        Self { tu1, tu2 }
    }
}

impl<TU1, TU2> ParticleMover for ChainedMover<TU1, TU2>
where
    TU1: ParticleMover,
    TU2: ParticleMover<TUnit = TU1::TUnit, TCommon = TU1::TCommon, TContext = TU1::TContext>,
{
    type TUnit = TU1::TUnit;
    type TCommon = TU1::TCommon;

    // Take the smaller leaf size so the heavier mover dictates split granularity.
    const PAR_LEAF_SIZE: usize = if TU1::PAR_LEAF_SIZE < TU2::PAR_LEAF_SIZE {
        TU1::PAR_LEAF_SIZE
    } else {
        TU2::PAR_LEAF_SIZE
    };

    fn update<R: Rng>(
        &self,
        common: &Self::TCommon,
        rng: &mut R,
        idx: usize,
        p: &mut ParticleRefMut<'_, Self::TUnit>,
    ) {
        self.tu1.update(common, rng, idx, p);
        self.tu2.update(common, rng, idx, p);
    }
}

impl<TU1, TU2> Contextful for ChainedMover<TU1, TU2>
where
    TU1: ParticleMover,
    TU2: ParticleMover<TUnit = TU1::TUnit, TCommon = TU1::TCommon, TContext = TU1::TContext>,
{
    type TContext = TU1::TContext;

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

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

/// Lifts a sub-unit mover into a parent-unit mover via [`ParticleRefFrom`]
/// (for the particle view) and `From<&TParentCommon>` (for the
/// per-iteration value).
#[derive(Clone, Debug)]
pub struct AdapterMover<TMover, TCommon>
where
    TMover: ParticleMover,
    TMover::TUnit: ParticleRefFrom,
    TMover::TCommon: for<'a> From<&'a TCommon> + Copy,
{
    mover: TMover,
    phantom: PhantomData<fn() -> TCommon>,
}

impl<TMover, TCommon> AdapterMover<TMover, TCommon>
where
    TMover: ParticleMover,
    TMover::TUnit: ParticleRefFrom,
    TMover::TCommon: for<'a> From<&'a TCommon> + Copy,
{
    /// Wrap `mover`.
    #[must_use]
    pub fn new(mover: TMover) -> Self {
        Self {
            mover,
            phantom: PhantomData,
        }
    }
}

impl<TMover, TCommon> ParticleMover for AdapterMover<TMover, TCommon>
where
    TMover: ParticleMover,
    TMover::TUnit: ParticleRefFrom,
    TMover::TCommon: for<'a> From<&'a TCommon> + Copy,
{
    type TUnit = <TMover::TUnit as ParticleRefFrom>::TSource;
    type TCommon = TCommon;

    const PAR_LEAF_SIZE: usize = TMover::PAR_LEAF_SIZE;

    fn update<R: Rng>(
        &self,
        common: &Self::TCommon,
        rng: &mut R,
        idx: usize,
        p: &mut ParticleRefMut<'_, Self::TUnit>,
    ) {
        self.mover.update(
            &<TMover::TCommon>::from(common),
            rng,
            idx,
            &mut <TMover::TUnit>::divide_from(p),
        );
    }
}

impl<TMover, TCommon> Contextful for AdapterMover<TMover, TCommon>
where
    TMover: ParticleMover,
    TMover::TUnit: ParticleRefFrom,
    TMover::TCommon: for<'a> From<&'a TCommon> + Copy,
{
    type TContext = TMover::TContext;

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

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

/// Runs an inner search inside one outer mover step.
///
/// Each `update`: clone+reseed the inner searcher, set its context to the
/// outer particle's position, init a fresh inner group, iterate to
/// `max_iteration`, write the inner best back via [`SetTo`]. Per-particle
/// parallelism (`PAR_LEAF_SIZE = 1`) since each outer particle runs a
/// whole inner PSO.
#[derive(Clone, Debug)]
pub struct NestedMover<TOuter, TInner, TSearcher, TInit>
where
    TOuter: Unit,
    TInner: Unit + SetTo<TTarget = TOuter>,
    TSearcher: Searcher<TUnit = TInner>,
    TInit: ParticleInitDependent<TIn = TOuter, TOut = TInner>,
{
    searcher: TSearcher,
    max_iteration: usize,
    particle_init: TInit,
}

impl<TOuter, TInner, TSearcher, TInit> NestedMover<TOuter, TInner, TSearcher, TInit>
where
    TOuter: Unit,
    TInner: Unit + SetTo<TTarget = TOuter>,
    TSearcher: Searcher<TUnit = TInner>,
    TInit: ParticleInitDependent<TIn = TOuter, TOut = TInner>,
{
    /// Panics if `max_iteration == 0`.
    #[must_use]
    #[track_caller]
    pub fn new(max_iteration: usize, searcher: TSearcher, particle_init: TInit) -> Self {
        assert!(
            max_iteration >= 1,
            "NestedMover requires max_iteration >= 1",
        );
        Self {
            searcher,
            max_iteration,
            particle_init,
        }
    }
}

impl<TOuter, TInner, TSearcher, TInit> ParticleMover
    for NestedMover<TOuter, TInner, TSearcher, TInit>
where
    TOuter: Unit,
    TInner: Unit + SetTo<TTarget = TOuter>,
    TSearcher: Searcher<TUnit = TInner, TContext = TOuter> + Clone + Sync,
    TInit: ParticleInitDependent<TIn = TOuter, TOut = TInner> + Sync,
{
    type TUnit = TOuter;
    type TCommon = Best<TOuter>;

    const PAR_LEAF_SIZE: usize = 1;

    fn update<R: Rng>(
        &self,
        _common: &Self::TCommon,
        rng: &mut R,
        _idx: usize,
        p: &mut ParticleRefMut<'_, Self::TUnit>,
    ) {
        let mut searcher = self.searcher.clone();
        searcher.reseed(rng);
        searcher.set_context(*p.pos);

        let mut group = self.particle_init.init_dep(rng, *p.pos);

        let iter: SearcherIter<'_, TInner, TSearcher> =
            SearcherIter::new(self.max_iteration, &mut searcher, &mut group, None);
        let best = iter
            .last()
            .expect("max_iteration >= 1 is enforced by NestedMover::new");

        best.best_pos.set_to_ref_mut(p);
    }
}

impl<TOuter, TInner, TSearcher, TInit> Contextful for NestedMover<TOuter, TInner, TSearcher, TInit>
where
    TOuter: Unit,
    TInner: Unit + SetTo<TTarget = TOuter>,
    TSearcher: Searcher<TUnit = TInner, TContext = TOuter>,
    TInit: ParticleInitDependent<TIn = TOuter, TOut = TInner>,
{
    type TContext = TOuter;

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