pso/optimizer/

swarm.rs

mod minima;
mod particle;
mod transient_model;

/// Configure behavior of Particle Swarms
pub mod swarm_config;

pub use minima::GlobalRecord;

use super::{Bounds, JobConfig, StopCondition};
use minima::{OwnedRecord, RecordBeatResult};
use particle::Particle;
use rand::rngs::SmallRng;
use rand::SeedableRng;
use std::sync::Arc;
use swarm_config::{GlobalBehavior, SwarmConfig, TransientBehavior};
use transient_model::TransientModel;

const DEFAULT_UPDATE_RATE: usize = 16;

#[derive(Clone, Debug)]
pub struct Swarm {
    name: String,
    particles: Vec<Particle>,
    transient_model: TransientModel,
    bound_state: BoundState,
    stop_state: StopState,
    global_updater: GlobalUpdater,
    synergic: bool,
    update_console: Option<usize>,
    num_variables: usize,
}

#[allow(dead_code)]
impl Swarm {
    pub fn new(name: String, sc: SwarmConfig, jc: &JobConfig) -> Self {
        let particles = Self::initialize_particles(&sc, &jc);

        let motion_coefficients = match sc.global_behavior {
            GlobalBehavior::Synergic(global_coeff, _) => {
                vec![sc.local_motion, sc.tribal_motion, global_coeff]
            }
            GlobalBehavior::Solitary => vec![sc.local_motion, sc.tribal_motion],
        };

        let transient_model = TransientModel::new(
            &sc.transient_behavior,
            &jc.stop_condition,
            jc.num_variables,
            motion_coefficients,
            sc.momentum,
        );

        let bound_state = BoundState::new(&sc, &jc);
        let stop_state = StopState::new(&jc.stop_condition);
        let (global_updater, synergic) = GlobalUpdater::new(&sc.global_behavior);
        let update_console = jc.update_console;
        let num_variables = jc.num_variables;

        Self {
            name,
            particles,
            transient_model,
            bound_state,
            stop_state,
            global_updater,
            synergic,
            update_console,
            num_variables,
        }
    }

    pub fn minimize_fn_mut(
        &mut self,
        mut objective: impl FnMut(&[f64]) -> f64,
        global_record: &mut GlobalRecord,
    ) {
        let mut tribal_record = OwnedRecord::new(self.num_variables);
        let mut owned_global_record = OwnedRecord::new(self.num_variables);

        let mut rng = SmallRng::from_entropy();

        while let Some(iteration) = self.stop_state.next_iteration(&owned_global_record) {
            let m = self.transient_model.momentum(iteration);
            let motion_coeffs = self.transient_model.motion_coeffs(iteration);

            let mut record_beat = false;
            let mut min_cost = tribal_record.cost;
            let mut min_particle_index = 0;

            for (i, p) in self.particles.iter_mut().enumerate() {
                let rand_vecs = self.transient_model.rand_vecs(iteration, &mut rng);

                // test cost function at current location
                if let Some(new_min_cost) = p.cost_fn_mut(&mut objective) {
                    if new_min_cost < min_cost {
                        record_beat = true;
                        min_particle_index = i;
                        min_cost = new_min_cost;
                    }
                }

                // update position and velocity
                if self.synergic {
                    p.update_syn(
                        m,
                        rand_vecs,
                        motion_coeffs,
                        tribal_record.pos(),
                        owned_global_record.pos(),
                    );
                } else {
                    p.update_sol(m, rand_vecs, motion_coeffs, tribal_record.pos());
                }

                // enforce search space and velocity bounds
                p.enforce_bounds(self.bound_state.values());
            }

            // bubble new records up the chain to tribal and global
            if record_beat {
                tribal_record = self.particles[min_particle_index].get_record();
                if tribal_record < owned_global_record {
                    owned_global_record = tribal_record.clone();
                }
            }

            // update or download global record as necessary
            self.global_updater
                .update_global(iteration, &mut owned_global_record, global_record);

            self.print_progress(iteration, &owned_global_record, &record_beat);
        }

        self.global_updater
            .update_global(0, &mut owned_global_record, global_record)
    }

    pub fn minimize_fn<T>(&mut self, objective: Arc<T>, global_record: &mut GlobalRecord)
    where
        T: Fn(&[f64]) -> f64,
    {
        let mut tribal_record = OwnedRecord::new(self.num_variables);
        let mut owned_global_record = OwnedRecord::new(self.num_variables);

        let mut rng = SmallRng::from_entropy();

        while let Some(iteration) = self.stop_state.next_iteration(&owned_global_record) {
            let m = self.transient_model.momentum(iteration);
            let motion_coeffs = self.transient_model.motion_coeffs(iteration);

            let mut record_beat = false;
            let mut min_cost = tribal_record.cost;
            let mut min_particle_index = 0;

            for (i, p) in self.particles.iter_mut().enumerate() {
                let rand_vecs = self.transient_model.rand_vecs(iteration, &mut rng);

                // test cost function at current location
                if let Some(new_min_cost) = p.cost_fn(&objective) {
                    if new_min_cost < min_cost {
                        record_beat = true;
                        min_particle_index = i;
                        min_cost = new_min_cost;
                    }
                }

                // update position and velocity
                if self.synergic {
                    p.update_syn(
                        m,
                        rand_vecs,
                        motion_coeffs,
                        tribal_record.pos(),
                        owned_global_record.pos(),
                    );
                } else {
                    p.update_sol(m, rand_vecs, motion_coeffs, tribal_record.pos());
                }

                // enforce search space and velocity bounds
                p.enforce_bounds(self.bound_state.values());
            }

            // bubble new records up the chain to tribal and global
            if record_beat {
                tribal_record = self.particles[min_particle_index].get_record();
                if tribal_record < owned_global_record {
                    owned_global_record = tribal_record.clone();
                }
            }

            // update or download global record as necessary
            self.global_updater
                .update_global(iteration, &mut owned_global_record, global_record);

            self.print_progress(iteration, &owned_global_record, &record_beat);
        }

        self.global_updater
            .update_global(0, &mut owned_global_record, global_record)
    }

