docs: Suggestions: network model

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<PersonId> = 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();
         }

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<File> {
     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<PersonId> {
+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::<Person, Id>();
     context.index_property::<Person, HouseholdId>();
 
-    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))
-        ));
     }
 }

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<PersonId> = 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();
 }

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<Person>), Edge);
+define_property!(struct Node2(EntityId<Person>), Edge);
 
-#[derive(Deserialize, Debug)]
-struct EdgeRecord {
-    v1: u16,
-    v2: u16,
+define_rng!(NetworkRng);
+
+fn add_bidi_edge(context: &mut Context, p1: EntityId<Person>, p2: EntityId<Person>, 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<EntityId<Person>> = 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<PersonId> = Vec::new();
-            context.with_query_results((household_id,), &mut |results| {
-                members = results.to_owned_vec()
-            });
+            let members: Vec<EntityId<Person>> =  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<ET: EdgeType<Person>>(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<EntityId<Person>> = 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<EntityId<Person>> = 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<Person>, duration: f64) -> Vec<EntityId<Person>> {
+    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<EntityId<Person>> = 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::<Person, Household>(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::<Person, Household>(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::<Person, Household>(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);
     }
 }