Speed up by 20x by never copying arrays and only constructing a dataframe after the model has run

Author: PhilMiller

src/app.py

@@ -37,7 +37,7 @@
 st.markdown("Projected number of **daily** COVID-19 admissions. \n\n _NOTE: Now including estimates of prior admissions for comparison._")
 admits_chart = build_admits_chart(alt=alt, admits_floor_df=m.admits_floor_df, max_y_axis=p.max_y_axis)
 st.altair_chart(admits_chart, use_container_width=True)
-st.markdown(build_descriptions(chart=admits_chart, labels=p.labels, suffix=" Admissions"))
+st.markdown(build_descriptions(chart=admits_chart, labels=p.labels, prefix="admits_", suffix=" Admissions"))
 display_download_link(
     st,
     filename=f"{p.current_date}_projected_admits.csv",
@@ -59,7 +59,7 @@
 st.markdown("Projected **census** of COVID-19 patients, accounting for arrivals and discharges \n\n _NOTE: Now including estimates of prior census for comparison._")
 census_chart = build_census_chart(alt=alt, census_floor_df=m.census_floor_df, max_y_axis=p.max_y_axis)
 st.altair_chart(census_chart, use_container_width=True)
-st.markdown(build_descriptions(chart=census_chart, labels=p.labels, suffix=" Census"))
+st.markdown(build_descriptions(chart=census_chart, labels=p.labels, prefix="census_", suffix=" Census"))
 display_download_link(
     st,
     filename=f"{p.current_date}_projected_census.csv",

src/penn_chime/charts.py

@@ -26,7 +26,7 @@ def build_admits_chart(
     # TODO fix the fold to allow any number of dispositions
     points = (
         alt.Chart()
-        .transform_fold(fold=["hospitalized", "icu", "ventilated"])
+        .transform_fold(fold=["admits_hospitalized", "admits_icu", "admits_ventilated"])
         .encode(x=alt.X(**x), y=alt.Y(**y), color=color, tooltip=tooltip)
         .mark_line(point=True)
         .encode(
@@ -65,7 +65,7 @@ def build_census_chart(
     # TODO fix the fold to allow any number of dispositions
     points = (
         alt.Chart()
-        .transform_fold(fold=["hospitalized", "icu", "ventilated"])
+        .transform_fold(fold=["census_hospitalized", "census_icu", "census_ventilated"])
         .encode(x=alt.X(**x), y=alt.Y(**y), color=color, tooltip=tooltip)
         .mark_line(point=True)
         .encode(
@@ -128,7 +128,11 @@ def build_sim_sir_w_date_chart(
 
 
 def build_descriptions(
-    *, chart: Chart, labels: Dict[str, str], suffix: str = ""
+    *,
+    chart: Chart,
+    labels: Dict[str, str],
+    prefix: str = "",
+    suffix: str = ""
 ) -> str:
     """
 
@@ -145,17 +149,17 @@ def build_descriptions(
     day = "date" if "date" in chart.data.columns else "day"
 
     for col in cols:
-        if chart.data[col].idxmax() + 1 == len(chart.data):
+        if chart.data[prefix+col].idxmax() + 1 == len(chart.data):
             asterisk = True
 
         # todo: bring this to an optional arg / i18n
-        on = datetime.strftime(chart.data[day][chart.data[col].idxmax()], "%b %d")
+        on = datetime.strftime(chart.data[day][chart.data[prefix+col].idxmax()], "%b %d")
 
         messages.append(
             "{}{} peaks at {:,} on {}{}".format(
                 labels[col],
                 suffix,
-                ceil(chart.data[col].max()),
+                ceil(chart.data[prefix+col].max()),
                 on,
                 "*" if asterisk else "",
             )

src/penn_chime/models.py

@@ -70,7 +70,7 @@ def __init__(self, p: Parameters):
 
             self.i_day = 0 # seed to the full length
             self.run_projection(p, [(self.beta, p.n_days)])
-            self.i_day = i_day = int(get_argmin_ds(self.census_df, p.current_hospitalized))
+            self.i_day = i_day = int(get_argmin_ds(self.raw["census_hospitalized"], p.current_hospitalized))
 
             self.run_projection(p, self.gen_policy(p))
 
@@ -120,6 +120,13 @@ def __init__(self, p: Parameters):
             )
             raise AssertionError('doubling_time or date_first_hospitalized must be provided.')
 
+        self.raw["date"] = self.raw["day"].astype("timedelta64[D]") + np.datetime64(p.current_date)
+
+        self.raw_df = pd.DataFrame(data=self.raw)
+        self.dispositions_df = self.raw_df
+        self.admits_df = self.raw_df
+        self.census_df = self.raw_df
+
         logger.info('len(np.arange(-i_day, n_days+1)): %s', len(np.arange(-self.i_day, p.n_days+1)))
         logger.info('len(raw_df): %s', len(self.raw_df))
 
@@ -139,9 +146,9 @@ def __init__(self, p: Parameters):
 
         self.sim_sir_w_date_df = build_sim_sir_w_date_df(self.raw_df, p.current_date, self.keys)
 
-        self.sim_sir_w_date_floor_df = build_floor_df(self.sim_sir_w_date_df, self.keys)
-        self.admits_floor_df = build_floor_df(self.admits_df, p.dispositions.keys())
-        self.census_floor_df = build_floor_df(self.census_df, p.dispositions.keys())
+        self.sim_sir_w_date_floor_df = build_floor_df(self.sim_sir_w_date_df, self.keys, "")
+        self.admits_floor_df = build_floor_df(self.admits_df, p.dispositions.keys(), "admits_")
+        self.census_floor_df = build_floor_df(self.census_df, p.dispositions.keys(), "census_")
 
         self.daily_growth_rate = get_growth_rate(p.doubling_time)
         self.daily_growth_rate_t = get_growth_rate(self.doubling_time_t)
@@ -156,7 +163,7 @@ def get_argmin_doubling_time(self, p: Parameters, dts):
             self.run_projection(p, self.gen_policy(p))
 
             # Skip values the would put the fit past peak
-            peak_admits_day = self.admits_df.hospitalized.argmax()
+            peak_admits_day = self.raw["admits_hospitalized"].argmax()
             if peak_admits_day < 0:
                 continue
 
@@ -186,7 +193,7 @@ def gen_policy(self, p: Parameters) -> Sequence[Tuple[float, int]]:
         ]
 
     def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
-        self.raw_df = sim_sir_df(
+        self.raw = sim_sir(
             self.susceptible,
             self.infected,
             p.recovered,
@@ -195,23 +202,24 @@ def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
             policy
         )
 
-        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)
-        self.census_df = build_census_df(self.admits_df, self.days)
-        self.current_infected = self.raw_df.infected.loc[self.i_day]
+        calculate_dispositions(self.raw, self.rates, p.market_share)
+        calculate_admits(self.rates, self.raw)
+        calculate_census(self.raw, self.days)
+
+        self.current_infected = self.raw["infected"][self.i_day]
 
     def get_loss(self) -> float:
         """Squared error: predicted vs. actual current hospitalized."""
-        predicted = self.census_df.hospitalized.loc[self.i_day]
+        predicted = self.raw["census_hospitalized"][self.i_day]
         return (self.current_hospitalized - predicted) ** 2.0
 
 
-def get_argmin_ds(census_df: pd.DataFrame, current_hospitalized: float) -> float:
+def get_argmin_ds(census, current_hospitalized: float) -> float:
     # By design, this forbids choosing a day after the peak
     # If that's a problem, see #381
-    peak_day = census_df.hospitalized.argmax()
-    losses_df = (census_df.hospitalized[:peak_day] - current_hospitalized) ** 2.0
-    return losses_df.argmin()
+    peak_day = census.argmax()
+    losses = (census[:peak_day] - current_hospitalized) ** 2.0
+    return losses.argmin()
 
 
 def get_beta(
@@ -259,31 +267,56 @@ def sir(
 
 def gen_sir(
     s: float, i: float, r: float, gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
-) -> Generator[Tuple[int, float, float, float], None, None]:
+):
     """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
+
+    total_days = 1
+    for beta, days in policies:
+        total_days += days
+
+    d_a = np.empty(total_days, "int")
+    s_a = np.empty(total_days, "float")
+    i_a = np.empty(total_days, "float")
+    r_a = np.empty(total_days, "float")
+
+    index = 0
     for beta, n_days in policies:
         for _ in range(n_days):
-            yield d, s, i, r
+            d_a[index] = d
+            s_a[index] = s
+            i_a[index] = i
+            r_a[index] = r
+            index += 1
+
             s, i, r = sir(s, i, r, beta, gamma, n)
             d += 1
-    yield d, s, i, r
 
+    d_a[index] = d
+    s_a[index] = s
+    i_a[index] = i
+    r_a[index] = r
+    return {
+        "day": d_a,
+        "susceptible": s_a,
+        "infected": i_a,
+        "recovered": r_a,
+        "ever_infected": i_a + r_a
+    }
 
-def sim_sir_df(
+
+def sim_sir(
     s: float, i: float, r: float,
     gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
 ) -> pd.DataFrame:
     """Simulate the SIR model forward in time."""
-    return pd.DataFrame(
-        data=gen_sir(s, i, r, gamma, i_day, policies),
-        columns=("day", "susceptible", "infected", "recovered"),
-    )
+    data = gen_sir(s, i, r, gamma, i_day, policies)
+    return data
 
 
 def build_sim_sir_w_date_df(
@@ -302,58 +335,50 @@ def build_sim_sir_w_date_df(
     })
 
 
-def build_floor_df(df, keys):
+def build_floor_df(df, keys, prefix):
     """Build floor sim sir w date."""
     return pd.DataFrame({
         "day": df.day,
         "date": df.date,
         **{
-            key: np.floor(df[key])
+            prefix + key: np.floor(df[prefix+key])
             for key in keys
         }
     })
 
 
-def build_dispositions_df(
-    raw_df: pd.DataFrame,
+def calculate_dispositions(
+    raw: Dict,
     rates: Dict[str, float],
     market_share: float,
-    current_date: datetime,
-) -> pd.DataFrame:
+):
     """Build dispositions dataframe of patients adjusted by rate and market_share."""
-    patients = raw_df.infected + raw_df.recovered
-    day = raw_df.day
-    return pd.DataFrame({
-        "day": day,
-        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
-        **{
-            key: patients * rate * market_share
-            for key, rate in rates.items()
-        }
-    })
+    for key, rate in rates.items():
+        raw["ever_" + key] = raw["ever_infected"] * rate * market_share
+        raw[key] = raw["ever_infected"] * rate * market_share
 
 
-def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
+def calculate_admits(rates, raw: Dict):
     """Build admits dataframe from dispositions."""
-    admits_df = dispositions_df - dispositions_df.shift(1)
-    admits_df.day = dispositions_df.day
-    admits_df.date = dispositions_df.date
-    return admits_df
+    for key in rates.keys():
+        ever = raw["ever_" + key]
+        admit = np.empty_like(ever)
+        admit[0] = np.nan
+        admit[1:] = ever[1:] - ever[:-1]
+        raw["admits_"+key] = admit
+        raw[key] = admit
 
 
-def build_census_df(
-    admits_df: pd.DataFrame,
+def calculate_census(
+    raw: Dict,
     lengths_of_stay: Dict[str, int],
-) -> pd.DataFrame:
+):
     """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
-    return pd.DataFrame({
-        'day': admits_df.day,
-        'date': admits_df.date,
-        **{
-            key: (
-                admits_df[key].cumsum()
-                - admits_df[key].cumsum().shift(los).fillna(0)
-            )
-            for key, los in lengths_of_stay.items()
-        }
-    })
+    n_days = raw["day"].shape[0]
+    for key, los in lengths_of_stay.items():
+        cumsum = np.empty(n_days + los)
+        cumsum[:los+1] = 0.0
+        cumsum[los+1:] = raw["admits_" + key][1:].cumsum()
+
+        census = cumsum[los:] - cumsum[:-los]
+        raw["census_" + key] = census

tests/by_doubling_time/2020-03-28_projected_admits.csv

@@ -1,4 +1,4 @@
-,day,date,hospitalized,icu,ventilated
+,day,date,admits_hospitalized,admits_icu,admits_ventilated
 0,-4,2020-03-24,,,
 1,-3,2020-03-25,2.5542297270266676,0.7662689181079996,0.5108459454053333
 2,-2,2020-03-26,2.8373214956844457,0.8511964487053332,0.5674642991368888

tests/by_doubling_time/2020-03-28_projected_census.csv

@@ -1,4 +1,4 @@
-,day,date,hospitalized,icu,ventilated
+,day,date,census_hospitalized,census_icu,census_ventilated
 0,-4,2020-03-24,,,
 1,-3,2020-03-25,3.0,1.0,1.0
 2,-2,2020-03-26,6.0,2.0,2.0

tests/conftest.py

@@ -112,7 +112,7 @@ def admits_df():
 
 @pytest.fixture
 def admits_floor_df(param, admits_df):
-    return build_floor_df(admits_df, param.dispositions.keys())
+    return build_floor_df(admits_df, param.dispositions.keys(), "admits_")
 
 
 @pytest.fixture
@@ -123,5 +123,5 @@ def census_df():
 
 @pytest.fixture
 def census_floor_df(param, census_df):
-    return build_floor_df(census_df, param.dispositions.keys())
+    return build_floor_df(census_df, param.dispositions.keys(), "census_")
 

tests/penn_chime/test_charts.py

@@ -19,7 +19,7 @@
 def test_admits_chart(admits_floor_df):
     chart = build_admits_chart(alt=alt, admits_floor_df=admits_floor_df)
     assert isinstance(chart, (alt.Chart, alt.LayerChart))
-    assert round(chart.data.iloc[40].icu, 0) == 38
+    assert round(chart.data.iloc[40].admits_icu, 0) == 38
 
     # test fx call with no params
     with pytest.raises(TypeError):
@@ -28,39 +28,39 @@ def test_admits_chart(admits_floor_df):
 
 def test_build_descriptions(admits_floor_df, param):
     chart = build_admits_chart(alt=alt, admits_floor_df=admits_floor_df)
-    description = build_descriptions(chart=chart, labels=param.labels)
+    description = build_descriptions(chart=chart, labels=param.labels, prefix="admits_")
 
     hosp, icu, vent = description.split("\n\n")  # break out the description into lines
 
-    max_hosp = chart.data["hospitalized"].max()
+    max_hosp = chart.data["admits_hospitalized"].max()
     assert str(ceil(max_hosp)) in hosp
 
 
 def test_no_asterisk(admits_floor_df, param):
     param.n_days = 600
 
     chart = build_admits_chart(alt=alt, admits_floor_df=admits_floor_df)
-    description = build_descriptions(chart=chart, labels=param.labels)
+    description = build_descriptions(chart=chart, labels=param.labels, prefix="admits_")
     assert "*" not in description
 
 
 def test_census(census_floor_df, param):
     chart = build_census_chart(alt=alt, census_floor_df=census_floor_df)
-    description = build_descriptions(chart=chart, labels=param.labels)
+    description = build_descriptions(chart=chart, labels=param.labels, prefix="census_")
 
-    assert str(ceil(chart.data["ventilated"].max())) in description
-    assert str(chart.data["icu"].idxmax()) not in description
+    assert str(ceil(chart.data["census_ventilated"].max())) in description
+    assert str(chart.data["census_icu"].idxmax()) not in description
     assert (
-        datetime.strftime(chart.data.iloc[chart.data["icu"].idxmax()].date, "%b %d")
+        datetime.strftime(chart.data.iloc[chart.data["census_icu"].idxmax()].date, "%b %d")
         in description
     )
 
 
 def test_census_chart(census_floor_df):
     chart = build_census_chart(alt=alt, census_floor_df=census_floor_df)
     assert isinstance(chart, (alt.Chart, alt.LayerChart))
-    assert chart.data.iloc[1].hospitalized == 3
-    assert chart.data.iloc[49].ventilated == 365
+    assert chart.data.iloc[1].census_hospitalized == 3
+    assert chart.data.iloc[49].census_ventilated == 365
 
     # test fx call with no params
     with pytest.raises(TypeError):