    fn initialize_particles(sc: &SwarmConfig, jc: &JobConfig) -> Vec<Particle> {
        let mut rng = rand::thread_rng();

        // position and velocity vectors
        let num_points = sc.num_particles * jc.num_variables;
        let mut positions = Vec::with_capacity(num_points);
        let mut velocities = Vec::with_capacity(num_points);

        // fill vectors randomly
        for _ in 0..num_points {
            positions.extend_from_slice(&jc.variable_bounds.sample_bounds(&mut rng));
            velocities.extend_from_slice(&jc.velocity_bounds.sample_bounds(&mut rng));
        }

        // partition random values into individual particles
        positions
            .chunks_exact(jc.num_variables)
            .zip(velocities.chunks_exact(jc.num_variables))
            .map(|(pos, vel)| Particle::new(pos, vel))
            .collect()
    }

    fn print_progress(&self, iteration: usize, record: &OwnedRecord, record_beat: &bool) {
        if let Some(console_update_rate) = self.update_console {
            if iteration % console_update_rate == 0 {
                println!(
                    "{} {} \t iter {} \t {}",
                    self.name, record_beat, iteration, record
                );
            }
        }
    }
}

#[derive(Clone, Debug)]
struct GlobalUpdater {
    update_rate: usize,
    retry: bool,
}

impl GlobalUpdater {
    pub fn new(global_behavior: &GlobalBehavior) -> (Self, bool) {
        let (update_rate, synergic) = match global_behavior {
            GlobalBehavior::Synergic(_, rate) => (*rate, true),
            GlobalBehavior::Solitary => (DEFAULT_UPDATE_RATE, false),
        };

        (
            Self {
                update_rate,
                retry: false,
            },
            synergic,
        )
    }

    pub fn update_global(
        &mut self,
        iteration: usize,
        local_record: &mut OwnedRecord,
        global_record: &mut GlobalRecord,
    ) {
        if iteration % self.update_rate == 0 || self.retry {
            self.retry = false;

            match global_record.beats_global(&local_record) {
                RecordBeatResult::Won => {
                    if let RecordBeatResult::Retry = global_record.update_global(&local_record) {
                        self.retry = true;
                    }
                }
                RecordBeatResult::Lost(new_global_rec) => {
                    *local_record = new_global_rec;
                }
                RecordBeatResult::Retry => {
                    self.retry = true;
                }
                RecordBeatResult::Same => (),
            }
        }
    }
}

#[derive(Clone, Debug)]
struct BoundState {
    pos_bounds: Vec<[f64; 2]>,
    vel_bounds: Vec<[f64; 2]>,
    wall_bounce: f64,
}

impl BoundState {
    pub fn new(sc: &SwarmConfig, jc: &JobConfig) -> Self {
        Self {
            pos_bounds: jc.variable_bounds.bound_array(),
            vel_bounds: jc.velocity_bounds.bound_array(),
            wall_bounce: sc.wall_bounce_factor,
        }
    }

    pub fn values(&self) -> (&[[f64; 2]], &[[f64; 2]], f64) {
        (&self.pos_bounds, &self.vel_bounds, self.wall_bounce)
    }
}

#[derive(Clone, Debug)]
struct StopState {
    iteration: usize,
    stop_condition: StopCondition,
}

impl StopState {
    pub fn new(stop_condition: &StopCondition) -> Self {
        Self {
            iteration: 0,
            stop_condition: stop_condition.clone(),
        }
    }

    pub fn next_iteration(&mut self, global_record: &OwnedRecord) -> Option<usize> {
        match self.stop_condition {
            StopCondition::Iterations(max_iter) => {
                if self.iteration <= max_iter {
                    self.iteration += 1;
                    return Some(self.iteration);
                }
            }
            StopCondition::Cost(min_cost) => {
                if global_record.cost > min_cost {
                    self.iteration += 1;
                    return Some(self.iteration);
                }
            }
            StopCondition::Both(max_iter, min_cost) => {
                if self.iteration <= max_iter && global_record.cost > min_cost {
                    self.iteration += 1;
                    return Some(self.iteration);
                }
            }
        }
        None
    }
}