use crate::molecule::Molecule;
use crate::tree::FenwickTree;
use rand::Rng;
use std::collections::HashMap;
pub struct Reactor {
molecules: Vec<Molecule>, population: Vec<usize>, population_tree: FenwickTree, index_map: HashMap<u64, usize>, }
impl Reactor {
pub fn new() -> Self {
Self {
molecules: Vec::new(),
population: Vec::new(),
population_tree: FenwickTree::new(),
index_map: HashMap::new(),
}
}
pub fn add_or_increment(&mut self, mol: Molecule, hash: u64, delta: usize) {
if let Some(&idx) = self.index_map.get(&hash) {
self.population[idx] += delta;
self.population_tree.update(idx, delta as isize);
} else {
let idx = self.molecules.len();
self.molecules.push(mol);
self.population.push(delta);
self.index_map.insert(hash, idx);
self.population_tree.push(delta);
}
}
pub fn decrement(&mut self, idx: usize, delta: usize) {
assert!(self.population[idx] >= delta);
self.population[idx] -= delta;
self.population_tree.update(idx, -(delta as isize));
}
pub fn sample_weighted<'a>(&'a self, rng: &mut impl Rng) -> Option<&'a Molecule> {
let total = self.population_tree.total_sum();
if total == 0 {
return None;
}
let r = rng.random_range(1..=total);
let idx = self.population_tree.lower_bound(r)?;
Some(&self.molecules[idx])
}
pub fn prune_zero_population(&mut self) {
let mut new_molecules = Vec::new();
let mut new_population = Vec::new();
let mut new_index_map = HashMap::new();
self.population_tree = FenwickTree::new();
for (i, mol) in self.molecules.iter().enumerate() {
if self.population[i] > 0 {
let idx = new_molecules.len();
new_molecules.push(mol.clone());
new_population.push(self.population[i]);
let hash = mol.wl_hash();
new_index_map.insert(hash, idx);
self.population_tree.push(self.population[i]);
}
}
self.molecules = new_molecules;
self.population = new_population;
self.index_map = new_index_map;
}
}