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#[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#[derive(Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd, Ord, Clone, Debug)]
32pub enum Cmd {
33 BecomeHospitalized(PersonID),
34 BecomeQuarantined(PersonID),
35}
36
37impl 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 pub(crate) fn initialize(&mut self, population: &[Person], _scheduler: &mut Scheduler) {
63 assert!(!self.initialized);
64 self.initialized = true;
65
66 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 }
100
101 pub fn count_exposed(&self) -> usize {
102 self.pop
103 .iter()
104 .filter(|(_, state)| matches!(state, State::Exposed(_)))
105 .count()
106 }
108
109 pub fn count_infected(&self) -> usize {
110 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 }
124
125 pub fn count_dead(&self) -> usize {
126 self.pop
127 .iter()
128 .filter(|(_, state)| matches!(state, State::Dead(_)))
129 .count()
130 }
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 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 match cmd {
196 Cmd::BecomeHospitalized(_person) => {
197 }
199 Cmd::BecomeQuarantined(_person) => {
200 }
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 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 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 }
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)] 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 }
312}
313
314#[derive(Clone)]
315struct SharedSpace<T: Ord> {
316 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 fn person_leaves_space(
341 &mut self,
342 now: Time,
343 person: PersonID,
344 space: T,
345 ) -> Option<Vec<(PersonID, Duration)>> {
346 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 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 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 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 assert_eq!(space.person_leaves_space(now, person3, bldg2), None);
409
410 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}