Updating people's infection status

examples/births-deaths/infection_manager.rs

@@ -0,0 +1,42 @@
+use ixa::context::Context;
+use ixa::define_rng;
+use ixa::global_properties::ContextGlobalPropertiesExt;
+use ixa::people::{ContextPeopleExt, PersonId, PersonPropertyChangeEvent};
+use ixa::random::ContextRandomExt;
+use rand_distr::Exp;
+
+use crate::population_manager::InfectionStatus;
+use crate::population_manager::InfectionStatusType;
+use crate::population_manager::Alive;
+use crate::Parameters;
+
+define_rng!(InfectionRng);
+
+fn schedule_recovery(context: &mut Context, person_id: PersonId) {
+    let parameters = context.get_global_property_value(Parameters).clone();
+    let infection_duration = parameters.infection_duration;
+    let recovery_time = context.get_current_time()
+        + context.sample_distr(InfectionRng, Exp::new(1.0 / infection_duration).unwrap());
+    if context.get_person_property(person_id, Alive)  {
+        context.add_plan(recovery_time, move |context| {
+            context.set_person_property(person_id, InfectionStatusType, InfectionStatus::R);
+        });
+    }
+}
+
+fn handle_infection_status_change(
+    context: &mut Context,
+    event: PersonPropertyChangeEvent<InfectionStatusType>,
+) {
+    if matches!(event.current, InfectionStatus::I) {
+        schedule_recovery(context, event.person_id);
+    }
+}
+
+pub fn init(context: &mut Context) {
+    context.subscribe_to_event(
+        move |context, event: PersonPropertyChangeEvent<InfectionStatusType>| {
+            handle_infection_status_change(context, event);
+        },
+    );
+}

examples/births-deaths/main.rs

@@ -10,6 +10,7 @@ mod population_manager;
 mod demographics_report;
 mod transmission_manager;
 mod incidence_report;
+mod infection_manager;
 
 use crate::parameters_loader::Parameters;
 
@@ -26,8 +27,10 @@ fn main() {
 
             demographics_report::init(&mut context);
             incidence_report::init(&mut context);
+
             population_manager::init(&mut context);
             transmission_manager::init(&mut context);
+            infection_manager::init(&mut context);
 
             context.add_plan(parameters.max_time, |context| {
                 context.shutdown();

examples/births-deaths/plot_output.R

@@ -1,5 +1,9 @@
 library(tidyverse)
 library(jsonlite)
+## Todo:
+## - Plot population changes
+## - Plot SIR
+## - Compare with theoretical foi with population change
 dir <- file.path("examples", "births-deaths")
 params <- read_json(file.path(dir, "input.json"))
 population <- params$population

examples/births-deaths/population_manager.rs

@@ -143,10 +143,8 @@ impl ContextPopulationExt for Context {
         let mut current_population = 0;
         for i in 0..self.get_current_population() {
             let person_id = self.get_person_id(i);
-            if self.get_person_property(person_id, Alive) == true {
-                if self.get_person_age_group(person_id) == age_group {
+            if self.get_person_property(person_id, Alive) && self.get_person_age_group(person_id) == age_group {
                     current_population += 1;
-                }
             }
         }
         current_population
@@ -157,13 +155,11 @@ impl ContextPopulationExt for Context {
         let mut people_vec = Vec::<PersonId>::new();
         for i in 0..self.get_current_population() {
             let person_id = self.get_person_id(i);
-            if self.get_person_property(person_id, Alive) == true {
-                if self.get_person_age_group(person_id) == age_group {
+            if self.get_person_property(person_id, Alive) && self.get_person_age_group(person_id) == age_group {
                     people_vec.push(person_id);
-                }
             }
         }
-        if people_vec.len() == 0 {
+        if people_vec.is_empty() {
             None
         } else {
             Some(people_vec[self.sample_range(PeopleRng, 0..people_vec.len())])
@@ -180,7 +176,7 @@ where <T as PersonProperty>::Value: PartialEq{
                     people_vec.push(person_id);
             }
         }
-        if people_vec.len() == 0 {
+    if people_vec.is_empty() {
             None
         } else {
             Some(people_vec[self.sample_range(PeopleRng, 0..people_vec.len())])

examples/births-deaths/transmission_manager.rs

@@ -23,7 +23,6 @@ fn attempt_infection(context: &mut Context, age_group: AgeGroupRisk) {
         .get_global_property_value(Foi)
         .get(&age_group)
         .unwrap();
-
     if population_size > 0 {
         let person_to_infect = context.sample_person(age_group).unwrap();
 
@@ -47,16 +46,12 @@ fn attempt_infection(context: &mut Context, age_group: AgeGroupRisk) {
 }
 
 pub fn init(context: &mut Context) {
-    // Need to convert to a more efficient way
-    context.add_plan(0.0, |context| {
-        attempt_infection(context, AgeGroupRisk::NewBorn);
-    });
-
-    context.add_plan(0.0, |context| {
-        attempt_infection(context, AgeGroupRisk::General);
-    });
-
-    context.add_plan(0.0, |context| {
-        attempt_infection(context, AgeGroupRisk::OldAdult);
-    });
+
+    let foi_age_groups = context
+        .get_global_property_value(Foi).clone();
+    for (age_group, _) in foi_age_groups {
+        context.add_plan(0.0, move |context| {
+            attempt_infection(context, age_group);
+        });
+    }
 }