diff --git a/game/src/common/overlays.rs b/game/src/common/overlays.rs index c4bfa54e5d..52d227d6dc 100644 --- a/game/src/common/overlays.rs +++ b/game/src/common/overlays.rs @@ -836,6 +836,7 @@ impl Overlays { Some(PandemicModel::calculate( app.primary.sim.get_analytics(), + app.primary.sim.get_all_people(), app.primary.sim.time(), &mut XorShiftRng::from_seed([42; 16]), )) diff --git a/sim/src/pandemic.rs b/sim/src/pandemic.rs index 01b36f855c..98a6e4ec27 100644 --- a/sim/src/pandemic.rs +++ b/sim/src/pandemic.rs @@ -1,17 +1,21 @@ -use crate::{Analytics, PersonID}; +use crate::{Analytics, Person, PersonID}; use geom::{Duration, Time}; use map_model::BuildingID; use rand::Rng; use rand_xorshift::XorShiftRng; -use std::collections::{BTreeMap, BTreeSet}; +use std::cmp::Ordering; +use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; pub struct PandemicModel { pub infected: BTreeSet, + hospitalized: BTreeSet, // Since when has a person been inside a building? // TODO This is an awkward data structure; abstutil::MultiMap is also bad, because key removal // would require knowing the time. Want something closer to // https://guava.dev/releases/19.0/api/docs/com/google/common/collect/Table.html. bldg_occupants: BTreeMap>, + + events: BinaryHeap, } impl PandemicModel { @@ -23,66 +27,121 @@ impl PandemicModel { // This recomputes everything every time the UI asks for it. That's fine for the scale of // simulations now; everything else in Analytics works the same way. The faster streaming // version is very straightforward -- cache this output and only process new events. - pub fn calculate(analytics: &Analytics, now: Time, rng: &mut XorShiftRng) -> PandemicModel { + pub fn calculate( + analytics: &Analytics, + population: &Vec, + now: Time, + rng: &mut XorShiftRng, + ) -> PandemicModel { let mut state = PandemicModel { infected: BTreeSet::new(), + hospitalized: BTreeSet::new(), bldg_occupants: BTreeMap::new(), + events: BinaryHeap::new(), }; - // Track people's movements through buildings + // Seed initially infected people. + for p in population { + if rng.gen_bool(0.1) { + state.infected.insert(p.id); + + if rng.gen_bool(0.1) { + state.events.push(Item { + time: Time::START_OF_DAY + + rand_duration(rng, Duration::hours(1), Duration::hours(3)), + event: Event::Hospitalized(p.id), + }); + } + } + } + + // Seed events in the pandemic model from the traffic simulaton. for (time, person, bldg, left) in &analytics.building_transitions { if *time > now { break; } - if *left { - // TODO Messy to mutate state inside a retain closure - let mut inside_since: Option