11"""Models."""
22
3- from typing import Generator , Tuple
3+ from __future__ import annotations
4+
5+ from typing import Dict , Generator , Tuple
46
57import numpy as np # type: ignore
68import pandas as pd # type: ignore
79
10+ from .parameters import Parameters
11+
12+
13+ class SimSirModel :
14+
15+ def __init__ (self , p : Parameters ) -> SimSirModel :
16+
17+ # Note: this should not be an integer.
18+ # We're appoximating infected from what we do know.
19+ # TODO market_share > 0, hosp_rate > 0
20+ self .infected = infected = (
21+ p .current_hospitalized / p .market_share / p .hospitalized .rate
22+ )
23+
24+ self .detection_probability = (
25+ p .known_infected / infected if infected > 1.0e-7 else None
26+ )
27+
28+ # TODO missing initial recovered value
29+ self .recovered = recovered = 0.0
30+
31+ self .intrinsic_growth_rate = intrinsic_growth_rate = \
32+ (2.0 ** (1.0 / p .doubling_time ) - 1.0 ) if p .doubling_time > 0.0 else 0.0
33+
34+ # TODO make this configurable, or more nuanced
35+ self .recovery_days = recovery_days = 14.0
36+
37+ self .gamma = gamma = 1.0 / recovery_days
38+
39+ # Contact rate, beta
40+ self .beta = beta = (
41+ (intrinsic_growth_rate + gamma )
42+ / p .susceptible
43+ * (1.0 - p .relative_contact_rate )
44+ ) # {rate based on doubling time} / {initial susceptible}
45+
46+ # r_t is r_0 after distancing
47+ self .r_t = beta / gamma * p .susceptible
48+
49+ # Simplify equation to avoid division by zero:
50+ # self.r_naught = r_t / (1.0 - relative_contact_rate)
51+ self .r_naught = (intrinsic_growth_rate + gamma ) / gamma
52+
53+ # doubling time after distancing
54+ # TODO constrain values np.log2(...) > 0.0
55+ self .doubling_time_t = 1.0 / np .log2 (
56+ beta * p .susceptible - gamma + 1 )
57+
58+ self .raw_df = raw_df = sim_sir_df (
59+ p .susceptible ,
60+ infected ,
61+ recovered ,
62+ beta ,
63+ gamma ,
64+ p .n_days ,
65+ )
66+
67+ rates = {
68+ key : d .rate
69+ for key , d in p .dispositions .items ()
70+ }
71+
72+ lengths_of_stay = {
73+ key : d .length_of_stay
74+ for key , d in p .dispositions .items ()
75+ }
76+
77+ i_dict_v = get_dispositions (raw_df .infected , rates , p .market_share )
78+ r_dict_v = get_dispositions (raw_df .recovered , rates , p .market_share )
79+
80+ self .dispositions = {
81+ key : value + r_dict_v [key ]
82+ for key , value in i_dict_v .items ()
83+ }
84+
85+ self .dispositions_df = pd .DataFrame (self .dispositions )
86+ self .admits_df = admits_df = build_admits_df (p .n_days , self .dispositions )
87+ self .census_df = build_census_df (admits_df , lengths_of_stay )
88+
889
990def sir (
1091 s : float , i : float , r : float , beta : float , gamma : float , n : float
@@ -30,91 +111,61 @@ def gen_sir(
30111 """Simulate SIR model forward in time yielding tuples."""
31112 s , i , r = (float (v ) for v in (s , i , r ))
32113 n = s + i + r
33- for _ in range (n_days + 1 ):
34- yield s , i , r
114+ for d in range (n_days + 1 ):
115+ yield d , s , i , r
35116 s , i , r = sir (s , i , r , beta , gamma , n )
36117
37118
38- def sim_sir (
39- s : float , i : float , r : float , beta : float , gamma : float , n_days : int
40- ) -> Tuple [ np . ndarray , np . ndarray , np . ndarray ] :
119+ def sim_sir_df (
120+ s : float , i : float , r : float , beta : float , gamma : float , n_days
121+ ) -> pd . DataFrame :
41122 """Simulate the SIR model forward in time."""
42- s , i , r = (float (v ) for v in (s , i , r ))
43- n = s + i + r
44- s_v , i_v , r_v = [s ], [i ], [r ]
45- for day in range (n_days ):
46- s , i , r = sir (s , i , r , beta , gamma , n )
47- s_v .append (s )
48- i_v .append (i )
49- r_v .append (r )
50-
51- return (
52- np .array (s_v ),
53- np .array (i_v ),
54- np .array (r_v ),
55- )
56-
57-
58- def sim_sir_df (p ) -> pd .DataFrame :
59- """Simulate the SIR model forward in time.
60-
61- p is a Parameters instance. for circuluar dependency reasons i can't annotate it.
62- """
63123 return pd .DataFrame (
64- data = gen_sir (p . susceptible , p . infected , p . recovered , p . beta , p . gamma , p . n_days ),
65- columns = ("Susceptible " , "Infected " , "Recovered " ),
124+ data = gen_sir (s , i , r , beta , gamma , n_days ),
125+ columns = ("day " , "susceptible " , "infected" , "recovered "
66126 )
67127
68128
69129def get_dispositions (
70- patient_state : np .ndarray , rates : Tuple [float , ...], market_share : float = 1.0
71- ) -> Tuple [np .ndarray , ...]:
72- """Get dispositions of infected adjusted by rate and market_share."""
73- return (* (patient_state * rate * market_share for rate in rates ),)
74-
75-
76- def build_admissions_df (p ) -> pd .DataFrame :
77- """Build admissions dataframe from Parameters."""
78- days = np .array (range (0 , p .n_days + 1 ))
79- data_dict = dict (
80- zip (
81- ["day" , "Hospitalized" , "ICU" , "Ventilated" ],
82- [days ] + [disposition for disposition in p .dispositions ],
83- )
84- )
85- projection = pd .DataFrame .from_dict (data_dict )
130+ patients : np .ndarray ,
131+ rates : Dict [str , float ],
132+ market_share : float ,
133+ ) -> Dict [str , np .ndarray ]:
134+ """Get dispositions of patients adjusted by rate and market_share."""
135+ return {
136+ key : patients * rate * market_share
137+ for key , rate in rates .items ()
138+ }
139+
140+
141+ def build_admits_df (n_days , dispositions ) -> pd .DataFrame :
142+ """Build admits dataframe from Parameters and Model."""
143+ days = np .arange (0 , n_days + 1 )
144+ projection = pd .DataFrame ({
145+ "day" : days ,
146+ ** dispositions ,
147+ })
86148 # New cases
87- projection_admits = projection .iloc [:- 1 , :] - projection .shift (1 )
88- projection_admits ["day" ] = range (projection_admits .shape [0 ])
89- return projection_admits
149+ admits_df = projection .iloc [:- 1 , :] - projection .shift (1 )
150+ admits_df ["day" ] = range (admits_df .shape [0 ])
151+ return admits_df
90152
91153
92- def build_census_df (projection_admits : pd .DataFrame , parameters ) -> pd .DataFrame :
154+ def build_census_df (
155+ admits_df : pd .DataFrame , lengths_of_stay
156+ ) -> pd .DataFrame :
93157 """ALOS for each category of COVID-19 case (total guesses)"""
94- n_days = np .shape (projection_admits )[0 ]
95- hosp_los , icu_los , vent_los = parameters .lengths_of_stay
96- los_dict = {
97- "Hospitalized" : hosp_los ,
98- "ICU" : icu_los ,
99- "Ventilated" : vent_los ,
100- }
101-
102- census_dict = dict ()
103- for k , los in los_dict .items ():
158+ n_days = np .shape (admits_df )[0 ]
159+ census_dict = {}
160+ for key , los in lengths_of_stay .items ():
104161 census = (
105- projection_admits .cumsum ().iloc [:- los , :]
106- - projection_admits .cumsum ().shift (los ).fillna (0 )
162+ admits_df .cumsum ().iloc [:- los , :]
163+ - admits_df .cumsum ().shift (los ).fillna (0 )
107164 ).apply (np .ceil )
108- census_dict [k ] = census [k ]
165+ census_dict [key ] = census [key ]
109166
110167 census_df = pd .DataFrame (census_dict )
111168 census_df ["day" ] = census_df .index
112- census_df = census_df [["day" , "Hospitalized" , "ICU" , "Ventilated" ]]
169+ census_df = census_df [["day" , * lengths_of_stay . keys () ]]
113170 census_df = census_df .head (n_days )
114- census_df = census_df .rename (
115- columns = {
116- disposition : f"{ disposition }
117- for disposition in ("Hospitalized" , "ICU" , "Ventilated" )
118- }
119- )
120171 return census_df
0 commit comments