Fix attribute and fitness computations

Author: sirno

flake.lock

@@ -76,7 +76,7 @@ pub trait AttributeProvider<S: Symbol>: Sync + Send + std::fmt::Debug {
 ///
 /// The attribute set definition contains the providers that compute the attributes of the set.
 /// It can be used to register new providers and to create an attribute set.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct AttributeSetDefinition<S: Symbol> {
     providers: HashMap<&'static str, Arc<dyn AttributeProvider<S> + Send + Sync>>,
     eager: Vec<&'static str>,
@@ -192,6 +192,30 @@ impl<S: Symbol> AttributeSet<S> {
         Self::new(self.definition.clone(), haplotype)
     }
 
+    /// Get or compute the value of an attribute.
+    pub fn get_or_compute_raw(&self, id: &'static str) -> Result<AttributeValue> {
+        let provider = self.definition.providers.get(&id).ok_or_else(|| {
+            VirolutionError::ImplementationError(format!("No provider found for attribute {}", id))
+        })?;
+
+        // First, try to read the value
+        {
+            let values = self.values.read().unwrap();
+            if let Some(value) = values.get(&id) {
+                return Ok(value.clone());
+            }
+        }
+
+        // Compute the attribute using the provider
+        let value = provider.compute(&self.haplotype.upgrade());
+
+        // Write the computed value
+        let mut values = self.values.write().unwrap();
+        values.insert(id, value.clone());
+
+        Ok(value)
+    }
+
     /// Get or compute the value of an attribute.
     pub fn get_or_compute(&self, id: &'static str) -> Result<AttributeValue> {
         let provider = self.definition.providers.get(&id).ok_or_else(|| {
@@ -207,20 +231,13 @@ impl<S: Symbol> AttributeSet<S> {
         }
 
         // Compute the attribute using the provider
-        if let Some(provider) = self.definition.providers.get(&id) {
-            let value = provider.compute(&self.haplotype.upgrade());
-
-            // Write the computed value
-            let mut values = self.values.write().unwrap();
-            values.insert(id, value.clone());
-
-            Ok(provider.map(value))
-        } else {
-            // If there is no provider, return None
-            Err(VirolutionError::ImplementationError(
-                "No provider found for attribute".to_string(),
-            ))
-        }
+        let value = provider.compute(&self.haplotype.upgrade());
+
+        // Write the computed value
+        let mut values = self.values.write().unwrap();
+        values.insert(id, value.clone());
+
+        Ok(provider.map(value))
     }
 
     /// Get the value of an already computed attribute.
@@ -260,7 +277,9 @@ impl<S: Symbol> Clone for AttributeSet<S> {
 impl<S: Symbol> std::fmt::Debug for AttributeSet<S> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let values = self.values.read().unwrap();
+        let haplotype = self.haplotype.upgrade();
         f.debug_struct("AttributeSet")
+            .field("haplotype", &haplotype)
             .field("values", &values)
             .finish()
     }

src/core/attributes.rs

@@ -287,10 +287,21 @@ impl<S: Symbol> Haplotype<S> {
         recombinant
     }
 
+    #[allow(dead_code)]
+    pub(crate) fn is_wildtype(&self) -> bool {
+        matches!(self, Haplotype::Wildtype(_))
+    }
+
+    #[allow(dead_code)]
     pub(crate) fn is_mutant(&self) -> bool {
         matches!(self, Haplotype::Mutant(_))
     }
 
+    #[allow(dead_code)]
+    pub(crate) fn is_recombinant(&self) -> bool {
+        matches!(self, Haplotype::Recombinant(_))
+    }
+
     /// Unwraps the haplotype into a mutant.
     ///
     /// This function will panic if the haplotype is not a mutant and is only intended for internal
@@ -466,6 +477,10 @@ impl<S: Symbol> Haplotype<S> {
         self.get_attributes().get(id)
     }
 
+    pub fn get_or_compute_attribute_raw(&self, id: &'static str) -> Result<AttributeValue> {
+        self.get_attributes().get_or_compute_raw(id)
+    }
+
     pub fn get_or_compute_attribute(&self, id: &'static str) -> Result<AttributeValue> {
         self.get_attributes().get_or_compute(id)
     }

src/core/haplotype.rs

@@ -60,12 +60,13 @@ impl<S: Symbol> HostSpecs<S> {
     }
 
     pub fn try_get_spec_from_index(&self, index: usize) -> Option<&HostSpec<S>> {
-        for spec in self.0.iter() {
-            if spec.range.contains(&index) {
-                return Some(spec);
-            }
-        }
-        None
+        self.0.iter().find(|&spec| spec.range.contains(&index))
+    }
+}
+
+impl<S: Symbol> Default for HostSpecs<S> {
+    fn default() -> Self {
+        Self::new()
     }
 }
 
@@ -111,6 +112,19 @@ pub struct HostMapBuffer {
     hosts: Vec<usize>,
 }
 
+impl std::fmt::Debug for HostMapBuffer {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let n_infections = self.infections.iter().filter(|x| x.is_some()).count();
+        let infectivity = n_infections as f64 / self.n_infectants as f64;
+        f.debug_struct("HostMapBuffer")
+            .field("n_hosts", &self.n_hosts)
+            .field("n_infectants", &self.n_infectants)
+            .field("n_infections", &n_infections)
+            .field("infectivity", &infectivity)
+            .finish()
+    }
+}
+
 /// A map from host index to infectant indices.
 ///
 /// The map is constructed from a list of infections between infectant and host indices. It uses
@@ -149,11 +163,11 @@ impl HostMapBuffer {
         }
     }
 
