graphmind_optimization/algorithms/
pso.rs1use 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, pub c1: f64, pub c2: f64, }
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 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 let mut velocities: Vec<Array1<f64>> = (0..self.config.population_size)
42 .map(|_| Array1::zeros(dim))
43 .collect();
44
45 let mut pbests = swarm.clone();
47
48 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 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 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}