use ixa::prelude::*;
use rand_distr::Exp;
use crate::{InfectionStatus, InfectionStatusValue, Parameters};
define_rng!(TransmissionRng);
fn attempt_infection(context: &mut Context) {
let population_size: usize = context.get_current_population();
let person_to_infect = context.sample_person(TransmissionRng, ()).unwrap();
let person_status: InfectionStatusValue =
context.get_person_property(person_to_infect, InfectionStatus);
let parameters = context
.get_global_property_value(Parameters)
.unwrap()
.clone();
if person_status == InfectionStatusValue::S {
context.set_person_property(person_to_infect, InfectionStatus, InfectionStatusValue::I);
}
#[allow(clippy::cast_precision_loss)]
let next_attempt_time = context.get_current_time()
+ context.sample_distr(TransmissionRng, Exp::new(parameters.foi).unwrap())
/ population_size as f64;
if next_attempt_time <= parameters.max_time {
context.add_plan(next_attempt_time, move |context| {
attempt_infection(context);
});
}
}
pub fn init(context: &mut Context) {
context.add_plan(0.0, |context| {
attempt_infection(context);
});
}
#[cfg(test)]
mod test {
use super::*;
use crate::parameters_loader::ParametersValues;
#[test]
fn test_attempt_infection() {
let p_values = ParametersValues {
population: 10,
max_time: 10.0,
seed: 42,
foi: 0.15,
infection_duration: 5.0,
output_dir: ".".to_string(),
output_file: ".".to_string(),
};
let mut context = Context::new();
context
.set_global_property_value(Parameters, p_values)
.unwrap();
context.init_random(123);
let pid = context.add_person(()).unwrap();
attempt_infection(&mut context);
let person_status = context.get_person_property(pid, InfectionStatus);
assert_eq!(person_status, InfectionStatusValue::I);
context.execute();
}
}