Extend prefit days for mitigation_date

Author: jlubken

setup.py

@@ -1,6 +1,6 @@
 """Setup file for chime
 """
-__version__ = "1.1.3"  # update VERSION in constants.py
+__version__ = "1.1.4"  # update VERSION in constants.py
 __author__ = "Predictive Healthcare @ Penn Medicine"
 
 from setuptools import setup, find_namespace_packages

src/penn_chime/constants.py

@@ -8,7 +8,7 @@
 re-run their reports
 """
 CHANGE_DATE = date(year=2020, month=4, day=8)
-VERSION = 'v1.1.3'
+VERSION = 'v1.1.4'
 
 DATE_FORMAT = "%b, %d"  # see https://strftime.org
 DOCS_URL = "https://code-for-philly.gitbook.io/chime"

src/penn_chime/model/sir.py

@@ -73,28 +73,38 @@ def __init__(self, p: Parameters):
                     (self.beta, p.n_days + PREFIT_ADDITIONAL_DAYS)])
                 self.i_day = i_day = int(get_argmin_ds(raw["census_hospitalized"], p.current_hospitalized))
 
-                self.raw = self.run_projection(p, self.gen_policy(p))
+                self.raw = self.run_projection(p, self.get_policies(p))
 
                 logger.info('Set i_day = %s', i_day)
             else:
-                projections = {}
                 best_i_day = -1
                 best_i_day_loss = float('inf')
-                for i_day in range(p.n_days):
-                    self.i_day = i_day
-                    raw = self.run_projection(p, self.gen_policy(p))
+                for self.i_day in range(p.n_days + PREFIT_ADDITIONAL_DAYS):
+                    mitigation_day = -(p.current_date - p.mitigation_date).days
+                    if mitigation_day < -self.i_day:
+                        mitigation_day = -self.i_day
+
+                    total_days = self.i_day + p.n_days + PREFIT_ADDITIONAL_DAYS
+                    pre_mitigation_days = self.i_day + mitigation_day
+                    post_mitigation_days = total_days - pre_mitigation_days
+
+                    raw = self.run_projection(p, [
+                            (self.beta,   pre_mitigation_days),
+                            (self.beta_t, post_mitigation_days),
+                        ]
+                    )
 
                     # Don't fit against results that put the peak before the present day
-                    if raw["census_hospitalized"].argmax() < i_day:
+                    if raw["census_hospitalized"].argmax() < self.i_day:
                         continue
 
-                    loss = get_loss(raw["census_hospitalized"][i_day], p.current_hospitalized)
+                    loss = get_loss(raw["census_hospitalized"][self.i_day], p.current_hospitalized)
                     if loss < best_i_day_loss:
                         best_i_day_loss = loss
-                        best_i_day = i_day
-                        self.raw = raw
+                        best_i_day = self.i_day
 
                 self.i_day = best_i_day
+                self.raw = self.run_projection(p, self.get_policies(p))
 
             logger.info(
                 'Estimated date_first_hospitalized: %s; current_date: %s; i_day: %s',
@@ -129,7 +139,7 @@ def __init__(self, p: Parameters):
             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.raw = self.run_projection(p, self.gen_policy(p))
+            self.raw = self.run_projection(p, self.get_policies(p))
 
             self.population = p.population
         else:
@@ -198,7 +208,7 @@ def get_argmin_doubling_time(self, p: Parameters, dts):
             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)
 
-            raw = self.run_projection(p, self.gen_policy(p))
+            raw = self.run_projection(p, self.get_policies(p))
 
             # Skip values the would put the fit past peak
             peak_admits_day = raw["admits_hospitalized"].argmax()
@@ -212,7 +222,7 @@ def get_argmin_doubling_time(self, p: Parameters, dts):
         min_loss = pd.Series(losses).argmin()
         return min_loss
 
-    def gen_policy(self, p: Parameters) -> Sequence[Tuple[float, int]]:
+    def get_policies(self, p: Parameters) -> Sequence[Tuple[float, int]]:
         if p.mitigation_date is not None:
             mitigation_day = -(p.current_date - p.mitigation_date).days
         else: