work in progress

Author: erik-rosenstrom

src/infection_propagation_loop.rs

@@ -6,8 +6,9 @@ use crate::infectiousness_manager::{
     Forecast, InfectionContextExt, InfectionStatus, evaluate_forecast, get_forecast,
     infection_attempt,
 };
+use crate::itinerary::{ContextItineraryExt, ItineraryUse};
 use crate::parameters::{ContextParametersExt, Params};
-use crate::population_loader::{Person, PersonId};
+use crate::population_loader::{IsolationStatus, Person, PersonId};
 use crate::rate_fns::{InfectiousnessRateExt, load_rate_fns};
 use ixa::profiling::{increment_named_count, open_span};
 use ixa::{Context, ContextRandomExt, IxaError, define_rng, trace};
@@ -80,6 +81,8 @@ pub fn init(context: &mut Context) -> Result<(), IxaError> {
     load_rate_fns(context)?;
     seed_initial_infections(context, initial_incidence);
 
+    context.register_person_interary_modifier(IsolationStatus::Isolated, ItineraryUse::Isolated, 1, Some(1.0), None, None, None);
+
     // Subscribe to the person becoming infectious to trigger the infection propagation loop
     context.subscribe_to_event(
         |context, event: PropertyChangeEvent<Person, InfectionStatus>| {
@@ -88,6 +91,21 @@ pub fn init(context: &mut Context) -> Result<(), IxaError> {
             }
             schedule_next_forecasted_infection(context, event.entity_id);
             schedule_recovery(context, event.entity_id);
+
+            context.add_plan(context.get_current_time() + 2.0, move |context| {
+                context.set_property::<Person, IsolationStatus>(event.entity_id, IsolationStatus::Isolated);
+            });
+        },
+    );
+
+    context.subscribe_to_event(
+        |context, event: PropertyChangeEvent<Person, InfectionStatus>| {
+            if event.current != InfectionStatus::Recovered {
+                context.add_plan(context.get_current_time(), move |context| {
+                context.set_property::<Person, IsolationStatus>(event.entity_id, IsolationStatus::NotIsolated);
+            });
+            }
+            
         },
     );
 

src/itinerary.rs

@@ -125,29 +125,61 @@ pub trait ContextItineraryExt: PluginContext + ContextRandomExt {
                         BelongsTo(person_id),
                         itinerary_use,
                         Ranking(ranking),
-                        Activity(true)
-                        HomeItinerary {
+                        Activity(true))).unwrap();
+                    if let Some(home_ratio) = home_ratio {
+                        context.set_property::<Itinerary, HomeItinerary>(modified_itinerary, HomeItinerary {
                             home_id: home_id.0,
-                            ratio: home_ratio,
-                        },
-                        SchoolItinerary {
-                            school_id: school_id.0,
-                            ratio: school_ratio,
-                        },
-                        WorkplaceItinerary {
+                            ratio: Some(home_ratio),
+                        });
+                    } else {
+                        context.set_property::<Itinerary, HomeItinerary>(modified_itinerary, HomeItinerary {
+                            home_id: None,
+                            ratio: None,
+                        });
+                    }
+                    if let Some(workplace_ratio) = workplace_ratio {
+                        context.set_property::<Itinerary, WorkplaceItinerary>(modified_itinerary, WorkplaceItinerary {
                             workplace_id: workplace_id.0,
-                            ratio: workplace_ratio,
-                        },
-                        CensusTractItinerary {
+                            ratio: Some(workplace_ratio),
+                        });
+                    } else {
+                        context.set_property::<Itinerary, WorkplaceItinerary>(modified_itinerary, WorkplaceItinerary {
+                            workplace_id: None,
+                            ratio: None,
+                        });
+                    }
+                    if let Some(school_ratio) = school_ratio {
+                        context.set_property::<Itinerary, SchoolItinerary>(modified_itinerary, SchoolItinerary {
+                            school_id: school_id.0,
+                            ratio: Some(school_ratio),
+                        });
+                    } else {
+                        context.set_property::<Itinerary, SchoolItinerary>(modified_itinerary, SchoolItinerary {
+                            school_id: None,
+                            ratio: None,
+                        });
+                    }
+                    if let Some(census_tract_ratio) = census_tract_ratio {                
+                    context.set_property::<Itinerary, CensusTractItinerary>(modified_itinerary, CensusTractItinerary {
                             census_tract_id: census_tract_id.0,
-                            ratio: census_tract_ratio,
-                        },
-                    ))
-                    .unwrap();
-    context.normalize_itinerary_ratios(person_id); 
+                            ratio: Some(census_tract_ratio),
+                    });
+                    } else {
+                        context.set_property::<Itinerary, CensusTractItinerary>(modified_itinerary, CensusTractItinerary {
+                            census_tract_id: None,
+                            ratio: None,
+                        });
+                    }
+                    context.normalize_itinerary_ratios(person_id); 
                 }
+                
                 if event.previous == entity_property {
-
+                    let active_itineraries = context.query_result_iterator::<Itinerary, _>((BelongsTo(event.entity_id), Activity(true))).collect::<Vec<_>>();
+                    for itinerary in active_itineraries {
+                        context.set_property::<Itinerary, Activity>(itinerary, Activity(false));
+                    }
+                    let default_itinerary = context.query_result_iterator::<Itinerary, _>((BelongsTo(event.entity_id), ItineraryUse::Default,)).next().unwrap();
+                    context.set_property::<Itinerary, Activity>(default_itinerary, Activity(true));
                 }
             },
         );

src/population_loader.rs

@@ -33,6 +33,13 @@ impl_property!(Age, Person);
 pub struct Alive(pub bool);
 impl_property!(Alive, Person, default_const = Alive(true));
 
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Hash)]
+pub enum IsolationStatus {
+    Isolated,
+    NotIsolated,
+}
+impl_property!(IsolationStatus, Person, default_const = IsolationStatus::NotIsolated);
+
 fn create_person_from_record(
     context: &mut Context,
     person_record: &PeopleRecord,