Skip to main content

graphmind_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!(
57                    "PSO Solver: Iteration {}/{}",
58                    iter, self.config.max_iterations
59                );
60            }
61
62            history.push(gbest.fitness);
63
64            // Update swarm
65            // Note: In parallel, we need to collect updates then apply?
66            // Or we can update particle i using its own pbest and the *current* gbest (read-only).
67            // Updating velocities requires mutable access to velocities[i].
68            // Updating positions requires mutable access to swarm[i].
69
70            // We'll compute new state in parallel and then replace.
71            let results: Vec<(Individual, Array1<f64>, Individual)> = swarm
72                .par_iter()
73                .zip(velocities.par_iter())
74                .zip(pbests.par_iter())
75                .map(|((particle, velocity), pbest)| {
76                    let mut local_rng = thread_rng();
77                    let mut new_vel = Array1::zeros(dim);
78                    let mut new_vars = Array1::zeros(dim);
79
80                    for j in 0..dim {
81                        let r1: f64 = local_rng.gen();
82                        let r2: f64 = local_rng.gen();
83
84                        let v = self.w * velocity[j]
85                            + self.c1 * r1 * (pbest.variables[j] - particle.variables[j])
86                            + self.c2 * r2 * (gbest.variables[j] - particle.variables[j]);
87
88                        new_vel[j] = v;
89                        new_vars[j] = (particle.variables[j] + v).clamp(lower[j], upper[j]);
90                    }
91
92                    let new_fitness = problem.fitness(&new_vars);
93                    let new_ind = Individual::new(new_vars, new_fitness);
94
95                    let new_pbest = if new_fitness < pbest.fitness {
96                        new_ind.clone()
97                    } else {
98                        pbest.clone()
99                    };
100
101                    (new_ind, new_vel, new_pbest)
102                })
103                .collect();
104
105            // Unpack results
106            for (i, (new_ind, new_vel, new_pbest)) in results.into_iter().enumerate() {
107                swarm[i] = new_ind;
108                velocities[i] = new_vel;
109                pbests[i] = new_pbest;
110
111                if swarm[i].fitness < gbest.fitness {
112                    gbest = swarm[i].clone();
113                }
114            }
115        }
116
117        OptimizationResult {
118            best_variables: gbest.variables.clone(),
119            best_fitness: gbest.fitness,
120            history,
121        }
122    }
123
124    fn find_best(&self, population: &[Individual]) -> usize {
125        let mut best_idx = 0;
126        for (i, ind) in population.iter().enumerate() {
127            if ind.fitness < population[best_idx].fitness {
128                best_idx = i;
129            }
130        }
131        best_idx
132    }
133}