diff --git a/src/itinerary_manager.rs b/src/itinerary_manager.rs index fb4fcd0..d38808e 100644 --- a/src/itinerary_manager.rs +++ b/src/itinerary_manager.rs @@ -5,7 +5,7 @@ use std::{ collections::HashMap, }; -use crate::population_loader::{Person, PersonId}; +use crate::{Age, error::ModelError, population_loader::{Alive, Person, PersonId}}; use crate::settings::ItineraryRatios; #[derive(Debug, PartialEq, Clone, Serialize, Copy, Eq, Hash)] @@ -187,6 +187,7 @@ pub trait ContextItineraryModifierExt: PluginContext + ContextEntitiesExt { fn get_itinerary_modifiers(&self, person_id: PersonId) -> Vec; fn get_dominant_itinerary_modifier(&self, person_id: PersonId) -> Option; + fn get_dominant_itinerary_modifier_ratios(&self, person_id: PersonId) -> ItineraryRatios; } impl ContextItineraryModifierExt for Context { // This needs to be here to have access to the concrete context type for the get_itinerary trait method @@ -206,6 +207,51 @@ impl ContextItineraryModifierExt for Context { fn get_dominant_itinerary_modifier(&self, person_id: PersonId) -> Option { self.get_itinerary_modifiers(person_id).into_iter().max() } + + fn get_dominant_itinerary_modifier_ratios(&self, person_id: PersonId) -> ItineraryRatios { + self.get_dominant_itinerary_modifier(person_id) + .map(|modifier| modifier.itinerary_ratios) + .unwrap_or(self.get_property::(person_id)) + } + +} + +pub fn init(context: &mut Context) -> Result<(), ModelError> { + let weekend_modifier = ItineraryModifier { + ranking: 2, + itinerary_ratios: ItineraryRatios { + itinerary_ratios: [0.5, 0.0, 0.0, 0.5], + }, + modifier_type: ItineraryModifierType::Weekend, + }; + + let school_closure_modifier = ItineraryModifier { + ranking: 1, + itinerary_ratios: ItineraryRatios { + itinerary_ratios: [0.75, 0.0, 0.0, 0.25], + }, + modifier_type: ItineraryModifierType::SchoolClosure, + }; + + for i in (0..=100).step_by(7) { + context.add_plan(i as f64, move |context| { + context.register_itinerary_modifier(Alive(true), weekend_modifier); + }); + context.add_plan(i as f64 + 2.0, move |context| { + context.remove_itinerary_modifier_by_property::(Alive(true)); + }); + } + + for age in 5..=18 { + context.add_plan(3.0, move |context| { + context.register_itinerary_modifier(Age(age), school_closure_modifier); + }); + context.add_plan(30.0, move |context| { + context.remove_itinerary_modifier_by_property::(Age(age)); + }); + } + + Ok(()) } #[cfg(test)] diff --git a/src/model.rs b/src/model.rs index 8b63965..e57b94a 100644 --- a/src/model.rs +++ b/src/model.rs @@ -3,8 +3,7 @@ use std::path::PathBuf; use ixa::{ExecutionPhase, prelude::*}; use crate::{ - abort_run, error::ModelError, infection_importation, infection_propagation_loop, parameters, - population_loader, reports, settings, symptom_status_manager, + abort_run, error::ModelError, infection_importation, infection_propagation_loop, itinerary_manager, parameters, population_loader, reports, settings, symptom_status_manager }; pub fn initialize_model( @@ -37,6 +36,8 @@ pub fn initialize_model( info!("Infection propagation loop initialized"); infection_importation::init(context)?; info!("Infection importation initialized"); + itinerary_manager::init(context)?; + info!("Itinerary manager initialized"); reports::init(context)?; info!("Reports initialized"); abort_run::init(context); diff --git a/src/settings.rs b/src/settings.rs index 7623c9a..0ff432d 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -7,6 +7,7 @@ use std::{ }; use strum::{EnumCount as EnumCountMacro, EnumIter, IntoEnumIterator}; +use crate::itinerary_manager::ContextItineraryModifierExt; pub(crate) use crate::{ ContextParametersExt, Params, error::ModelError, @@ -126,7 +127,7 @@ impl ContextSettingExtPrivate for Context {} #[allow(private_bounds)] pub trait ContextSettingExt: - PluginContext + ContextEntitiesExt + ContextSettingExtPrivate + ContextParametersExt + PluginContext + ContextEntitiesExt + ContextSettingExtPrivate + ContextParametersExt + ContextItineraryModifierExt { fn register_setting_global_properties(&mut self) { let Params { @@ -179,7 +180,7 @@ pub trait ContextSettingExt: ) -> Result, ModelError> { let mut active_settings = Vec::new(); let setting_ids = self.get_property::(person_id); - let itinerary_ratios = self.get_property::(person_id); + let itinerary_ratios = self.get_dominant_itinerary_modifier_ratios(person_id); for category in SettingCategory::iter() { if let Some(id) = setting_ids.setting_ids[category] {