Skip to main content

samyama_optimization/algorithms/
pso.rs

1use crate::common::{Individual, OptimizationResult, Problem, SolverConfig};
2use ndarray::Array1;
3use rand::prelude::*;
4use rayon::prelude::*;
5
6pub struct PSOSolver {
7    pub config: SolverConfig,
8    pub w: f64,  // Inertia weight
9    pub c1: f64, // Cognitive weight (pbest)
10    pub c2: f64, // Social weight (gbest)
11}
12
13impl PSOSolver {
14    pub fn new(config: SolverConfig) -> Self {
15        Self { 
16            config,
17            w: 0.7,
18            c1: 1.5,
19            c2: 1.5,
20        }
21    }
22
23    pub fn solve<P: Problem>(&self, problem: &P) -> OptimizationResult {
24        let mut rng = thread_rng();
25        let dim = problem.dim();
26        let (lower, upper) = problem.bounds();
27
28        // Initialize population (swarm)
29        let mut swarm: Vec<Individual> = (0..self.config.population_size)
30            .map(|_| {
31                let mut vars = Array1::zeros(dim);
32                for i in 0..dim {
33                    vars[i] = rng.gen_range(lower[i]..upper[i]);
34                }
35                let fitness = problem.fitness(&vars);
36                Individual::new(vars, fitness)
37            })
38            .collect();
39
40        // Initialize velocities
41        let mut velocities: Vec<Array1<f64>> = (0..self.config.population_size)
42            .map(|_| Array1::zeros(dim))
43            .collect();
44
45        // Initialize personal bests (pbest)
46        let mut pbests = swarm.clone();
47
48        // Initialize global best (gbest)
49        let gbest_idx = self.find_best(&swarm);
50        let mut gbest = swarm[gbest_idx].clone();
51
52        let mut history = Vec::with_capacity(self.config.max_iterations);
53
54        for iter in 0..self.config.max_iterations {
55            if iter % 10 == 0 {
56                println!("PSO Solver: Iteration {}/{}", iter, self.config.max_iterations);
57            }
58            
59            history.push(gbest.fitness);
60
61            // Update swarm
62            // Note: In parallel, we need to collect updates then apply? 
63            // Or we can update particle i using its own pbest and the *current* gbest (read-only).
64            // Updating velocities requires mutable access to velocities[i].
65            // Updating positions requires mutable access to swarm[i].
66            
67            // We'll compute new state in parallel and then replace.
68            let results: Vec<(Individual, Array1<f64>, Individual)> = swarm.par_iter().zip(velocities.par_iter()).zip(pbests.par_iter())
69                .map(|((particle, velocity), pbest)| {
70                    let mut local_rng = thread_rng();
71                    let mut new_vel = Array1::zeros(dim);
72                    let mut new_vars = Array1::zeros(dim);
73
74                    for j in 0..dim {
75                        let r1: f64 = local_rng.gen();
76                        let r2: f64 = local_rng.gen();
77                        
78                        let v = self.w * velocity[j] 
79                              + self.c1 * r1 * (pbest.variables[j] - particle.variables[j])
80                              + self.c2 * r2 * (gbest.variables[j] - particle.variables[j]);
81                        
82                        new_vel[j] = v;
83                        new_vars[j] = (particle.variables[j] + v).clamp(lower[j], upper[j]);
84                    }
85
86                    let new_fitness = problem.fitness(&new_vars);
87                    let new_ind = Individual::new(new_vars, new_fitness);
88                    
89                    let new_pbest = if new_fitness < pbest.fitness {
90                        new_ind.clone()
91                    } else {
92                        pbest.clone()
93                    };
94
95                    (new_ind, new_vel, new_pbest)
96                })
97                .collect();
98
99            // Unpack results
100            for (i, (new_ind, new_vel, new_pbest)) in results.into_iter().enumerate() {
101                swarm[i] = new_ind;
102                velocities[i] = new_vel;
103                pbests[i] = new_pbest;
104                
105                if swarm[i].fitness < gbest.fitness {
106                    gbest = swarm[i].clone();
107                }
108            }
109        }
110
111        OptimizationResult {
112            best_variables: gbest.variables.clone(),
113            best_fitness: gbest.fitness,
114            history,
115        }
116    }
117
118    fn find_best(&self, population: &[Individual]) -> usize {
119        let mut best_idx = 0;
120        for (i, ind) in population.iter().enumerate() {
121            if ind.fitness < population[best_idx].fitness {
122                best_idx = i;
123            }
124        }
125        best_idx
126    }
127}