Make census counts inclusive over the relevant lengths of stay

src/penn_chime/models.py

@@ -296,7 +296,7 @@ def build_dispositions_df(
 
 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 = dispositions_df - dispositions_df.shift(1)
     admits_df.day = dispositions_df.day
     admits_df.date = dispositions_df.date
     return admits_df
@@ -307,14 +307,30 @@ def build_census_df(
     lengths_of_stay: Dict[str, int],
 ) -> pd.DataFrame:
     """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
+
+    hospital_admissions_exclusive = admits_df.hospitalized - admits_df.icu
+    icu_admissions_exclusive = admits_df.icu - admits_df.ventilated
+
+    census_vent = (admits_df.ventilated.cumsum() -
+                   admits_df.ventilated.cumsum().shift(lengths_of_stay['ventilated']).fillna(0)
+    ).apply(np.ceil)
+
+    # For the census to be inclusive, it needs to count each admission
+    # for the length-of-stay applicable to the level of care needed
+    census_icu =  (icu_admissions_exclusive.cumsum()
+                   - icu_admissions_exclusive.cumsum().shift(lengths_of_stay['icu']).fillna(0)
+                   + census_vent
+    ).apply(np.ceil)
+
+    census_hosp =  (hospital_admissions_exclusive.cumsum()
+                   - hospital_admissions_exclusive.cumsum().shift(lengths_of_stay['hospitalized']).fillna(0)
+                   + census_icu
+    ).apply(np.ceil)
+
     return pd.DataFrame({
         'day': admits_df.day,
         'date': admits_df.date,
-        **{
-            key: (
-                admits_df[key].cumsum()
-                - admits_df[key].cumsum().shift(los).fillna(0)
-            ).apply(np.ceil)
-            for key, los in lengths_of_stay.items()
-        }
+        'hospitalized': census_hosp,
+        'icu': census_icu,
+        'ventilated': census_vent,
     })