Skip to main content

pso_rs/
model.rs

1use rand::{thread_rng, Rng};
2use rayon::prelude::*;
3use std::fmt;
4pub type Particle = Vec<f64>;
5pub type Population = Vec<Particle>;
6
7/// Model struct
8///
9/// It takes in a `Config` instance and `fn` pointer to an objective function and defines a `run` method for running Particle Swarm Optimization.
10pub struct Model {
11    pub config: Config,
12    pub flat_dim: usize,
13    pub population: Population,
14    pub population_f_scores: Vec<f64>,
15    pub x_best: Particle,
16    pub f_best: f64,
17    obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
18}
19
20impl Model {
21    /// Creates a new Model instance
22    pub fn new(
23        config: Config,
24        obj_f: fn(p: &Particle, flat_dim: usize, dim: &Vec<usize>) -> f64,
25    ) -> Model {
26        // init population
27        let mut rng = thread_rng();
28        let mut flat_dim = 1;
29        for d in config.dimensions.clone() {
30            flat_dim *= d;
31        }
32        let mut population: Population = vec![];
33
34        for _ in 0..config.population_size {
35            let mut particle: Particle = vec![];
36            for flat_i in 0..flat_dim {
37                let true_i = flat_i % config.dimensions[config.dimensions.len() - 1];
38                particle.push(rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1));
39            }
40            population.push(particle);
41        }
42        let population_f_scores = vec![f64::INFINITY; config.population_size];
43        let x_best = population[0].clone();
44        let f_best = population_f_scores[0].clone();
45        let mut model = Model {
46            config,
47            flat_dim,
48            population,
49            population_f_scores,
50            x_best,
51            f_best,
52            obj_f: obj_f,
53        };
54        model.get_f_values();
55        model
56    }
57
58    /// Computes the value of the objective function for each particle and updates best found
59    ///
60    /// Returns the objective function values for all particles
61    ///
62    /// Uses the rayon crate for parallel computation
63    pub fn get_f_values(&mut self) -> Vec<f64> {
64        // find the objective function value for each member of the population
65        if self.config.parallelize {
66            let iter = self.population.par_iter();
67            self.population_f_scores = iter
68                .map(|particle| {
69                    (self.obj_f)(particle, self.flat_dim, &self.config.dimensions)
70                    // self.population_f_scores[i] = f_score;
71                })
72                .collect();
73        } else {
74            let iter = self.population.iter();
75            self.population_f_scores = iter
76                .map(|particle| {
77                    (self.obj_f)(particle, self.flat_dim, &self.config.dimensions)
78                    // self.population_f_scores[i] = f_score;
79                })
80                .collect();
81        }
82        // update best
83        let mut f_best = self.f_best;
84        let mut x_best = self.x_best.clone();
85        for (index, &score) in self.population_f_scores.iter().enumerate() {
86            if score < f_best {
87                f_best = score;
88                x_best = self.population[index].clone();
89            }
90        }
91        self.f_best = f_best;
92        self.x_best = x_best;
93        self.population_f_scores.to_owned()
94    }
95
96    /// Returns the best found objective function value
97    pub fn get_f_best(&self) -> f64 {
98        self.f_best
99    }
100
101    /// Returns the best found minimizer
102    pub fn get_x_best(&self) -> Particle {
103        self.x_best.clone()
104    }
105}
106
107/// Configuration struct
108///
109/// Used to define model parameters
110#[derive(Debug)]
111pub struct Config {
112    pub dimensions: Vec<usize>,
113    pub population_size: usize,
114    pub neighborhood_type: NeighborhoodType,
115    pub rho: usize,
116    pub alpha: f64,
117    pub c1: f64,
118    pub c2: f64,
119    pub lr: f64,
120    pub bounds: Vec<(f64, f64)>,
121    pub t_max: usize,
122    pub progress_bar: bool,
123    pub parallelize: bool,
124}
125
126impl Config {
127    pub fn new() -> Config {
128        Self::default()
129    }
130}
131
132impl Default for Config {
133    fn default() -> Self {
134        Self {
135            dimensions: vec![2],
136            population_size: 1000,
137            neighborhood_type: NeighborhoodType::Lbest,
138            rho: 2,
139            alpha: 0.1,
140            lr: 0.5,
141            c1: 2.05,
142            c2: 2.05,
143            bounds: vec![(-1.0, 1.0); 2],
144            t_max: 1000,
145            progress_bar: true,
146            parallelize: true,
147        }
148    }
149}
150
151#[derive(Debug)]
152pub enum NeighborhoodType {
153    Lbest,
154    Gbest,
155}
156
157impl fmt::Display for NeighborhoodType {
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        match self {
160            NeighborhoodType::Lbest => write!(f, "Local neighborhood (lbest)"),
161            NeighborhoodType::Gbest => write!(f, "Global neighborhood (gbest)"),
162        }
163    }
164}