9
9
from datetime import date , datetime , timedelta
10
10
from logging import INFO , basicConfig , getLogger
11
11
from sys import stdout
12
- from typing import Dict , Generator , Tuple , Sequence ,Optional
12
+ from typing import Dict , Generator , Tuple , Sequence , Optional
13
13
14
14
import numpy as np
15
15
import pandas as pd
@@ -145,7 +145,7 @@ def __init__(self, p: Parameters):
145
145
self .daily_growth_rate = get_growth_rate (p .doubling_time )
146
146
self .daily_growth_rate_t = get_growth_rate (self .doubling_time_t )
147
147
148
- def gen_policy (self , p : Parameters ) -> List [Tuple [float , int ]]:
148
+ def gen_policy (self , p : Parameters ) -> Sequence [Tuple [float , int ]]:
149
149
if p .mitigation_date is not None :
150
150
mitigation_day = - (p .current_date - p .mitigation_date ).days
151
151
else :
@@ -164,7 +164,7 @@ def gen_policy(self, p: Parameters) -> List[Tuple[float, int]]:
164
164
(self .beta_t , post_mitigation_days ),
165
165
]
166
166
167
- def run_projection (self , p : Parameters , policy : List [Tuple [float , int ]]):
167
+ def run_projection (self , p : Parameters , policy : Sequence [Tuple [float , int ]]):
168
168
self .raw_df = sim_sir_df (
169
169
self .susceptible ,
170
170
self .infected ,
@@ -237,7 +237,7 @@ def sir(
237
237
238
238
239
239
def gen_sir (
240
- s : float , i : float , r : float , gamma : float , i_day : int , policies : List [Tuple [float , int ]]
240
+ s : float , i : float , r : float , gamma : float , i_day : int , policies : Sequence [Tuple [float , int ]]
241
241
) -> Generator [Tuple [int , float , float , float ], None , None ]:
242
242
"""Simulate SIR model forward in time yielding tuples.
243
243
Parameter order has changed to allow multiple (beta, n_days)
@@ -256,7 +256,7 @@ def gen_sir(
256
256
257
257
def sim_sir_df (
258
258
s : float , i : float , r : float ,
259
- gamma : float , i_day : int , policies : List [Tuple [float , int ]]
259
+ gamma : float , i_day : int , policies : Sequence [Tuple [float , int ]]
260
260
) -> pd .DataFrame :
261
261
"""Simulate the SIR model forward in time."""
262
262
return pd .DataFrame (
0 commit comments