1use crate::steps::EngineStep;
2use radiate_core::{
3 Chromosome, Ecosystem, Executor, MetricSet, Objective, Phenotype, Population, RateSet, Species,
4 diversity::Diversity, math::distribution, metric_names, random_provider,
5};
6use radiate_error::Result;
7use std::sync::{Arc, Mutex, RwLock};
8
9type SpeciesAssignments = Vec<Option<(usize, f32)>>;
10
11pub struct SpeciateStep<C>
12where
13 C: Chromosome,
14{
15 pub(crate) threshold: RateSet,
16 pub(crate) objective: Objective,
17 pub(crate) distance: Arc<dyn Diversity<C>>,
18 pub(crate) executor: Arc<Executor>,
19 pub(crate) distances: Vec<f32>,
20 pub(crate) assignments: Arc<Mutex<SpeciesAssignments>>,
21}
22
23impl<C: Chromosome> SpeciateStep<C> {
24 pub fn new(
25 threshold: impl Into<RateSet>,
26 objective: Objective,
27 distance: Arc<dyn Diversity<C>>,
28 executor: Arc<Executor>,
29 ) -> Self {
30 Self {
31 threshold: threshold.into(),
32 objective,
33 distance,
34 executor,
35 distances: Vec::new(),
36 assignments: Arc::new(Mutex::new(Vec::new())),
37 }
38 }
39}
40
41impl<C> SpeciateStep<C>
42where
43 C: Chromosome + 'static,
44{
45 fn assign_species(
46 &mut self,
47 generation: usize,
48 threshold: f32,
49 ecosystem: &mut Ecosystem<C>,
50 mascots: Arc<Vec<Phenotype<C>>>,
51 assignments: Arc<Mutex<SpeciesAssignments>>,
52 ) -> Result<()>
53 where
54 C: Clone,
55 {
56 let pop_len = ecosystem.population().len();
57 let num_threads = self.executor.num_workers().max(1);
58 let chunk_size = (pop_len as f32 / num_threads as f32).ceil() as usize;
59
60 let mut batches = Vec::new();
61
62 let mut empty_population = Population::empty();
63 std::mem::swap(ecosystem.population_mut(), &mut empty_population);
64 let population = Arc::new(RwLock::new(empty_population));
65
66 for chunk_start in (0..pop_len).step_by(chunk_size) {
67 let chunk_end = (chunk_start + chunk_size).min(pop_len);
68
69 let distance = Arc::clone(&self.distance);
70 let assignments = Arc::clone(&assignments);
71 let population = Arc::clone(&population);
72 let species_snapshot = Arc::clone(&mascots);
73
74 batches.push(move || {
75 Self::process_chunk(
76 population,
77 species_snapshot,
78 threshold,
79 distance,
80 assignments,
81 chunk_start..chunk_end,
82 );
83 });
84 }
85
86 self.executor.submit_blocking(batches);
87
88 std::mem::swap(ecosystem.population_mut(), &mut population.write().unwrap());
89 self.assign_unassigned(
90 generation,
91 threshold,
92 ecosystem,
93 &assignments.lock().unwrap(),
94 );
95
96 Ok(())
97 }
98
99 #[inline]
100 fn assign_unassigned(
101 &mut self,
102 generation: usize,
103 threshold: f32,
104 ecosystem: &mut Ecosystem<C>,
105 assignments: &[Option<(usize, f32)>],
106 ) where
107 C: Clone,
108 {
109 let pop_len = ecosystem.population().len();
110 let mut new_count = 0;
111
112 for (i, assignment) in assignments.iter().enumerate().take(pop_len) {
113 if let Some((species_id, dist)) = assignment {
114 ecosystem.add_species_member(*species_id, i);
115 self.distances[i] = *dist;
116 continue;
117 }
118
119 let mut best_dist = f32::MAX;
120 let phenotype = ecosystem.get_phenotype(i).unwrap();
121 let maybe_idx = ecosystem.species().and_then(|specs| {
122 for (species_idx, species) in specs.iter().enumerate() {
123 let dist = self.distance.measure(phenotype, species.mascot());
124
125 best_dist = best_dist.min(dist);
126
127 if dist < threshold {
128 self.distances[i] = dist;
129 return Some(species_idx);
130 }
131 }
132
133 None
134 });
135
136 match maybe_idx {
137 Some(idx) => ecosystem.add_species_member(idx, i),
138 None => {
139 if let Some(pheno) = ecosystem.get_phenotype_mut(i) {
140 let new_species = Species::new(generation, pheno.clone());
141 let species_idx = ecosystem.push_species(new_species);
142
143 ecosystem.add_species_member(species_idx, i);
144 self.distances[i] = best_dist;
145
146 new_count += 1;
147 }
148 }
149 }
150 }
151
152 if new_count == pop_len {
153 ecosystem.clear_species();
154 self.distances.clear();
155 self.distances.push(threshold);
156 let idx = random_provider::range(0..pop_len);
157 if let Some(phenotype) = ecosystem.get_phenotype(idx) {
158 let new_species = Species::new(generation, (*phenotype).clone());
159 let new_species_idx = ecosystem.push_species(new_species);
160
161 for i in 0..pop_len {
162 if i != idx {
163 ecosystem.add_species_member(new_species_idx, i);
164 }
165 }
166 }
167 }
168 }
169
170 #[inline]
171 fn process_chunk(
172 population: Arc<RwLock<Population<C>>>,
173 species_mascots: Arc<Vec<Phenotype<C>>>,
174 threshold: f32,
175 distance: Arc<dyn Diversity<C>>,
176 assignments: Arc<Mutex<SpeciesAssignments>>,
177 range: std::ops::Range<usize>,
178 ) {
179 let mut inner_assignments = Vec::new();
180
181 let start = range.start;
182 let reader = population.read().unwrap();
183 for (idx, individual) in reader[range].iter().enumerate() {
184 let mut assigned = None;
185 for (spec_idx, sp) in species_mascots.iter().enumerate() {
186 let dist = distance.measure(individual, sp);
187
188 if dist < threshold {
189 assigned = Some((spec_idx, dist));
190 break;
191 }
192 }
193
194 if assigned.is_some() {
195 inner_assignments.push((start + idx, assigned));
196 }
197 }
198
199 {
200 let mut assignments = assignments.lock().unwrap();
201 for (idx, assigned) in inner_assignments {
202 assignments[idx] = assigned;
203 }
204 }
205 }
206
207 #[inline]
208 fn generate_mascots(ecosystem: &mut Ecosystem<C>) -> Arc<Vec<Phenotype<C>>>
209 where
210 C: Clone,
211 {
212 let (species, population) = ecosystem.species_population_mut();
213
214 if let Some(species) = species {
215 for spec in species.iter_mut() {
216 let species_members = population
217 .iter_species(spec.id())
218 .collect::<Vec<&Phenotype<C>>>();
219
220 if species_members.is_empty() {
221 continue;
222 }
223
224 let idx = random_provider::range(0..species_members.len());
225 if let Some(phenotype) = species_members.get(idx) {
226 spec.set_new_mascot((*phenotype).clone());
227 }
228 }
229 }
230
231 Arc::new(
232 ecosystem
233 .species_mascots()
234 .into_iter()
235 .cloned()
236 .collect::<Vec<Phenotype<C>>>(),
237 )
238 }
239
240 fn calc_species_metrics(
241 &mut self,
242 generation: usize,
243 ecosystem: &Ecosystem<C>,
244 metrics: &mut MetricSet,
245 ) {
246 let Some(species) = ecosystem.species() else {
247 return;
248 };
249
250 let pop_len = ecosystem.population().len().max(1);
251 let mut new_species_count = 0;
252 let mut ages = Vec::with_capacity(species.len());
253 let mut sizes = Vec::with_capacity(species.len());
254 let mut max_size = 0;
255
256 for spec in species.iter() {
257 let age = spec.age(generation);
258 let len = spec.len();
259
260 if age == 0 {
261 new_species_count += 1;
262 }
263
264 ages.push(age);
265 sizes.push(len);
266 max_size = max_size.max(len);
267 }
268
269 let largest_share = max_size as f32 / pop_len as f32;
270 let evenness = distribution::evenness(&sizes);
271 let s_count = species.len();
272
273 let churn = if s_count > 0 {
274 new_species_count as f32 / s_count as f32
275 } else {
276 0.0
277 };
278
279 metrics.upsert(metric_names::SPECIES_AGE, &ages);
280 metrics.upsert(metric_names::SPECIES_SIZE, &sizes);
281 metrics.upsert(metric_names::SPECIES_COUNT, s_count);
282 metrics.upsert(metric_names::SPECIES_CREATED, new_species_count);
283 metrics.upsert(metric_names::SPECIES_EVENNESS, evenness);
284 metrics.upsert(metric_names::SPECIES_NEW_RATIO, churn);
285 metrics.upsert(metric_names::LARGEST_SPECIES_SHARE, largest_share);
286 }
287}
288
289impl<C> EngineStep<C> for SpeciateStep<C>
290where
291 C: Chromosome + PartialEq + Clone + 'static,
292{
293 #[inline]
294 fn execute(
295 &mut self,
296 generation: usize,
297 ecosystem: &mut Ecosystem<C>,
298 metrics: &mut MetricSet,
299 ) -> Result<()> {
300 let pop_len = ecosystem.population().len();
301 if pop_len == 0 {
302 return Ok(());
303 }
304
305 let threshold = self.threshold.calculate_control_rate(generation, metrics)?;
306 let mascots = Self::generate_mascots(ecosystem);
307
308 self.distances.clear();
309 self.distances.resize(pop_len, 0.0);
310
311 let assignments = {
312 let mut assignments_guard = self.assignments.lock().unwrap();
313 assignments_guard.clear();
314 assignments_guard.resize(pop_len, None);
315 Arc::clone(&self.assignments)
316 };
317
318 self.assign_species(
319 generation,
320 threshold,
321 ecosystem,
322 Arc::clone(&mascots),
323 Arc::clone(&assignments),
324 )?;
325
326 let rm_species_count = ecosystem.remove_dead_species();
327
328 metrics.upsert(metric_names::SPECIES_DISTANCE_DIST, &self.distances);
329 metrics.upsert(metric_names::SPECIES_DIED, rm_species_count);
330
331 self.calc_species_metrics(generation, ecosystem, metrics);
332
333 ecosystem.fitness_share(&self.objective);
334
335 Ok(())
336 }
337}