Fully abstracted facemask from intervention mgr

CDCgov · Dec 13, 2024 · fdb9fc6 · fdb9fc6
1 parent 440bf35
commit fdb9fc6
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 25 deletions.
diff --git a/src/facemask_manager.rs b/src/facemask_manager.rs
@@ -1,20 +1,30 @@
-use crate::intervention_manager::{ContextInterventionExt, FacemaskStatus, FacemaskStatusType};
+use crate::intervention_manager::ContextInterventionExt;
 use crate::parameters::Parameters;
 use crate::population_loader::Alive;
 use crate::transmission_manager::InfectiousStatusType;
 use ixa::{
-    define_rng, Context, ContextGlobalPropertiesExt, ContextPeopleExt, ContextRandomExt, PersonId,
+    define_person_property, define_person_property_with_default, define_rng, Context,
+    ContextGlobalPropertiesExt, ContextPeopleExt, ContextRandomExt, PersonId,
 };
+use std::hash::Hash;
+
+#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)]
+pub enum FacemaskStatusType {
+    None,
+    Wearing,
+}
+
+define_person_property_with_default!(FacemaskStatus, FacemaskStatusType, FacemaskStatusType::None);
 
 define_rng!(FacemaskRng);
 
 pub fn init(context: &mut Context) {
-    context.register_facemask(
+    context.register_intervention(
         InfectiousStatusType::Susceptible,
         FacemaskStatusType::Wearing,
         0.5,
     );
-    context.register_facemask(
+    context.register_intervention(
         InfectiousStatusType::Infectious,
         FacemaskStatusType::Wearing,
         0.25,

diff --git a/src/intervention_manager.rs b/src/intervention_manager.rs
@@ -1,15 +1,7 @@
 use crate::transmission_manager::{InfectiousStatus, InfectiousStatusType};
-use ixa::{
-    define_data_plugin, define_person_property, define_person_property_with_default,
-    people::ContextPeopleExt, Context, PersonId, PersonProperty,
-};
-use std::{collections::HashMap, hash::Hash};
+use ixa::{define_data_plugin, people::ContextPeopleExt, Context, PersonId, PersonProperty};
+use std::collections::HashMap;
 
-#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)]
-pub enum FacemaskStatusType {
-    None,
-    Wearing,
-}
 struct InterventionContainer {
     intervention_map: HashMap<InfectiousStatusType, HashMap<String, f64>>,
 }
@@ -22,8 +14,6 @@ define_data_plugin!(
     }
 );
 
-define_person_property_with_default!(FacemaskStatus, FacemaskStatusType, FacemaskStatusType::None);
-
 pub fn init(context: &mut Context) {
     let _ = context.get_data_container_mut(InterventionPlugin);
 }
@@ -34,10 +24,10 @@ pub trait ContextInterventionExt {
         person_id: PersonId,
         intervention_type: T,
     ) -> f64;
-    fn register_facemask(
+    fn register_intervention<T: std::fmt::Debug>(
         &mut self,
         infectious_status: InfectiousStatusType,
-        facemask_status: FacemaskStatusType,
+        intervention_status: T,
         relative_transmission: f64,
     );
 }
@@ -62,25 +52,26 @@ impl ContextInterventionExt for Context {
             .unwrap_or(&1.0)
     }
 
-    fn register_facemask(
+    fn register_intervention<T: std::fmt::Debug>(
         &mut self,
         infectious_status: InfectiousStatusType,
-        facemask_status: FacemaskStatusType,
+        intervention_status: T,
         relative_transmission: f64,
     ) {
-        let mut facemask_map = HashMap::new();
+        let mut transmission_map = HashMap::new();
 
-        facemask_map.insert(format!("{facemask_status:?}"), relative_transmission);
+        transmission_map.insert(format!("{intervention_status:?}"), relative_transmission);
 
         self.get_data_container_mut(InterventionPlugin)
             .intervention_map
-            .insert(infectious_status, facemask_map);
+            .insert(infectious_status, transmission_map);
     }
 }
 
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::facemask_manager::{FacemaskStatus, FacemaskStatusType};
     use ixa::{people::ContextPeopleExt, Context};
 
     #[test]
@@ -89,7 +80,7 @@ mod test {
         let mut context = Context::new();
         let contact_id = context.add_person(()).unwrap();
 
-        context.register_facemask(
+        context.register_intervention(
             InfectiousStatusType::Susceptible,
             FacemaskStatusType::Wearing,
             0.5,

diff --git a/src/transmission_manager.rs b/src/transmission_manager.rs
@@ -8,7 +8,8 @@ use ixa::{
 };
 use statrs::distribution::{ContinuousCDF, Exp, Poisson};
 
-use crate::intervention_manager::{ContextInterventionExt, FacemaskStatus};
+use crate::facemask_manager::FacemaskStatus;
+use crate::intervention_manager::ContextInterventionExt;
 use crate::parameters::Parameters;
 use crate::{contact::ContextContactExt, population_loader::Alive};