-    pub fn build<F: FnMut(&mut Option<usize>)>(&mut self, builder: F) {
+    pub fn build<F: FnMut((usize, &mut Option<usize>))>(&mut self, builder: F) {
         let infection_slice = &mut self.infections[0..self.n_infectants];
 
         // set new infections
-        infection_slice.iter_mut().for_each(builder);
+        infection_slice.iter_mut().enumerate().for_each(builder);
 
         // reset offsets
         self.offsets[0..=self.n_hosts].fill(0);

src/core/hosts.rs

@@ -180,9 +180,11 @@ impl LognormalParameters {
 
 impl FileParameters {
     pub fn load_table(&self) -> Vec<f64> {
-        let reader = NpyFile::new(std::fs::File::open(&self.path).unwrap_or_else(
-            |_e| panic!("Could not open file: {}", self.path),
-        )).unwrap();
+        let reader = NpyFile::new(
+            std::fs::File::open(&self.path)
+                .unwrap_or_else(|_e| panic!("Could not open file: {}", self.path)),
+        )
+        .unwrap();
         reader
             .data::<f64>()
             .unwrap()

src/init/fitness.rs

@@ -216,8 +216,8 @@ impl<S: Symbol> FitnessProvider<S> {
                     haplotype
                         .try_get_ancestor()
                         .expect("Could not find ancestor during fitness update...")
-                        .get_or_compute_attribute(self.name)
-                        .unwrap(),
+                        .get_or_compute_attribute_raw(self.name)
+                        .unwrap()
                 )
                 .unwrap();
                 let update = self.update_fitness(haplotype);

src/providers/fitness.rs

@@ -73,7 +73,7 @@ impl Runner {
 
         let wildtype = Haplotype::load_wildtype(sequence, &attribute_definitions);
 
-       // perform sanity checks
+        // perform sanity checks
         if !settings
             .schedule
             .check_transfer_table_sizes(args.n_compartments)
@@ -334,7 +334,7 @@ impl Runner {
 
         for generation in 0..=self.args.generations {
             // logging
-            log::debug!("Generate logging message for generation {generation}...");
+            log::debug!("Running generation {generation}...");
             let population_sizes: Vec<usize> = self
                 .simulations
                 .iter()

src/runner.rs

@@ -366,20 +366,20 @@ impl<S: Symbol> Simulation<S> for BasicSimulation<S> {
         let host_sampler = Uniform::new(0, self.parameters.host_population_size)
             .expect("Invalid host population size");
         let mut rng = rand::rng();
-        self.host_map_buffer.build(|ref mut infectant| {
-            let host_candidate = host_sampler.sample(&mut rng);
-            **infectant = self
-                .host_specs
-                .try_get_spec_from_index(host_candidate)
-                .map(|spec| {
-                    if spec.host.infect(&self.wildtype, &mut rng) {
-                        Some(host_candidate)
-                    } else {
-                        None
-                    }
-                })
-                .flatten();
-        });
+        self.host_map_buffer
+            .build(|(infectant, ref mut infection)| {
+                let host_candidate = host_sampler.sample(&mut rng);
+                **infection = self
+                    .host_specs
+                    .try_get_spec_from_index(host_candidate)
+                    .and_then(|spec| {
+                        if spec.host.infect(&self.population.get(&infectant), &mut rng) {
+                            Some(host_candidate)
+                        } else {
+                            None
+                        }
+                    })
+            });
     }
 
     #[cfg(feature = "parallel")]