Updated README and added test for infection manager with deaths

examples/births-deaths/README.md

@@ -3,105 +3,141 @@
 This example describes the process of loading a population, adding new births and deaths. It has a simple sir model that represents infections from a constant force of infection that differs by age category (0,2m, 6m, 65y).
 
 # Features
-
-  * Adding people to the simulation with person properties that are set by default for newborns
-  * Looking up people based on their current age
-  * Introduces the concept of updated person properties. For instance, if a person's risk property depends on their age, it should change when they age.
-  * Introduces deaths and how they are removed from the population.
+  * Adding newborns to the simulation with some default person properties,
+  * Age-Groups that update with time,
+  * Person look-up based on their properties, or Age-Groups,
+  * Deaths and plan removal to ensure that dead agents are not executing plans.
 
 # Simulation overview
-The simulation loads parameters that contain a seed, population size, birth and death rates, as well as foi rates. Then, a the population manager module is initialized to set up the initial population and their person properties. This module is also in charge of scheduling new births and deaths based on the specified rates in parameters. Finally, the transmission manager module is initialized. Infection attempts are scheduled to occur based on the constant age-specific force of infection. Once an infection event is scheduled, a susceptible individual is selected to be infected. This individual is required to be alive at the time of the infection attempt. After an infection attempt is finished, the next infection event is scheduled based on a constant force of infection.
-
-Infected individuals schedule their recovery at time `t + infection_period`. The infection status of recovered individuals remains as recovered for the rest of the simulation.
+This simulation builds on the basic infection example. Mainly, susceptible individuals are exposed to a pathogen at a risk determined by a constant force of infection, which varies by age group. Three age groups are defined. Namely, `Newborns (< 1yr)`, `General (1-65yr)`, and `OldAdult( 65+yr)`. At initialization, people are added to the simulation with an age from 0 - 100 yrs and assigned one of the age groups. However, membership to these groups can change over time as people age. People can also be added to the simulation as newborns, as they can die at any point during the simulation.
 
+The exposure process is initialized by scheduling an initial infection attempt for each age group. This depends on the force of infection defined for each age group, and the current size of the group, which only includes people who are alive. Individuals who are infected by the pathogen may recover based on a recovery period (`t + infection_period`) specified at initialization. These recovery times are scheduled at the time of infection as a `plan`. The `id` for this plan is saved in a data structure that holds the current plans for recovery scheduled for each individual. Once the person recovers, these plans are removed from the data structure. However, if the person dies before recovering, these plans are canceled at the time of death. The infection status of recovered individuals remains as recovered for the rest of the simulation, which will stop when the plan queue is empty.
 
-```rust main.rs
-fn main() {
-	let mut context = Context::new();
-	let current_dir = Path::new(file!()).parent().unwrap();
-	let file_path = current_dir
-		.join("input.json");
+# Population manager
+At initialization, the population manager adds people to the simulation as defined in the input parameters file, and initializes each person with two person properties: `Birth` time and `Alive` status. Birth time is estimated as a random number between 0 and -100 to represent ages from 0 to 100 years. To trait `get_person_age` can return a person's age based on their time of birth.
 
-	parameters_loader::init();
+## Births
+New people are constantly added to the simulation based on a birth rate, which is defined by the input parameters. The creation of new individuals is performed as any other person added to the simulation, and their person property `Birth` is assigned to the current simulation time `context.get_current_time()`.
 
-	let parameters = context.get_global_property_value(Parameters).clone();
-	context.init_random(parameters.seed);
-	population_manager::init(&mut context);
-	transmission_manager::init(&mut context);
+```rust
+fn create_new_person(&mut self, birth_time: f64) -> PersonId {
+    let person = self.add_person();
+    self.initialize_person_property(person, Birth, birth_time);
+    self.initialize_person_property(person, Alive, true);
+    person
 }
-
 ```
 
-# People and person properties
-When the `Population manager` module initializes, a number of persons are created and given a unique person id (from `0` to `population_size`). This functionality is provided by an `create_person` method from the `people` module in `ixa`, which adds them to a `People` data container. This function also defines a special person property that determines people's life status in the simulation `define_person_property!(Alive, bool, true)`.
-
-The population manager also defines an infection status person property and an age property, which is assigned randomly based on a uniform distribution.
+## Deaths
+People are constantly removed from the simulation based on a death rate, which is defined by the input parameters. Every time a death is scheduled to occur, the function `attempt_death` is called, which will set person property `Alive` to false. **Plans are not directly canceled by the population manager, this is done directly in the module that schedules the plan (e.g., `infection_manager`) by subscribing to events related to changes in the person property `Alive`. It is important to keep in mind that dead individuals should not be counted for the force of infection or other transmission events.
 
 ```rust
