diff --git a/examples/network-hhmodel/incidence_report.rs b/examples/network-hhmodel/incidence_report.rs index 874d9232..7aad3ae9 100644 --- a/examples/network-hhmodel/incidence_report.rs +++ b/examples/network-hhmodel/incidence_report.rs @@ -146,7 +146,7 @@ mod test { .unwrap(); let people = loader::init(&mut context); - network::init(&mut context, &people); + network::init(&mut context); incidence_report::init(&mut context).unwrap(); context.subscribe_to_event( @@ -158,7 +158,7 @@ mod test { let to_infect: Vec = vec![context.sample_entity(MainRng, Person).unwrap()]; #[allow(clippy::vec_init_then_push)] - seir::init(&mut context, &to_infect); + seir::init(&mut context, &to_infect, 1.0); context.execute(); } diff --git a/examples/network-hhmodel/loader.rs b/examples/network-hhmodel/loader.rs index f3b0220c..7e980b76 100644 --- a/examples/network-hhmodel/loader.rs +++ b/examples/network-hhmodel/loader.rs @@ -5,7 +5,7 @@ use ixa::impl_property; use ixa::prelude::*; use serde::{Deserialize, Serialize}; -use crate::{example_dir, Person, PersonId}; +use crate::{example_dir, Person}; #[derive(Serialize, Deserialize, Copy, Clone, PartialEq, Eq, Debug, Hash)] pub struct Id(pub u16); @@ -39,38 +39,30 @@ struct PeopleRecord { household_id: HouseholdId, } -fn create_person_from_record(context: &mut Context, record: &PeopleRecord) -> PersonId { - context - .add_entity((record.id, record.age_group, record.sex, record.household_id)) - .unwrap() -} - pub fn open_csv(file_name: &str) -> Reader { let current_dir = example_dir(); let file_path = current_dir.join(file_name); csv::Reader::from_path(file_path).unwrap() } -pub fn init(context: &mut Context) -> Vec { +pub fn init(context: &mut Context) { // Load csv and deserialize records let mut reader = open_csv("Households.csv"); - let mut people = Vec::new(); for result in reader.deserialize() { let record: PeopleRecord = result.expect("Failed to parse record"); - people.push(create_person_from_record(context, &record)); + let _ = context.add_entity((record.id, record.age_group, record.sex, record.household_id)); } context.index_property::(); context.index_property::(); - people } #[cfg(test)] mod tests { use ixa::context::Context; - use ixa::random::ContextRandomExt; + use ixa::random::{ContextRandomExt}; use super::*; @@ -87,26 +79,17 @@ mod tests { #[test] fn test_some_people_load_correctly() { let mut context = Context::new(); - context.init_random(42); + init(&mut context); + + // only one with the given id + assert_eq!(context.query_entity_count(with!(Person, Id(676))), 1); + assert_eq!(context.query_entity_count(with!(Person, Id(213))), 1); + assert_eq!(context.query_entity_count(with!(Person, Id(1591))), 1); + + // test exist fully specified + assert_eq!(context.query_entity_count((Id(676), AgeGroup::Age18to64, Sex::Female, HouseholdId(1))), 1); + assert_eq!(context.query_entity_count((Id(213), AgeGroup::AgeUnder5, Sex::Female, HouseholdId(162))), 1); + assert_eq!(context.query_entity_count((Id(1591), AgeGroup::Age65Plus, Sex::Male, HouseholdId(496))), 1); - let people = init(&mut context); - - let person = people[0]; - assert!(context.match_entity( - person, - (Id(676), AgeGroup::Age18to64, Sex::Female, HouseholdId(1)) - )); - - let person = people[246]; - assert!(context.match_entity( - person, - (Id(213), AgeGroup::AgeUnder5, Sex::Female, HouseholdId(162)) - )); - - let person = people[1591]; - assert!(context.match_entity( - person, - (Id(1591), AgeGroup::Age65Plus, Sex::Male, HouseholdId(496)) - )); } } diff --git a/examples/network-hhmodel/main.rs b/examples/network-hhmodel/main.rs index fee887dc..ee730082 100644 --- a/examples/network-hhmodel/main.rs +++ b/examples/network-hhmodel/main.rs @@ -27,14 +27,14 @@ fn initialize(context: &mut Context) { context.init_random(1); // Load people from csv and set up some base properties - let people = loader::init(context); + loader::init(context); // Load parameters from json let file_path = example_dir().join("config.json"); context.load_global_properties(&file_path).unwrap(); // Load network - network::init(context, &people); + network::init(context); // Initialize incidence report incidence_report::init(context).unwrap(); @@ -43,6 +43,6 @@ fn initialize(context: &mut Context) { let to_infect: Vec = vec![context.sample_entity(MainRng, Person).unwrap()]; #[allow(clippy::vec_init_then_push)] - seir::init(context, &to_infect); + seir::init(context, &to_infect, 1.0); context.execute(); } diff --git a/examples/network-hhmodel/network.rs b/examples/network-hhmodel/network.rs index b5e1cac8..dd067f42 100644 --- a/examples/network-hhmodel/network.rs +++ b/examples/network-hhmodel/network.rs @@ -1,111 +1,151 @@ -use ixa::network::edge::EdgeType; use ixa::prelude::*; use ixa::{HashSet, HashSetExt}; -use serde::Deserialize; +use rand_distr::Bernoulli; use crate::loader::{open_csv, HouseholdId, Id}; -use crate::{Person, PersonId}; +use crate::parameters::Parameters; +use crate::Person; -define_edge_type!(struct Household, Person); -define_edge_type!(struct AgeUnder5, Person); -define_edge_type!(struct Age5to17, Person); +define_entity!(Edge); +// relative transmission rate +define_property!(struct RR(f64), Edge); +define_property!(struct Node1(EntityId), Edge); +define_property!(struct Node2(EntityId), Edge); -#[derive(Deserialize, Debug)] -struct EdgeRecord { - v1: u16, - v2: u16, +define_rng!(NetworkRng); + +fn add_bidi_edge(context: &mut Context, p1: EntityId, p2: EntityId, rr: RR) { + context + .add_entity((Node1(p1), Node2(p2), rr)) + .unwrap(); + context + .add_entity((Node1(p2), Node2(p1), rr)) + .unwrap(); } -fn create_household_networks(context: &mut Context, people: &[PersonId]) { +fn create_household_networks(context: &mut Context, rr: RR) { let mut households = HashSet::new(); - for person_id in people { - let household_id: HouseholdId = context.get_property(*person_id); + let people: Vec> = context.get_entity_iterator().collect(); + for person in people { + let household_id: HouseholdId = context.get_property(person); if households.insert(household_id) { - let mut members: Vec = Vec::new(); - context.with_query_results((household_id,), &mut |results| { - members = results.to_owned_vec() - }); + let members: Vec> = context.query((household_id,)).into_iter().collect(); + // create a dense network - while let Some(person) = members.pop() { - for other_person in &members { - context - .add_edge_bidi(person, *other_person, 1.0, Household) - .unwrap(); + for i in 0..(members.len() - 1) { + for j in (i + 1)..members.len() { + debug_assert!(i != j); + add_bidi_edge(context, members[i], members[j], rr); } } } } } -fn load_edge_list>(context: &mut Context, file_name: &str, inner: ET) { +fn load_edge_list(context: &mut Context, file_name: &str, rr: RR) { let mut reader = open_csv(file_name); for result in reader.deserialize() { - let record: EdgeRecord = result.expect("Failed to parse edge"); - let mut p1_vec = Vec::new(); - context.with_query_results((Id(record.v1),), &mut |people| { - p1_vec = people.to_owned_vec() - }); + let record: (u16, u16) = result.expect("Failed to parse edge"); + let p1_vec: Vec> = context.query(((Id(record.0)),)).into_iter().collect(); assert_eq!(p1_vec.len(), 1); - let p1 = p1_vec[0]; - let mut p2_vec = Vec::new(); - context.with_query_results((Id(record.v2),), &mut |people| { - p2_vec = people.to_owned_vec() - }); + let p2_vec: Vec> = context.query(((Id(record.1)),)).into_iter().collect(); assert_eq!(p2_vec.len(), 1); - let p2 = p2_vec[0]; - context.add_edge_bidi(p1, p2, 1.0, inner.clone()).unwrap(); + add_bidi_edge(context, p1_vec[0], p2_vec[0], rr); } } -pub fn init(context: &mut Context, people: &[PersonId]) { - // Create dense household networks - create_household_networks(context, people); +fn sar_to_beta(sar: f64, infectious_period: f64) -> f64 { + 1.0 - (1.0 - sar).powf(1.0 / infectious_period) +} + +/// Get all the effective contacts a person will have over a certain duration +pub fn get_contacts(context: &Context, person: EntityId, duration: f64) -> Vec> { + let parameters = context + .get_global_property_value(Parameters) + .unwrap() + .clone(); + + // Base probability of contact during the duration. Note that this assumes that the duration is not too high! + let base_p = duration * sar_to_beta(parameters.sar, parameters.incubation_period); + + let mut contacts: Vec> = Vec::new(); + + for edge in context.query((Node1(person), )) { + let RR(rr): RR = context.get_property(edge); + let Node2(person2): Node2 = context.get_property(edge); + + if context.sample_distr(NetworkRng, Bernoulli::new(base_p * rr).unwrap()) { + if !contacts.contains(&person2) { + contacts.push(person2); + } + } + } - // Add U5 edges from csv - load_edge_list(context, "AgeUnder5Edges.csv", AgeUnder5); + contacts +} + +pub fn init(context: &mut Context) { + let parameters = context + .get_global_property_value(Parameters) + .unwrap() + .clone(); + + // relative risk of transmission between (vs. within) households + let rr = 1.0 / parameters.between_hh_transmission_reduction; - // Add U18 edges from csv - load_edge_list(context, "Age5to17Edges.csv", Age5to17); + // Create dense household networks + create_household_networks(context, RR(1.0)); + // Add other edges from csv's with lower transmission rate + load_edge_list(context, "AgeUnder5Edges.csv", RR(rr)); + load_edge_list(context, "Age5to17Edges.csv", RR(rr)); } #[cfg(test)] mod tests { - use super::*; - use crate::{loader, network}; - - const N_SIZE_12: usize = 1; - const N_SIZE_11: usize = 1; - const N_SIZE_3: usize = 122; + use crate::{loader, network, Person}; + use crate::parameters::ParametersValues; - #[test] - fn test_expected_12_member_household() { + // Assert that person with `id` has `n` contacts (i.e., edges going from + // them, and also edges going to them) + fn assert_has_n_contacts(id: u16, n: usize) { let mut context = Context::new(); context.init_random(42); - let people = loader::init(&mut context); - network::init(&mut context, &people); - let deg11 = context.find_entities_by_degree::(11); - assert_eq!(deg11.len(), 12 * N_SIZE_12); + loader::init(&mut context); + let parameters = ParametersValues { + incubation_period: 8.0, + infectious_period: 27.0, + sar: 1.0, + shape: 15.0, + infection_duration: 5.0, + between_hh_transmission_reduction: 1.0, + data_dir: "examples/network-hhmodel/tests".to_owned(), + output_dir: "examples/network-hhmodel/tests".to_owned(), + }; + context + .set_global_property_value(Parameters, parameters) + .unwrap(); + network::init(&mut context); + + let person = context.query((Id(id),)).into_iter().next().unwrap(); + let n_to = context.query_entity_count((Node1(person),)); + let n_from = context.query_entity_count((Node2(person), )); + assert_eq!(n_to, n); + assert_eq!(n_from, n); } #[test] - fn test_expected_11_member_household() { - let mut context = Context::new(); - context.init_random(42); - let people = loader::init(&mut context); - network::init(&mut context, &people); - let deg10 = context.find_entities_by_degree::(10); - assert_eq!(deg10.len(), 11 * N_SIZE_11); + fn test_person_826() { + // Person 826 is in a household of 5 with no other contacts. + // There should be 4 edges going from them, and 4 going to them. + assert_has_n_contacts(826, 4); } #[test] - fn test_expected_3_member_household() { - let mut context = Context::new(); - context.init_random(42); - let people = loader::init(&mut context); - network::init(&mut context, &people); - let deg10 = context.find_entities_by_degree::(2); - assert_eq!(deg10.len(), 3 * N_SIZE_3); + fn test_person_243() { + // Person 243 is in a household of size 6, with 3 other contacts, + // for 9 total contacts. + assert_has_n_contacts(243, 9); } } diff --git a/examples/network-hhmodel/seir.rs b/examples/network-hhmodel/seir.rs index cc6be56f..ae8bffcb 100644 --- a/examples/network-hhmodel/seir.rs +++ b/examples/network-hhmodel/seir.rs @@ -1,11 +1,10 @@ use ixa::log::info; -use ixa::network::edge::EdgeType; use ixa::prelude::*; -use ixa::{impl_property, ExecutionPhase}; -use rand_distr::{Bernoulli, Gamma}; +use ixa::{impl_property, ExecutionPhase, HashSet, HashSetExt}; +use rand_distr::Gamma; use serde::{Deserialize, Serialize}; -use crate::network::{Age5to17, AgeUnder5, Household}; +use crate::network::get_contacts; use crate::parameters::Parameters; use crate::{Person, PersonId}; @@ -27,35 +26,20 @@ define_property!( default_const = InfectedBy(None) ); -fn sar_to_beta(sar: f64, infectious_period: f64) -> f64 { - 1.0 - (1.0 - sar).powf(1.0 / infectious_period) -} - fn calculate_waiting_time(context: &Context, shape: f64, mean_period: f64) -> f64 { let d = Gamma::new(shape, mean_period / shape).unwrap(); context.sample_distr(SeirRng, d) } -fn expose_network>(context: &mut Context, beta: f64) { - let infectious_people = context.query((DiseaseStatus::I,)).to_owned_vec(); - - for infectious in infectious_people { - let edges = context.get_matching_edges::(infectious, |context, edge| { - context.match_entity(edge.neighbor, (DiseaseStatus::S,)) - }); - - for e in edges { - if context.sample_distr(SeirRng, Bernoulli::new(beta).unwrap()) { - context.set_property(e.neighbor, DiseaseStatus::E); - info!( - "Person {} exposed person {} at time {}.", - infectious, - e.neighbor, - context.get_current_time() - ); - context.set_property(e.neighbor, InfectedBy(Some(infectious))); - } - } +fn expose(context: &mut Context, infector: PersonId, infectee: PersonId) { + let infectee_status: DiseaseStatus = context.get_property(infectee); + if infectee_status == DiseaseStatus::S { + info!( + "{infector:?} exposed {infectee:?} at time {}.", + context.get_current_time() + ); + context.set_property(infectee, DiseaseStatus::E); + context.set_property(infectee, InfectedBy(Some(infector))); } } @@ -66,15 +50,15 @@ fn schedule_waiting_event( mean_period: f64, new_status: DiseaseStatus, ) { - let ct = context.get_current_time(); - let waiting_time = calculate_waiting_time(context, shape, mean_period); + let t = context.get_current_time() + calculate_waiting_time(context, shape, mean_period); - context.add_plan(ct + waiting_time, move |context| { + context.add_plan(t, move |context| { + trace!("{person_id:?} changed to disease state {new_status:?} at t={t:?}"); context.set_property(person_id, new_status); }); } -fn schedule_infection(context: &mut Context, person_id: PersonId) { +fn schedule_infectiousness(context: &mut Context, person_id: PersonId) { let parameters = context .get_global_property_value(Parameters) .unwrap() @@ -104,41 +88,29 @@ fn schedule_recovery(context: &mut Context, person_id: PersonId) { ); } -pub fn init(context: &mut Context, initial_infections: &Vec) { +pub fn init(context: &mut Context, initial_infections: &Vec, period: f64) { context.add_periodic_plan_with_phase( - 1.0, - |context| { - let parameters = context - .get_global_property_value(Parameters) - .unwrap() - .clone(); - - // infect the networks - expose_network::( - context, - sar_to_beta(parameters.sar, parameters.incubation_period), - ); - expose_network::( - context, - sar_to_beta( - parameters.sar / parameters.between_hh_transmission_reduction, - parameters.incubation_period, - ), - ); - expose_network::( - context, - sar_to_beta( - parameters.sar / parameters.between_hh_transmission_reduction, - parameters.incubation_period, - ), - ); + period, + move |context| { + // get all infector-infectee pairs + let mut pairs = HashSet::new(); + for infector in context.query(with!(Person, DiseaseStatus::I)) { + for infectee in get_contacts(context, infector, period) { + pairs.insert((infector, infectee)); + } + } + + // do the exposures + for (infector, infectee) in pairs { + expose(context, infector, infectee) + } }, ExecutionPhase::Normal, ); context.subscribe_to_event( move |context, event: PropertyChangeEvent| match event.current { - DiseaseStatus::E => schedule_infection(context, event.entity_id), + DiseaseStatus::E => schedule_infectiousness(context, event.entity_id), DiseaseStatus::I => schedule_recovery(context, event.entity_id), _ => (), }, @@ -166,7 +138,7 @@ mod tests { context.init_random(42); - let people = loader::init(&mut context); + loader::init(&mut context); // set sar and between_hh_transmission_reduction to 1.0 so that // beta is 1.0 @@ -184,14 +156,14 @@ mod tests { .set_global_property_value(Parameters, parameters) .unwrap(); - network::init(&mut context, &people); + network::init(&mut context); let mut to_infect = Vec::::new(); context.with_query_results((Id(71),), &mut |people| { to_infect.extend(people); }); - init(&mut context, &to_infect); + init(&mut context, &to_infect, 1.0); context.execute();