99from datetime import date , datetime , timedelta
1010from logging import INFO , basicConfig , getLogger
1111from sys import stdout
12- from typing import Dict , Generator , Tuple , Optional
12+ from typing import Dict , Generator , Tuple , Sequence , Optional
1313
1414import numpy as np
1515import pandas as pd
@@ -40,6 +40,8 @@ def __init__(self, p: Parameters):
4040 for key , d in p .dispositions .items ()
4141 }
4242
43+ self .keys = ("susceptible" , "infected" , "recovered" )
44+
4345 # Note: this should not be an integer.
4446 # We're appoximating infected from what we do know.
4547 # TODO market_share > 0, hosp_rate > 0
@@ -153,7 +155,11 @@ def __init__(self, p: Parameters):
153155 self .beta_t * susceptible - gamma + 1 )
154156 self .doubling_time_t = doubling_time_t
155157
156- self .sim_sir_w_date_df = build_sim_sir_w_date_df (self .raw_df , p .current_date )
158+ self .sim_sir_w_date_df = build_sim_sir_w_date_df (self .raw_df , p .current_date , self .keys )
159+
160+ self .sim_sir_w_date_floor_df = build_floor_df (self .sim_sir_w_date_df , self .keys )
161+ self .admits_floor_df = build_floor_df (self .admits_df , p .dispositions .keys ())
162+ self .census_floor_df = build_floor_df (self .census_df , p .dispositions .keys ())
157163
158164 self .daily_growth_rate = get_growth_rate (p .doubling_time )
159165 self .daily_growth_rate_t = get_growth_rate (self .doubling_time_t )
@@ -252,7 +258,8 @@ def gen_sir(
252258
253259
254260def sim_sir_df (
255- s : float , i : float , r : float , gamma : float , i_day : int , * args
261+ s : float , i : float , r : float ,
262+ gamma : float , i_day : int , * args
256263) -> pd .DataFrame :
257264 """Simulate the SIR model forward in time."""
258265 return pd .DataFrame (
@@ -264,14 +271,28 @@ def sim_sir_df(
264271def build_sim_sir_w_date_df (
265272 raw_df : pd .DataFrame ,
266273 current_date : datetime ,
274+ keys : Sequence [str ],
267275) -> pd .DataFrame :
268276 day = raw_df .day
269277 return pd .DataFrame ({
270278 "day" : day ,
271279 "date" : day .astype ('timedelta64[D]' ) + np .datetime64 (current_date ),
272- "susceptible" : raw_df .susceptible ,
273- "infected" : raw_df .infected ,
274- "recovered" : raw_df .recovered ,
280+ ** {
281+ key : raw_df [key ]
282+ for key in keys
283+ }
284+ })
285+
286+
287+ def build_floor_df (df , keys ):
288+ """Build floor sim sir w date."""
289+ return pd .DataFrame ({
290+ "day" : df .day ,
291+ "date" : df .date ,
292+ ** {
293+ key : np .floor (df [key ])
294+ for key in keys
295+ }
275296 })
276297
277298
@@ -314,7 +335,7 @@ def build_census_df(
314335 key : (
315336 admits_df [key ].cumsum ()
316337 - admits_df [key ].cumsum ().shift (los ).fillna (0 )
317- ). apply ( np . ceil )
338+ )
318339 for key , los in lengths_of_stay .items ()
319340 }
320341 })
0 commit comments