sim/pandemic/
mod.rs

1//! An experimental SEIR model by https://github.com/omalaspinas/ glued to the traffic simulation.
2//! Transmission may occur when people spend time in shared spaces like buildings, bus stops, and
3//! buses.
4
5use std::ops;
6
7use anyhow::Result;
8
9pub use model::{Cmd, PandemicModel};
10use rand::Rng;
11use rand_distr::{Distribution, Exp, Normal};
12use rand_xorshift::XorShiftRng;
13
14use geom::{Duration, Time};
15
16mod model;
17
18#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
19pub struct AnyTime(f64);
20
21impl AnyTime {
22    fn inner_seconds(&self) -> f64 {
23        self.0
24    }
25
26    pub fn is_finite(&self) -> bool {
27        self.0.is_finite()
28    }
29}
30
31impl ops::Add<Duration> for AnyTime {
32    type Output = AnyTime;
33
34    fn add(self, other: Duration) -> AnyTime {
35        AnyTime(self.0 + other.inner_seconds())
36    }
37}
38
39impl ops::AddAssign<Duration> for AnyTime {
40    fn add_assign(&mut self, other: Duration) {
41        *self = *self + other;
42    }
43}
44
45impl ops::Sub<Duration> for AnyTime {
46    type Output = AnyTime;
47
48    fn sub(self, other: Duration) -> AnyTime {
49        AnyTime(self.0 - other.inner_seconds())
50    }
51}
52
53impl ops::Sub for AnyTime {
54    type Output = Duration;
55
56    fn sub(self, other: AnyTime) -> Duration {
57        Duration::seconds(self.0 - other.0)
58    }
59}
60
61impl From<Time> for AnyTime {
62    fn from(t: Time) -> AnyTime {
63        AnyTime(t.inner_seconds())
64    }
65}
66
67impl From<AnyTime> for Time {
68    fn from(time: AnyTime) -> Self {
69        Time::START_OF_DAY + Duration::seconds(time.inner_seconds())
70    }
71}
72
73impl From<f64> for AnyTime {
74    fn from(t: f64) -> AnyTime {
75        AnyTime(t)
76    }
77}
78
79#[derive(Debug, Clone)]
80pub enum StateEvent {
81    Exposition,
82    Incubation,
83    Hospitalization,
84    Recovery,
85    Death,
86}
87
88#[derive(Debug, Clone)]
89pub struct Event {
90    s: StateEvent,
91    p_hosp: f64,  // probability of people being hospitalized after infection
92    p_death: f64, // probability of dying after hospitalizaion
93    t: AnyTime,
94}
95
96impl Event {
97    fn next(&self, now: AnyTime, rng: &mut XorShiftRng) -> State {
98        match self.s {
99            StateEvent::Exposition => State::Exposed((
100                Event {
101                    s: StateEvent::Incubation,
102                    p_hosp: self.p_hosp,
103                    p_death: self.p_death,
104                    t: now + State::get_time_normal(State::T_INC, State::T_INC / 2.0, rng),
105                },
106                now.into(),
107            )),
108            StateEvent::Incubation => {
109                if rng.gen_bool(self.p_death) {
110                    State::Infectious((
111                        Event {
112                            s: StateEvent::Recovery,
113                            p_hosp: self.p_hosp,
114                            p_death: self.p_death,
115                            t: now + State::get_time_normal(State::T_INF, State::T_INF / 2.0, rng),
116                        },
117                        now.into(),
118                    ))
119                } else {
120                    State::Infectious((
121                        Event {
122                            s: StateEvent::Hospitalization,
123                            p_hosp: self.p_hosp,
124                            p_death: self.p_death,
125                            t: now + State::get_time_normal(State::T_INF, State::T_INF / 2.0, rng),
126                        },
127                        now.into(),
128                    ))
129                }
130            }
131            StateEvent::Hospitalization => {
132                if rng.gen_bool(self.p_hosp) {
133                    State::Hospitalized((
134                        Event {
135                            s: StateEvent::Recovery,
136                            p_hosp: self.p_hosp,
137                            p_death: self.p_death,
138                            t: now + State::get_time_normal(State::T_INF, State::T_INF / 2.0, rng),
139                        },
140                        now.into(),
141                    ))
142                } else {
143                    State::Hospitalized((
144                        Event {
145                            s: StateEvent::Death,
146                            p_hosp: self.p_hosp,
147                            p_death: self.p_death,
148                            t: now + State::get_time_normal(State::T_INF, State::T_INF / 2.0, rng),
149                        },
150                        now.into(),
151                    ))
152                }
153            }
154            StateEvent::Death => State::Dead(now.into()),
155            StateEvent::Recovery => State::Recovered(now.into()),
156        }
157    }
158}
159
160#[derive(Debug, Clone)]
161pub enum State {
162    Sane((Event, Time)),
163    Exposed((Event, Time)),
164    Infectious((Event, Time)),
165    Hospitalized((Event, Time)),
166    Recovered(Time),
167    Dead(Time),
168}
169
170impl State {
171    const T_INF: f64 = 360.0 * 10.0; // TODO dummy values
172    const T_INC: f64 = 3600.0; // TODO dummy values
173    const R_0: f64 = 2.5;
174    // const S_RATIO: f64 = 0.985;
175    const E_RATIO: f64 = 0.01;
176    const I_RATIO: f64 = 0.05;
177    // const R_RATIO: f64 = 0.0;
178
179    pub fn ini_infectious_ratio() -> f64 {
180        Self::I_RATIO
181    }
182
183    pub fn ini_exposed_ratio() -> f64 {
184        Self::E_RATIO
185    }
186
187    fn new(p_hosp: f64, p_death: f64) -> Self {
188        Self::Sane((
189            Event {
190                s: StateEvent::Exposition,
191                p_hosp,
192                p_death,
193                t: AnyTime::from(std::f64::INFINITY),
194            },
195            Time::START_OF_DAY,
196        ))
197    }
198
199    fn get_time_exp(lambda: f64, rng: &mut XorShiftRng) -> geom::Duration {
200        let normal = Exp::new(lambda).unwrap();
201        Duration::seconds(normal.sample(rng))
202    }
203
204    fn get_time_normal(mu: f64, sigma: f64, rng: &mut XorShiftRng) -> geom::Duration {
205        let normal = Normal::new(mu, sigma).unwrap();
206        Duration::seconds(normal.sample(rng))
207    }
208
209    fn is_sane(&self) -> bool {
210        match self {
211            State::Sane((ev, _)) => !ev.t.is_finite(),
212            _ => false,
213        }
214    }
215
216    fn is_exposed(&self) -> bool {
217        matches!(self, State::Exposed(_))
218    }
219
220    fn is_infectious(&self) -> bool {
221        matches!(self, State::Infectious(_) | State::Hospitalized(_))
222    }
223
224    fn is_recovered(&self) -> bool {
225        matches!(self, State::Recovered(_))
226    }
227
228    fn is_dead(&self) -> bool {
229        matches!(self, State::Dead(_))
230    }
231
232    pub fn get_time(&self) -> Option<Time> {
233        match self {
234            Self::Sane(_) => None,
235            Self::Recovered(t)
236            | Self::Dead(t)
237            | Self::Exposed((_, t))
238            | Self::Infectious((_, t))
239            | Self::Hospitalized((_, t)) => Some(*t),
240        }
241    }
242
243    pub fn get_event_time(&self) -> Option<AnyTime> {
244        match self {
245            Self::Sane((ev, _))
246            | Self::Exposed((ev, _))
247            | Self::Infectious((ev, _))
248            | Self::Hospitalized((ev, _)) => Some(ev.t),
249            Self::Recovered(_) | Self::Dead(_) => None,
250        }
251    }
252
253    // pub fn set_time(self, new_time: AnyTime) -> Self {
254    //     match self {
255    //         Self::Sane(Event {
256    //             s,
257    //             p_hosp,
258    //             p_death,
259    //             t: _,
260    //         }) => Self::Sane(Event {
261    //             s,
262    //             p_hosp,
263    //             p_death,
264    //             t: new_time,
265    //         }),
266    //         _ => unreachable!(),
267    //     }
268    // }
269
270    // TODO: not sure if we want an option here...
271    pub fn next_default(self, default: AnyTime, rng: &mut XorShiftRng) -> Option<Self> {
272        // TODO: when #![feature(bindings_after_at)] reaches stable
273        // rewrite this part with it
274        match self {
275            Self::Sane((ev, _)) => Some(Self::Sane((ev, default.into()))),
276            Self::Exposed((ev, _)) => Some(ev.next(default, rng)),
277            Self::Infectious((ev, _)) => Some(ev.next(default, rng)),
278            Self::Hospitalized((ev, _)) => Some(ev.next(default, rng)),
279            Self::Recovered(_) => Some(Self::Recovered(default.into())),
280            Self::Dead(_) => Some(Self::Dead(default.into())),
281        }
282    }
283
284    // TODO: not sure if we want an option here...
285    pub fn next(self, now: AnyTime, rng: &mut XorShiftRng) -> Option<Self> {
286        // TODO: when #![feature(bindings_after_at)] reaches stable
287        // rewrite this part with it
288        match self {
289            Self::Sane((ev, t)) => Some(Self::Sane((ev, t))),
290            Self::Exposed((ev, t)) => {
291                if ev.t <= now {
292                    Some(ev.next(now, rng))
293                } else {
294                    Some(Self::Exposed((ev, t)))
295                }
296            }
297            Self::Infectious((ev, t)) => {
298                if ev.t <= now {
299                    Some(ev.next(now, rng))
300                } else {
301                    Some(Self::Infectious((ev, t)))
302                }
303            }
304            Self::Hospitalized((ev, t)) => {
305                if ev.t <= now {
306                    Some(ev.next(now, rng))
307                } else {
308                    Some(Self::Hospitalized((ev, t)))
309                }
310            }
311            Self::Recovered(t) => Some(Self::Recovered(t)),
312            Self::Dead(t) => Some(Self::Dead(t)),
313        }
314    }
315
316    // TODO: not sure if we want an option here... I guess here we want because we could have
317    pub fn start(self, now: AnyTime, overlap: Duration, rng: &mut XorShiftRng) -> Result<Self> {
318        // rewrite this part with it
319        match self {
320            Self::Sane((ev, t)) => {
321                if overlap >= Self::get_time_exp(State::R_0 / State::T_INF, rng) {
322                    Ok(ev.next(now, rng))
323                } else {
324                    Ok(Self::Sane((ev, t)))
325                }
326            }
327            _ => bail!("impossible to start from a non-sane situation.",),
328        }
329    }
330}