1use 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, p_death: f64, 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; const T_INC: f64 = 3600.0; const R_0: f64 = 2.5;
174 const E_RATIO: f64 = 0.01;
176 const I_RATIO: f64 = 0.05;
177 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 next_default(self, default: AnyTime, rng: &mut XorShiftRng) -> Option<Self> {
272 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 pub fn next(self, now: AnyTime, rng: &mut XorShiftRng) -> Option<Self> {
286 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 pub fn start(self, now: AnyTime, overlap: Duration, rng: &mut XorShiftRng) -> Result<Self> {
318 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}