Merge pull request #256 from CodeForPhilly/fix_dfs

Fix dispositions and census construction

Author: quinn-dougherty

src/penn_chime/models.py

@@ -18,75 +18,81 @@
 class SimSirModel:
 
     def __init__(self, p: Parameters) -> SimSirModel:
+        # TODO missing initial recovered value
+        susceptible = p.susceptible
+        recovered = 0.0
+        recovery_days = p.recovery_days
+
+        rates = {
+            key: d.rate
+            for key, d in p.dispositions.items()
+        }
+
+        lengths_of_stay = {
+            key: d.length_of_stay
+            for key, d in p.dispositions.items()
+        }
 
         # Note: this should not be an integer.
         # We're appoximating infected from what we do know.
         # TODO market_share > 0, hosp_rate > 0
-        self.infected = infected = (
+        infected = (
             p.current_hospitalized / p.market_share / p.hospitalized.rate
         )
 
-        self.detection_probability = (
+        detection_probability = (
             p.known_infected / infected if infected > 1.0e-7 else None
         )
 
-        # TODO missing initial recovered value
-        self.recovered = recovered = 0.0
-
-        self.intrinsic_growth_rate = intrinsic_growth_rate = \
+        intrinsic_growth_rate = \
             (2.0 ** (1.0 / p.doubling_time) - 1.0) if p.doubling_time > 0.0 else 0.0
 
-        self.gamma = gamma = 1.0 / p.recovery_days
+        gamma = 1.0 / recovery_days
 
         # Contact rate, beta
-        self.beta = beta = (
+        beta = (
             (intrinsic_growth_rate + gamma)
-            / p.susceptible
+            / susceptible
             * (1.0 - p.relative_contact_rate)
         )  # {rate based on doubling time} / {initial susceptible}
 
         # r_t is r_0 after distancing
-        self.r_t = beta / gamma * p.susceptible
+        r_t = beta / gamma * susceptible
 
         # Simplify equation to avoid division by zero:
         # self.r_naught = r_t / (1.0 - relative_contact_rate)
-        self.r_naught = (intrinsic_growth_rate + gamma) / gamma
-
-        # doubling time after distancing
-        # TODO constrain values np.log2(...) > 0.0
-        self.doubling_time_t = 1.0 / np.log2(
+        r_naught = (intrinsic_growth_rate + gamma) / gamma
+        doubling_time_t = 1.0 / np.log2(
             beta * p.susceptible - gamma + 1)
 
-        self.raw_df = raw_df = sim_sir_df(
+        raw_df = sim_sir_df(
             p.susceptible,
             infected,
             recovered,
             beta,
             gamma,
             p.n_days,
         )
-
-        rates = {
-            key: d.rate
-            for key, d in p.dispositions.items()
-        }
-
-        lengths_of_stay = {
-            key: d.length_of_stay
-            for key, d in p.dispositions.items()
-        }
-
-        i_dict_v = get_dispositions(raw_df.infected, rates, p.market_share)
-        r_dict_v = get_dispositions(raw_df.recovered, rates, p.market_share)
-
-        self.dispositions = {
-            key: value + r_dict_v[key]
-            for key, value in i_dict_v.items()
-        }
-
-        self.dispositions_df = pd.DataFrame(self.dispositions)
-        self.admits_df = admits_df = build_admits_df(p.n_days, self.dispositions)
-        self.census_df = build_census_df(admits_df, lengths_of_stay)
+        dispositions_df = build_dispositions_df(raw_df, rates, p.market_share)
+        admits_df = build_admits_df(dispositions_df)
+        census_df = build_census_df(admits_df, lengths_of_stay)
+
+        self.susceptible = susceptible
+        self.infected = infected
+        self.recovered = recovered
+
+        self.detection_probability = detection_probability
+        self.recovered = recovered
+        self.intrinsic_growth_rate = intrinsic_growth_rate
+        self.gamma = gamma
+        self.beta = beta
+        self.r_t = r_t
+        self.r_naught = r_naught
+        self.doubling_time_t = doubling_time_t
+        self.raw_df = raw_df
+        self.dispositions_df = dispositions_df
+        self.admits_df = admits_df
+        self.census_df = census_df
 
 
 def sir(
@@ -119,55 +125,49 @@ def gen_sir(
 
 
 def sim_sir_df(
-    s: float, i: float, r: float, beta: float, gamma: float, n_days
+    s: float, i: float, r: float, beta: float, gamma: float, n_days: int
 ) -> pd.DataFrame:
     """Simulate the SIR model forward in time."""
     return pd.DataFrame(
         data=gen_sir(s, i, r, beta, gamma, n_days),
         columns=("day", "susceptible", "infected", "recovered"),
     )
 
-
-def get_dispositions(
-    patients: np.ndarray,
+def build_dispositions_df(
+    sim_sir_df: pd.DataFrame,
     rates: Dict[str, float],
     market_share: float,
-) -> Dict[str, np.ndarray]:
+) -> pd.DataFrame:
     """Get dispositions of patients adjusted by rate and market_share."""
-    return {
-        key: patients * rate * market_share
-        for key, rate in rates.items()
-    }
-
-
-def build_admits_df(n_days, dispositions) -> pd.DataFrame:
-    """Build admits dataframe from Parameters and Model."""
-    days = np.arange(0, n_days + 1)
-    projection = pd.DataFrame({
-        "day": days,
-        **dispositions,
+    patients = sim_sir_df.infected + sim_sir_df.recovered
+    return pd.DataFrame({
+        "day": sim_sir_df.day,
+        **{
+            key: patients * rate * market_share
+            for key, rate in rates.items()
+        }
     })
-    # New cases
-    admits_df = projection.iloc[:-1, :] - projection.shift(1)
-    admits_df["day"] = range(admits_df.shape[0])
+
+
+def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
+    """Build admits dataframe from dispositions."""
+    admits_df = dispositions_df.iloc[:-1, :] - dispositions_df.shift(1)
+    admits_df.day = dispositions_df.day
     return admits_df
 
 
 def build_census_df(
-    admits_df: pd.DataFrame, lengths_of_stay
+    admits_df: pd.DataFrame,
+    lengths_of_stay: Dict[str, int],
 ) -> pd.DataFrame:
-    """ALOS for each category of COVID-19 case (total guesses)"""
-    n_days = np.shape(admits_df)[0]
-    census_dict = {}
-    for key, los in lengths_of_stay.items():
-        census = (
-            admits_df.cumsum().iloc[:-los, :]
-            - admits_df.cumsum().shift(los).fillna(0)
-        ).apply(np.ceil)
-        census_dict[key] = census[key]
-
-    census_df = pd.DataFrame(census_dict)
-    census_df["day"] = census_df.index
-    census_df = census_df[["day", *lengths_of_stay.keys()]]
-    census_df = census_df.head(n_days)
-    return census_df
+    """ALOS for each disposition of COVID-19 case (total guesses)"""
+    return pd.DataFrame({
+        'day': admits_df.day,
+        **{
+            key: (
+                admits_df[key].cumsum().iloc[:-los]
+                - admits_df[key].cumsum().shift(los).fillna(0)
+            ).apply(np.ceil)
+            for key, los in lengths_of_stay.items()
+        }
+    })

tests/test_app.py

@@ -249,13 +249,12 @@ def test_model(model=MODEL, param=PARAM):
     assert round(last.susceptible, 0) == 67202
     assert round(raw_df.recovered[30], 0) == 224048
 
-    assert [d[0] for d in model.dispositions.values()] == [100.0, 40.0, 20.0]
-    assert [round(d[60], 0) for d in model.dispositions.values()] == [1182.0, 473.0, 236.0]
+    assert list(model.dispositions_df.iloc[0, :]) == [0, 100.0, 40.0, 20.0]
+    assert [round(i, 0) for i in model.dispositions_df.iloc[60, :]] == [60, 1182.0, 473.0, 236.0]
 
     # test that admissions are being properly calculated (thanks @PhilMiller)
-    admissions = build_admits_df(param.n_days, model.dispositions)
-    cumulative_admissions = admissions.cumsum()
-    diff = cumulative_admissions["hospitalized"][1:-1] - (
+    cumulative_admits = model.admits_df.cumsum()
+    diff = cumulative_admits.hospitalized[1:-1] - (
         0.05 * 0.05 * (raw_df.infected[1:-1] + raw_df.recovered[1:-1]) - 100
     )
     assert (diff.abs() < 0.1).all()