sim/pandemic/
model.rs

1use std::collections::BTreeMap;
2
3use rand::Rng;
4use rand_xorshift::XorShiftRng;
5use serde::{Deserialize, Serialize};
6
7use geom::{Duration, Time};
8use map_model::{BuildingID, TransitStopID};
9
10use crate::pandemic::{AnyTime, State};
11use crate::{CarID, Event, Person, PersonID, Scheduler, TripPhaseType};
12
13// TODO This does not model transmission by surfaces; only person-to-person.
14// TODO If two people are in the same shared space indefinitely and neither leaves, we don't model
15// transmission. It only occurs when people leave a space.
16
17#[derive(Clone)]
18pub struct PandemicModel {
19    pop: BTreeMap<PersonID, State>,
20
21    bldgs: SharedSpace<BuildingID>,
22    bus_stops: SharedSpace<TransitStopID>,
23    buses: SharedSpace<CarID>,
24    person_to_bus: BTreeMap<PersonID, CarID>,
25
26    rng: XorShiftRng,
27    initialized: bool,
28}
29
30// You can schedule callbacks in the future by doing scheduler.push(future time, one of these)
31#[derive(Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd, Ord, Clone, Debug)]
32pub enum Cmd {
33    BecomeHospitalized(PersonID),
34    BecomeQuarantined(PersonID),
35}
36
37// TODO Pretend handle_event and handle_cmd also take in some object that lets you do things like:
38//
39// - replace_future_trips(PersonID, Vec<IndividTrip>)
40//
41// I'm not exactly sure how this should work yet. Any place you want to change the rest of the
42// simulation, just add a comment describing what you want to do exactly, and we'll figure it out
43// from there.
44
45impl PandemicModel {
46    pub fn new(rng: XorShiftRng) -> PandemicModel {
47        PandemicModel {
48            pop: BTreeMap::new(),
49
50            bldgs: SharedSpace::new(),
51            bus_stops: SharedSpace::new(),
52            buses: SharedSpace::new(),
53            person_to_bus: BTreeMap::new(),
54
55            rng,
56            initialized: false,
57        }
58    }
59
60    // Sorry, initialization order of simulations is still a bit messy. This'll be called at
61    // Time::START_OF_DAY after all of the people have been created from a Scenario.
62    pub(crate) fn initialize(&mut self, population: &[Person], _scheduler: &mut Scheduler) {
63        assert!(!self.initialized);
64        self.initialized = true;
65
66        // Seed initially infected people.
67        // TODO the intial time is not well set. it should start "before"
68        // the beginning of the day. Also
69        for p in population {
70            let state = State::new(0.5, 0.5);
71            let state = if self.rng.gen_bool(State::ini_exposed_ratio()) {
72                let next_state = state
73                    .start(
74                        AnyTime::from(Time::START_OF_DAY),
75                        Duration::seconds(std::f64::MAX),
76                        &mut self.rng,
77                    )
78                    .unwrap();
79                if self.rng.gen_bool(State::ini_infectious_ratio()) {
80                    next_state
81                        .next_default(AnyTime::from(Time::START_OF_DAY), &mut self.rng)
82                        .unwrap()
83                } else {
84                    next_state
85                }
86            } else {
87                state
88            };
89            self.pop.insert(p.id, state);
90        }
91    }
92
93    pub fn count_sane(&self) -> usize {
94        self.pop
95            .iter()
96            .filter(|(_, state)| matches!(state, State::Sane(_)))
97            .count()
98        // self.sane.len()
99    }
100
101    pub fn count_exposed(&self) -> usize {
102        self.pop
103            .iter()
104            .filter(|(_, state)| matches!(state, State::Exposed(_)))
105            .count()
106        // self.exposed.len()
107    }
108
109    pub fn count_infected(&self) -> usize {
110        // self.infected.len()
111        self.pop
112            .iter()
113            .filter(|(_, state)| matches!(state, State::Infectious(_) | State::Hospitalized(_)))
114            .count()
115    }
116
117    pub fn count_recovered(&self) -> usize {
118        self.pop
119            .iter()
120            .filter(|(_, state)| matches!(state, State::Recovered(_)))
121            .count()
122        // self.recovered.len()
123    }
124
125    pub fn count_dead(&self) -> usize {
126        self.pop
127            .iter()
128            .filter(|(_, state)| matches!(state, State::Dead(_)))
129            .count()
130        // self.recovered.len()
131    }
132
133    pub fn count_total(&self) -> usize {
134        self.count_sane()
135            + self.count_exposed()
136            + self.count_infected()
137            + self.count_recovered()
138            + self.count_dead()
139    }
140
141    pub(crate) fn handle_event(&mut self, now: Time, ev: &Event, scheduler: &mut Scheduler) {
142        assert!(self.initialized);
143
144        match ev {
145            Event::PersonEntersBuilding(person, bldg) => {
146                self.bldgs.person_enters_space(now, *person, *bldg);
147            }
148            Event::PersonLeavesBuilding(person, bldg) => {
149                if let Some(others) = self.bldgs.person_leaves_space(now, *person, *bldg) {
150                    self.transmission(now, *person, others, scheduler);
151                } else {
152                    panic!("{} left {}, but they weren't inside", person, bldg);
153                }
154            }
155            Event::TripPhaseStarting(_, p, _, tpt) => {
156                let person = *p;
157                match tpt {
158                    TripPhaseType::WaitingForBus(_, stop) => {
159                        self.bus_stops.person_enters_space(now, person, *stop);
160                    }
161                    TripPhaseType::RidingBus(_, stop, bus) => {
162                        let others = self
163                            .bus_stops
164                            .person_leaves_space(now, person, *stop)
165                            .unwrap();
166                        self.transmission(now, person, others, scheduler);
167
168                        self.buses.person_enters_space(now, person, *bus);
169                        self.person_to_bus.insert(person, *bus);
170                    }
171                    TripPhaseType::Walking => {
172                        // A person can start walking for many reasons, but the only possible state
173                        // transition after riding a bus is walking, so use this to detect the end
174                        // of a bus ride.
175                        if let Some(car) = self.person_to_bus.remove(&person) {
176                            let others = self.buses.person_leaves_space(now, person, car).unwrap();
177                            self.transmission(now, person, others, scheduler);
178                        }
179                    }
180                    _ => {
181                        self.transition(now, person, scheduler);
182                    }
183                }
184            }
185            _ => {}
186        }
187    }
188
189    pub(crate) fn handle_cmd(&mut self, _now: Time, cmd: Cmd, _scheduler: &mut Scheduler) {
190        assert!(self.initialized);
191
192        // TODO Here we might enforce policies. Like severe -> become hospitalized
193        // Symptomatic -> stay quaratined, and/or track contacts to quarantine them too (or test
194        // them)
195        match cmd {
196            Cmd::BecomeHospitalized(_person) => {
197                // self.hospitalized.insert(person);
198            }
199            Cmd::BecomeQuarantined(_person) => {
200                // self.quarantined.insert(person);
201            }
202        }
203    }
204
205    pub fn get_time(&self, person: PersonID) -> Option<Time> {
206        match self.pop.get(&person) {
207            Some(state) => state.get_time(),
208            None => unreachable!(),
209        }
210    }
211
212    pub fn is_sane(&self, person: PersonID) -> bool {
213        match self.pop.get(&person) {
214            Some(state) => state.is_sane(),
215            None => unreachable!(),
216        }
217    }
218
219    pub fn is_infectious(&self, person: PersonID) -> bool {
220        match self.pop.get(&person) {
221            Some(state) => state.is_infectious(),
222            None => unreachable!(),
223        }
224    }
225
226    pub fn is_exposed(&self, person: PersonID) -> bool {
227        match self.pop.get(&person) {
228            Some(state) => state.is_exposed(),
229            None => unreachable!(),
230        }
231    }
232
233    pub fn is_recovered(&self, person: PersonID) -> bool {
234        match self.pop.get(&person) {
235            Some(state) => state.is_recovered(),
236            None => unreachable!(),
237        }
238    }
239
240    pub fn is_dead(&self, person: PersonID) -> bool {
241        match self.pop.get(&person) {
242            Some(state) => state.is_dead(),
243            None => unreachable!(),
244        }
245    }
246
247    fn infectious_contact(&self, person: PersonID, other: PersonID) -> Option<PersonID> {
248        if self.is_sane(person) && self.is_infectious(other) {
249            return Some(person);
250        } else if self.is_infectious(person) && self.is_sane(other) {
251            return Some(other);
252        }
253        None
254    }
255
256    fn transmission(
257        &mut self,
258        now: Time,
259        person: PersonID,
260        other_occupants: Vec<(PersonID, Duration)>,
261        scheduler: &mut Scheduler,
262    ) {
263        // person has spent some duration in the same space as other people. Does transmission
264        // occur?
265        for (other, overlap) in other_occupants {
266            if let Some(pid) = self.infectious_contact(person, other) {
267                self.become_exposed(now, overlap, pid, scheduler);
268            }
269        }
270    }
271
272    // transition from a state to another without interaction with others
273    fn transition(&mut self, now: Time, person: PersonID, _scheduler: &mut Scheduler) {
274        let state = self.pop.remove(&person).unwrap();
275        let state = state.next(AnyTime::from(now), &mut self.rng).unwrap();
276        self.pop.insert(person, state);
277
278        // if self.rng.gen_bool(0.1) {
279        //     scheduler.push(
280        //         now + self.rand_duration(Duration::hours(1), Duration::hours(3)),
281        //         Command::Pandemic(Cmd::BecomeHospitalized(person)),
282        //     );
283        // }
284    }
285
286    fn become_exposed(
287        &mut self,
288        now: Time,
289        overlap: Duration,
290        person: PersonID,
291        _scheduler: &mut Scheduler,
292    ) {
293        #![allow(clippy::float_cmp)] // false positive
294                                     // When people become exposed
295        let state = self.pop.remove(&person).unwrap();
296        assert_eq!(
297            state.get_event_time().unwrap().inner_seconds(),
298            std::f64::INFINITY
299        );
300        let state = state
301            .start(AnyTime::from(now), overlap, &mut self.rng)
302            .unwrap();
303        self.pop.insert(person, state);
304
305        // if self.rng.gen_bool(0.1) {
306        //     scheduler.push(
307        //         now + self.rand_duration(Duration::hours(1), Duration::hours(3)),
308        //         Command::Pandemic(Cmd::BecomeHospitalized(person)),
309        //     );
310        // }
311    }
312}
313
314#[derive(Clone)]
315struct SharedSpace<T: Ord> {
316    // Since when has a person been in some shared space?
317    // TODO This is an awkward data structure; abstutil::MultiMap is also bad, because key removal
318    // would require knowing the time. Want something closer to
319    // https://guava.dev/releases/19.0/api/docs/com/google/common/collect/Table.html.
320    occupants: BTreeMap<T, Vec<(PersonID, Time)>>,
321}
322
323impl<T: Ord> SharedSpace<T> {
324    fn new() -> SharedSpace<T> {
325        SharedSpace {
326            occupants: BTreeMap::new(),
327        }
328    }
329
330    fn person_enters_space(&mut self, now: Time, person: PersonID, space: T) {
331        self.occupants
332            .entry(space)
333            .or_insert_with(Vec::new)
334            .push((person, now));
335    }
336
337    // Returns a list of all other people that the person was in the shared space with, and how
338    // long their time overlapped. If it returns None, then a bug must have occurred, because
339    // somebody has left a space they never entered.
340    fn person_leaves_space(
341        &mut self,
342        now: Time,
343        person: PersonID,
344        space: T,
345    ) -> Option<Vec<(PersonID, Duration)>> {
346        // TODO Messy to mutate state inside a retain closure
347        let mut inside_since: Option<Time> = None;
348        let occupants = self.occupants.entry(space).or_insert_with(Vec::new);
349        occupants.retain(|(p, t)| {
350            if *p == person {
351                inside_since = Some(*t);
352                false
353            } else {
354                true
355            }
356        });
357        // TODO Bug!
358        let inside_since = inside_since?;
359
360        Some(
361            occupants
362                .iter()
363                .map(|(p, t)| (*p, now - (*t).max(inside_since)))
364                .collect(),
365        )
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372
373    fn time(x: usize) -> Time {
374        Time::START_OF_DAY + Duration::hours(x)
375    }
376
377    #[test]
378    fn test_overlap() {
379        let mut space = SharedSpace::new();
380        let mut now = time(0);
381
382        let bldg1 = BuildingID(1);
383        let bldg2 = BuildingID(2);
384
385        let person1 = PersonID(1);
386        let person2 = PersonID(2);
387        let person3 = PersonID(3);
388
389        // Only one person
390        space.person_enters_space(now, person1, bldg1);
391        now = time(1);
392        assert_eq!(
393            space.person_leaves_space(now, person1, bldg1),
394            Some(Vec::new())
395        );
396
397        // Two people at the same time
398        now = time(2);
399        space.person_enters_space(now, person1, bldg2);
400        space.person_enters_space(now, person2, bldg2);
401        now = time(3);
402        assert_eq!(
403            space.person_leaves_space(now, person1, bldg2),
404            Some(vec![(person2, Duration::hours(1))])
405        );
406
407        // Bug
408        assert_eq!(space.person_leaves_space(now, person3, bldg2), None);
409
410        // Different times
411        now = time(5);
412        space.person_enters_space(now, person1, bldg1);
413        now = time(6);
414        space.person_enters_space(now, person2, bldg1);
415        now = time(7);
416        space.person_enters_space(now, person3, bldg1);
417        now = time(10);
418        assert_eq!(
419            space.person_leaves_space(now, person1, bldg1),
420            Some(vec![
421                (person2, Duration::hours(4)),
422                (person3, Duration::hours(3))
423            ])
424        );
425        now = time(12);
426        assert_eq!(
427            space.person_leaves_space(now, person2, bldg1),
428            Some(vec![(person3, Duration::hours(5))])
429        );
430    }
431}