-InfectionStatus = enum(
-    S,
-    I,
-    R
-);
-
-pub enum RiskCategory {
-    High,
-    Low,
+fn attempt_death(&mut self, person_id) {
+	self.set_person_property(person_id, Alive, false);
 }
+```
+
+## Age Groups
+Age groups are defined in the population manager as an `enum`. These groups are determined for each person using the trait `get_person_age_group`. This function estimates the current age group based on the time of the simulation and the time of birth. In this example, the force of infection varies for each of the three age groups defined below. Hence, a hash map contains the force of infection for each of these groups and is saved as a global property.
 
-define_person_property!(Age, u8); // Age in days
-define_person_property!(RiskCategoryType, RiskCategory); // This is a derived property that depends on age.
-define_person_property!(InfectionStatusType, InfectionStatus, InfectionStatus::S);
+```rust
+pub enum AgeGroupRisk {
+    NewBorn,
+    General,
+    OldAdult,
+}
 
-for (person_id in 0..parameters.get_parameter(population_size)) {
-    context.create_person(person_id = person_id)
-	let age_in_days = context.sample_unif(PeopleRng, 0, 100 * 365);
-	context.initialize_person_property(person, Age, age_in_days);
-	let	risk_category = RiskCategory.get_risk_category(age_in_days);
-    context.initialize_person_property(person, RiskCategoryType, risk_category);
+fn get_person_age_group(&mut self, person_id: PersonId) -> AgeGroupRisk {
+	let current_age = self.get_person_age(person_id);
+    if current_age <= 1.0 {
+		AgeGroupRisk::NewBorn
+    } else if current_age <= 65.0 {
+        AgeGroupRisk::General
+    } else {
+        AgeGroupRisk::OldAdult
+    }
 }
 ```
-## Births and deaths
-
-### Births
-Some requirements are necessary to include births in the simulation. Namely,
-  * Newborns increase the population,
-  * Person properties are set for newborns at the time of creation, including derived properties,
-  * Newborns become available to look up and should be considered alive after their time of birth,
-  * A person created events should be emitted.
-
-### Deaths
-Requirements for deaths include removing from simulation and canceling any plans for the `person_id`.
-  * Deaths reduce current population, but not for the purposes of the person id for next new born. This means that a counter for total population is required as well as newborns and deaths,
-  * Alive person property should be set to `false`,
-  * All plans should be canceled for `person_id` when they are removed from population. This should happen inside `people.rs` so that modules aren't required to continuously observe for death events,
-  * Death people should not be counted for the force of infection or other transmission events.
-
-
-# Transmission manager
-Infections are spread throughout the population based on a constant force of infection, which differs for age groups 0-12m, 1-65, and 65+. Infection attempts are scheduled based on each age group force of infection. This requires the implementation of an Ixa functionality to look up individuals based on their current age.
-
-```rust transmission_manager.rs
-fn schedule_infection(context, age_group, foi_age) {
-    transmission_rng = rng.get_rng(id = transmission);
-    population = context.get_population(age_group);
-    person_to_infect = context.get_random_person(age_group);
-
-    if (context.get_infection_status(person_to_infect) == Susceptible) {
-        context.set_infection_status(person_to_infect, Infected);
+
+
+## Person look-up based on properties
+This example implements a function to sample a random person from a group of people with the same person property. For instance, to sample a random person who's alive, one can filter by the Alive property `sample_person_by_property(Alive, true)`. A similar function is implemented to select a random person from a specific age group. For instance, to sample someone a Newborn, one can call the function  `sample_person(AgeGroupRisk::NewBorn`.
+
+```rust
+fn sample_person_by_property<T: PersonProperty + 'static>(
+    &mut self,
+    property: T,
+    value: T::Value,
+) -> Option<PersonId>
+where
+     <T as PersonProperty>::Value: PartialEq,
+{
+	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, property) == value {
+            people_vec.push(person_id);
+        }
     }
+    if people_vec.is_empty() {
+        None
+    } else {
+        Some(people_vec[self.sample_range(PeopleRng, 0..people_vec.len())])
+    }
+}
 
-    time_next_infection = transmission_rng.draw_exponential(foi_age) / population;
-    context.add_plan(attempt_infection(context, age_group, foi_age), time = context.get_time() + time_next_infection);
+fn sample_person(&mut self, age_group: AgeGroupRisk) -> Option<PersonId> {
+    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)
+            && self.get_person_age_group(person_id) == age_group
+        {
+            people_vec.push(person_id);
+        }
+    }
+    if people_vec.is_empty() {
+        None
+    } else {
+        Some(people_vec[self.sample_range(PeopleRng, 0..people_vec.len())])
+    }
 }
+```
+
+# Transmission  & infection progression
+Infections are spread throughout the population based on a constant force of infection, which differs for age groups 0-12m, 1-65, and 65+. Infection attempts are scheduled based on each age group force of infection. To spread the pathogen in the population, a random person is selected for each age group using `sample_person(age_group)`, if this person is susceptible to infection.
+
+Infected individuals are scheduled to recover based on the infection period. These are the only type of plans that are scheduled for an individual in this simulation. Hence, when recovery is scheduled using `context.add_plan()`, the `plan id` is stored in a data container named `InfectionPlansPlugin`.
 
-//initialization
-init(context) {
-    context.add_rng(id = transmission);
-    age_groups = parameters.age_groups;
-    vec_foi_age = parameters.foi_age;
-    for n in range(vec_foi_age) {
-        context.add_plan(attempt_infection(context, age_groups[n], vec_foi_age[n]), time = 0);
+```rust
+let plan_id = context
+    .add_plan(recovery_time, move |context| {
+context.set_person_property(person_id, InfectionStatusType, InfectionStatus::R);
+    })
+    .clone();
+let plans_data_container = context.get_data_container_mut(InfectionPlansPlugin);
+plans_data_container
+    .plans_map
+    .entry(person_id)
+    .or_default()
+    .insert(plan_id.clone());
+
+```
+
+These plan ids are removed from the data container once the individual recovers form infection or dies. However, if the person dies during the simulation, the upcoming plans need to be canceled; hence, a special function is used to handle person removal. This function cancels plans to recover and removes their ids from the data container for the person.
+
+```rust
+fn cancel_recovery_plans(context: &mut Context, person_id: PersonId) {
+    let plans_data_container = context.get_data_container_mut(InfectionPlansPlugin);
+    let plans_set = plans_data_container
+        .plans_map
+        .get(&person_id)
+        .unwrap_or(&HashSet::<plan::Id>::new())
+        .clone();
+
+    for plan_id in plans_set {
+        context.cancel_plan(&plan_id);
     }
+    remove_recovery_plan_data(context, person_id);
 }
 ```

examples/births-deaths/incidence_report.rs

@@ -6,10 +6,9 @@ use ixa::{create_report_trait, report::Report};
 use std::path::Path;
 use std::path::PathBuf;
 
-use crate::population_manager::AgeGroupRisk;
-use crate::population_manager::ContextPopulationExt;
-use crate::population_manager::InfectionStatus;
-use crate::population_manager::InfectionStatusType;
+use crate::population_manager::{
+    AgeGroupRisk, ContextPopulationExt, InfectionStatus, InfectionStatusType,
+};
 
 use serde::{Deserialize, Serialize};
 

examples/births-deaths/infection_manager.rs

@@ -33,7 +33,6 @@ fn schedule_recovery(context: &mut Context, person_id: PersonId) {
         + context.sample_distr(InfectionRng, Exp::new(1.0 / infection_duration).unwrap());
 
     if context.get_person_property(person_id, Alive) {
-        println!("Person {person_id:?} infected, scheduling recovery");
         let plan_id = context
             .add_plan(recovery_time, move |context| {
                 context.set_person_property(person_id, InfectionStatusType, InfectionStatus::R);
@@ -45,19 +44,15 @@ fn schedule_recovery(context: &mut Context, person_id: PersonId) {
             .entry(person_id)
             .or_default()
             .insert(plan_id.clone());
-
-        println!("Person {:?} - plan Id: {:?}", person_id, plan_id.clone());
     }
 }
 
 fn remove_recovery_plan_data(context: &mut Context, person_id: PersonId) {
     let plans_data_container = context.get_data_container_mut(InfectionPlansPlugin);
     plans_data_container.plans_map.remove(&person_id);
-    println!("Person {person_id:?} - plans removed");
 }
 
 fn cancel_recovery_plans(context: &mut Context, person_id: PersonId) {
-    println!("Attempting to cancel plans for Person {person_id:?}");
     let plans_data_container = context.get_data_container_mut(InfectionPlansPlugin);
     let plans_set = plans_data_container
         .plans_map
@@ -66,7 +61,6 @@ fn cancel_recovery_plans(context: &mut Context, person_id: PersonId) {
         .clone();
 
     for plan_id in plans_set {
-        println!("Canceling plan {:?}", plan_id.clone());
         context.cancel_plan(&plan_id);
     }
 
@@ -81,7 +75,6 @@ fn handle_infection_status_change(
         schedule_recovery(context, event.person_id);
     }
     if matches!(event.current, InfectionStatus::R) {
-        println!("Person {:?} has recovered", event.person_id);
         remove_recovery_plan_data(context, event.person_id);
     }
 }
@@ -102,3 +95,85 @@ pub fn init(context: &mut Context) {
         handle_person_removal(context, event);
     });
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::population_manager::ContextPopulationExt;
+    use ixa::context::Context;
+    use ixa::define_data_plugin;
+    use ixa::global_properties::ContextGlobalPropertiesExt;
+    use ixa::people::{ContextPeopleExt, PersonPropertyChangeEvent};
+    use ixa::random::ContextRandomExt;
+
+    use crate::parameters_loader::{FoiAgeGroups, ParametersValues};
+    define_data_plugin!(RecoveryPlugin, usize, 0);
+    define_data_plugin!(PlansPlugin, usize, 0);
+
+    #[test]
+    fn test_handle_infection_change_with_deaths() {
+        let p_values = ParametersValues {
+            population: 10,
+            max_time: 10.0,
+            seed: 42,
+            birth_rate: 0.0,
+            death_rate: 0.1,
+            foi_groups: Vec::<FoiAgeGroups>::new(),
+            infection_duration: 5.0,
+            output_file: ".".to_string(),
+            demographic_output_file: ".".to_string(),
+        };
+
+        let mut context = Context::new();
+
+        context.set_global_property_value(Parameters, p_values.clone());
+        context.init_random(p_values.seed);
+        init(&mut context);
+
+        context.subscribe_to_event(
+            move |context, event: PersonPropertyChangeEvent<InfectionStatusType>| {
+                if matches!(event.current, InfectionStatus::R) {
+                    *context.get_data_container_mut(RecoveryPlugin) += 1;
+                }
+            },
+        );
+
+        let population_size: usize = 10;
+        for _ in 0..population_size {
+            let person = context.create_new_person(0.0);
+
+            context.add_plan(1.0, move |context| {
+                context.set_person_property(person, InfectionStatusType, InfectionStatus::I);
+            });
+        }
+
+        context.add_plan(1.1, move |context| {
+            context.attempt_death(context.get_person_id(0));
+        });
+
+        context.execute();
+        assert_eq!(population_size, context.get_current_population());
+        let recovered_size: usize = *context.get_data_container(RecoveryPlugin).unwrap();
+
+        assert_eq!(recovered_size, population_size - 1);
+    }
+
+    #[test]
+    fn test_cancel_null_plan() {
+        let mut context = Context::new();
+
+        context.init_random(42);
+        init(&mut context);
+
+        let person = context.create_new_person(0.0);
+        context.add_plan(1.1, move |context| {
+            cancel_recovery_plans(context, person);
+        });
+
+        context.add_plan(1.2, move |context| {
+            cancel_recovery_plans(context, person);
+        });
+
+        context.execute();
+    }
+}

examples/births-deaths/population_manager.rs

@@ -126,12 +126,6 @@ pub trait ContextPopulationExt {
 
 impl ContextPopulationExt for Context {
     fn attempt_death(&mut self, person_id: PersonId) {
-        // Where should we assign all the person properties to be dead and cancel plans? people.rs?
-        println!(
-            "Attempting to Kill {:?} - at time: {:?}",
-            person_id,
-            self.get_current_time()
-        );
         self.set_person_property(person_id, Alive, false);
     }
 
@@ -157,9 +151,6 @@ impl ContextPopulationExt for Context {
         (self.get_current_time() - birth_time) / 365.0
     }
     fn get_current_group_population(&mut self, age_group: AgeGroupRisk) -> usize {
-        // loop through all population
-        // filter those who are alive
-        // filter those with age group risk = age_group
         let mut current_population = 0;
         for i in 0..self.get_current_population() {
             let person_id = self.get_person_id(i);
@@ -233,7 +224,7 @@ impl ContextPopulationExt for Context {
 mod test {
     use super::*;
     use crate::parameters_loader::{FoiAgeGroups, ParametersValues};
-    use ixa::context::{self, Context};
+    use ixa::context::Context;
 
     #[test]
     fn test_birth_death() {