Put social distancing effect back in. Starts at current day.

Author: Corey Chivers

src/penn_chime/models.py

@@ -80,21 +80,19 @@ def __init__(self, p: Parameters):
         # Simplify equation to avoid division by zero:
         # self.r_naught = r_t / (1.0 - relative_contact_rate)
         r_naught = (intrinsic_growth_rate + gamma) / gamma
-        doubling_time_t = 1.0 / np.log2(
-            beta * susceptible - gamma + 1)
+
 
         self.susceptible = susceptible
         self.infected = infected
         self.recovered = p.recovered
 
         self.detection_probability = detection_probability
-        self.doubling_time_t = doubling_time_t
+
         self.beta = beta
         self.gamma = gamma
+        self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
         self.intrinsic_growth_rate = intrinsic_growth_rate
-        self.r_t = r_t
-        self.r_naught = r_naught
-        self.infected = 1.0 / p.hospitalized.rate / p.market_share
+        self.infected = infected
 
         if p.date_first_hospitalized is None and p.doubling_time is not None:
             logger.info('Using doubling_time.')
@@ -104,13 +102,19 @@ def __init__(self, p: Parameters):
                 / susceptible
             )
 
+            self.i_day = 0 # seed to the full length
+            self.beta_t = self.beta
             self.run_projection(p)
             self.i_day = i_day = int(get_argmin_ds(self.census_df, p.current_hospitalized))
+
+            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
             self.run_projection(p)
             self.infected = self.raw_df['infected'].values[i_day]
             self.susceptible = self.raw_df['susceptible'].values[i_day]
             self.recovered = self.raw_df['recovered'].values[i_day]
             self.n_days_since = i_day
+            self.r_t = self.beta_t / gamma * susceptible
+            self.r_naught = self.beta / gamma * susceptible
             logger.info('Set i_day = %s', i_day)
 
         elif p.date_first_hospitalized is not None and p.doubling_time is None:
@@ -123,6 +127,7 @@ def __init__(self, p: Parameters):
             for i, i_dt in enumerate(dts):
                 intrinsic_growth_rate = get_growth_rate(i_dt)
                 self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
+                self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
 
                 self.run_projection(p)
                 loss = self.get_loss()
@@ -132,17 +137,20 @@ def __init__(self, p: Parameters):
             logger.info('Set doubling_time = %s', doubling_time)
             intrinsic_growth_rate = get_growth_rate(p.doubling_time)
             self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
+            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
             self.run_projection(p)
 
             self.intrinsic_growth_rate = intrinsic_growth_rate
-            self.beta = beta
-            self.doubling_time_t = doubling_time_t
             self.population = p.population
-            self.r_t = r_t
-            self.r_naught = r_naught
         else:
             raise AssertionError('doubling_time or date_first_hospitalized must be provided.')
 
+        self.r_t = self.beta_t / gamma * susceptible
+        self.r_naught = self.beta / gamma * susceptible
+
+        doubling_time_t = 1.0 / np.log2(
+            self.beta_t * susceptible - gamma + 1)
+        self.doubling_time_t = doubling_time_t
 
         self.sim_sir_w_date_df = build_sim_sir_w_date_df(self.raw_df, p.current_date)
 
@@ -154,10 +162,12 @@ def run_projection(self, p):
             self.susceptible,
             self.infected,
             p.recovered,
-            self.beta,
             self.gamma,
-            p.n_days + self.i_day,
-            -self.i_day
+            -self.i_day,
+            self.beta,
+            self.i_day,
+            self.beta_t,
+            p.n_days
         )
         self.dispositions_df = build_dispositions_df(self.raw_df, self.rates, p.market_share, p.current_date)
         self.admits_df = build_admits_df(self.dispositions_df)
@@ -225,28 +235,30 @@ def sir(
 
 
 def gen_sir(
-    s: float, i: float, r: float,
-    beta: float, gamma: float, n_days: int, i_day: int = 0
+    s: float, i: float, r: float, gamma: float, i_day: int = 0, *args
 ) -> Generator[Tuple[int, float, float, float], None, None]:
-    """Simulate SIR model forward in time yielding tuples."""
+    """Simulate SIR model forward in time yielding tuples.
+    Parameter order has changed to allow multiple (beta, n_days)
+    to reflect multiple changing social distancing policies.
+    """
     s, i, r = (float(v) for v in (s, i, r))
     n = s + i + r
     d = i_day
-    # TODO: Ask corey if n_days is really required for the sim
-    # while i >= 0.5:
-    for _ in range(n_days):
-        yield d, s, i, r
-        s, i, r = sir(s, i, r, beta, gamma, n)
-        d += 1
+    while args:
+        beta, n_days, *args = args
+        for _ in range(n_days):
+            yield d, s, i, r
+            s, i, r = sir(s, i, r, beta, gamma, n)
+            d += 1
     yield d, s, i, r
 
 
 def sim_sir_df(
-    s: float, i: float, r: float, beta: float, gamma: float, n_days: int, i_day: int = 0
+    s: float, i: float, r: float, gamma: float, i_day: int = 0, *args
 ) -> pd.DataFrame:
     """Simulate the SIR model forward in time."""
     return pd.DataFrame(
-        data=gen_sir(s, i, r, beta, gamma, n_days, i_day),
+        data=gen_sir(s, i, r, gamma, i_day, *args),
         columns=("day", "susceptible", "infected", "recovered"),
     )