Repository: CodeForPhilly/chime
Pull request #471 (2020-04-03 23:23:08)

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


Patch:
diff --git a/src/app.py b/src/app.py
index 5fad23e..94b7da5 100644
--- a/src/app.py
+++ b/src/app.py
@@ -36,7 +36,7 @@ st.subheader("New Admissions")
 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",
@@ -58,7 +58,7 @@ st.subheader("Admitted Patients (Census)")
 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",
diff --git a/src/chime_dash/app/services/callbacks.py b/src/chime_dash/app/services/callbacks.py
index 347ff6a..d512359 100644
--- a/src/chime_dash/app/services/callbacks.py
+++ b/src/chime_dash/app/services/callbacks.py
@@ -38,10 +38,14 @@ class IndexCallbacks(ComponentCallbacks):
         if sidebar_data:
             pars = parameters_deserializer(sidebar_data["parameters"])
             model = SimSirModel(pars)
+            vis = i.components.get("visualizations", None) if i else None
+            vis_content = vis.content if vis else None
+
             viz_kwargs = dict(
                 labels=pars.labels,
                 table_mod=7,
                 max_y_axis=pars.max_y_axis,
+                content=vis_content
             )
         result.extend(i.components["intro"].build(model, pars))
         for df_key in ["admits_df", "census_df", "sim_sir_w_date_df"]:
diff --git a/src/chime_dash/app/services/plotting.py b/src/chime_dash/app/services/plotting.py
index c8c6c5c..8a9609c 100644
--- a/src/chime_dash/app/services/plotting.py
+++ b/src/chime_dash/app/services/plotting.py
@@ -7,9 +7,17 @@ from typing import Dict, Any
 from pandas import DataFrame
 
 
-def plot_dataframe(dataframe: DataFrame, max_y_axis: int = None,) -> Dict[str, Any]:
-    """
+def plot_dataframe(
+    dataframe: DataFrame,
+    max_y_axis: int = None,
+) -> Dict[str, Any]:
+    """Returns dictionary used for plotly graphs
+
+    Arguments:
+        dataframe: The dataframe to plot. Plots all columns as y, index is x.
+        max_y_axis: Maximal value on y-axis.
     """
+
     if max_y_axis is None:
         yaxis = {}
     else:
diff --git a/src/chime_dash/app/templates/en/visualizations.yml b/src/chime_dash/app/templates/en/visualizations.yml
index 5211a77..31a7775 100644
--- a/src/chime_dash/app/templates/en/visualizations.yml
+++ b/src/chime_dash/app/templates/en/visualizations.yml
@@ -16,5 +16,20 @@ hosp: Hospitalized
 icu: ICU
 vent: Ventalized
 
+# Admits df column translations
+admits_hospitalized: Hospitalized
+admits_icu: ICU
+admits_ventilated: Ventilated
+
+# Admits df column translations
+census_hospitalized: Hospitalized
+census_icu: ICU
+census_ventilated: Ventilated
+
+# SIR df column translations
+susceptible: Susceptible
+infected: Infected
+recovered: Recovered
+
 # Date Localization
 date-format: "%m%d%Y"
diff --git a/src/chime_dash/app/utils/__init__.py b/src/chime_dash/app/utils/__init__.py
index 612f470..b9f8519 100644
--- a/src/chime_dash/app/utils/__init__.py
+++ b/src/chime_dash/app/utils/__init__.py
@@ -116,18 +116,50 @@ def get_n_switch_values(input_value, elements_to_update) -> List[bool]:
 
 def prepare_visualization_group(df: DataFrame = None, **kwargs) -> List[Any]:
     """Creates plot, table and download link for data frame.
+
+    Arguments:
+        df: The Dataframe to plot
+        content: Dict[str, str]
+            Mapping for translating columns and index.
+        max_y_axis:  int
+            Maximal value on y-axis
+        labels: List[str]
+            Columns to display
+        table_mod: int
+            Displays only each `table_mod` row in table
+
     """
     result = [{}, None, None]
     if df is not None and isinstance(df, DataFrame):
+
+        date_column = "date"
+        day_column = "day"
+
+        # Translate column and index if specified
+        content = kwargs.get("content", None)
+        if content:
+            columns = {col: content[col] for col in df.columns if col in content}
+            index = (
+                {df.index.name: content[df.index.name]}
+                if df.index.name and df.index.name in content
+                else None
+            )
+            df = df.rename(columns=columns, index=index)
+            date_column = content.get(date_column, date_column)
+            day_column = content.get(day_column, day_column)
+
         plot_data = plot_dataframe(
-            df.dropna().set_index("date").drop(columns=["day"]),
+            df.dropna().set_index(date_column).drop(columns=[day_column]),
             max_y_axis=kwargs.get("max_y_axis", None),
         )
 
+
+        # translate back for backwards compability of build_table
+        column_map = {day_column: "day", date_column: "date"}
         table = (
             df_to_html_table(
                 build_table(
-                    df=df,
+                    df=df.rename(columns=column_map),
                     labels=kwargs.get("labels", df.columns),
                     modulo=kwargs.get("table_mod", 7),
                 ),
@@ -140,7 +172,9 @@ def prepare_visualization_group(df: DataFrame = None, **kwargs) -> List[Any]:
             # else None
         )
 
-        csv = build_csv_download(df)
+        # Convert columnnames to lowercase
+        column_map = {col: col.lower() for col in df.columns}
+        csv = build_csv_download(df.rename(columns=column_map))
         result = [plot_data, table, csv]
 
     return result
@@ -153,4 +187,5 @@ def singleton(class_):
         if class_ not in instances:
             instances[class_] = class_(*args, **kwargs)
         return instances[class_]
+
     return get_instance
diff --git a/src/penn_chime/charts.py b/src/penn_chime/charts.py
index d928973..1d21c13 100644
--- a/src/penn_chime/charts.py
+++ b/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 "",
             )
diff --git a/src/penn_chime/models.py b/src/penn_chime/models.py
index 136f49b..e47bf5e 100644
--- a/src/penn_chime/models.py
+++ b/src/penn_chime/models.py
@@ -69,10 +69,10 @@ class SimSirModel:
             self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
 
             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))
+            raw = self.run_projection(p, [(self.beta, p.n_days)])
+            self.i_day = i_day = int(get_argmin_ds(raw["census_hospitalized"], p.current_hospitalized))
 
-            self.run_projection(p, self.gen_policy(p))
+            self.raw = self.run_projection(p, self.gen_policy(p))
 
             logger.info('Set i_day = %s', i_day)
             p.date_first_hospitalized = p.current_date - timedelta(days=i_day)
@@ -109,7 +109,7 @@ class SimSirModel:
             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.gen_policy(p))
+            self.raw = self.run_projection(p, self.gen_policy(p))
 
             self.population = p.population
         else:
@@ -120,6 +120,31 @@ class SimSirModel:
             )
             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 = pd.DataFrame(data={
+            'day': self.raw['day'],
+            'date': self.raw['date'],
+            'ever_hospitalized': self.raw['ever_hospitalized'],
+            'ever_icu': self.raw['ever_icu'],
+            'ever_ventilated': self.raw['ever_ventilated'],
+        })
+        self.admits_df = pd.DataFrame(data={
+            'day': self.raw['day'],
+            'date': self.raw['date'],
+            'admits_hospitalized': self.raw['admits_hospitalized'],
+            'admits_icu': self.raw['admits_icu'],
+            'admits_ventilated': self.raw['admits_ventilated'],
+        })
+        self.census_df = pd.DataFrame(data={
+            'day': self.raw['day'],
+            'date': self.raw['date'],
+            'census_hospitalized': self.raw['census_hospitalized'],
+            'census_icu': self.raw['census_icu'],
+            'census_ventilated': self.raw['census_ventilated'],
+        })
+
         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 +164,9 @@ class SimSirModel:
 
         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)
@@ -153,14 +178,15 @@ class SimSirModel:
             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.gen_policy(p))
+            raw = 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 = raw["admits_hospitalized"].argmax()
             if peak_admits_day < 0:
                 continue
 
-            loss = self.get_loss()
+            predicted = raw["census_hospitalized"][self.i_day]
+            loss = get_loss(self.current_hospitalized, predicted)
             losses[i] = loss
 
         min_loss = pd.Series(losses).argmin()
@@ -186,7 +212,7 @@ class SimSirModel:
         ]
 
     def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
-        self.raw_df = sim_sir_df(
+        raw = sim_sir(
             self.susceptible,
             self.infected,
             p.recovered,
@@ -195,23 +221,24 @@ class SimSirModel:
             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(raw, self.rates, p.market_share)
+        calculate_admits(raw, self.rates)
+        calculate_census(raw, self.days)
+
+        return raw
 
-    def get_loss(self) -> float:
-        """Squared error: predicted vs. actual current hospitalized."""
-        predicted = self.census_df.hospitalized.loc[self.i_day]
-        return (self.current_hospitalized - predicted) ** 2.0
 
+def get_loss(current_hospitalized, predicted) -> float:
+    """Squared error: predicted vs. actual current hospitalized."""
+    return (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(
@@ -257,33 +284,49 @@ def sir(
     return s_n * scale, i_n * scale, r_n * scale
 
 
-def gen_sir(
+def sim_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.
+):
+    """Simulate SIR model forward in time, returning a dictionary of daily arrays
     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
 
-
-def sim_sir_df(
-    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"),
-    )
+    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 build_sim_sir_w_date_df(
@@ -302,58 +345,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(raw: Dict, rates):
     """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


Test patch:
diff --git a/tests/by_doubling_time/2020-03-28_projected_admits.csv b/tests/by_doubling_time/2020-03-28_projected_admits.csv
index d872896..80fc5cd 100644
--- a/tests/by_doubling_time/2020-03-28_projected_admits.csv
+++ b/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
diff --git a/tests/by_doubling_time/2020-03-28_projected_census.csv b/tests/by_doubling_time/2020-03-28_projected_census.csv
index 789baf0..a627d2a 100644
--- a/tests/by_doubling_time/2020-03-28_projected_census.csv
+++ b/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
diff --git a/tests/conftest.py b/tests/conftest.py
index b7cf01f..0b98962 100644
--- a/tests/conftest.py
+++ b/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_")
 
diff --git a/tests/penn_chime/test_charts.py b/tests/penn_chime/test_charts.py
index 69c9576..0bcd119 100644
--- a/tests/penn_chime/test_charts.py
+++ b/tests/penn_chime/test_charts.py
@@ -19,7 +19,7 @@ DISPOSITION_KEYS = ("hospitalized", "icu", "ventilated")
 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,11 +28,11 @@ 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
 
 
@@ -40,18 +40,18 @@ 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
     )
 
@@ -59,8 +59,8 @@ def test_census(census_floor_df, param):
 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):
diff --git a/tests/penn_chime/test_models.py b/tests/penn_chime/test_models.py
index 4b99a0d..e159874 100644
--- a/tests/penn_chime/test_models.py
+++ b/tests/penn_chime/test_models.py
@@ -7,7 +7,7 @@ from datetime import timedelta
 
 from src.penn_chime.models import (
     sir,
-    sim_sir_df,
+    sim_sir,
     get_growth_rate,
     SimSirModel,
 )
@@ -65,21 +65,16 @@ def test_sim_sir():
     """
     Rounding to move fast past decimal place issues
     """
-    raw_df = sim_sir_df(
+    raw = sim_sir(
         5, 6, 7, 0.1, 0, [(0.1, 40)],  # s  # i  # r  # gamma  # i_day  # beta1  # n_days1
     )
 
-    first = raw_df.iloc[0, :]
-    last = raw_df.iloc[-1, :]
-
-    assert round(first.susceptible, 0) == 5
-    assert round(first.infected, 2) == 6
-    assert round(first.recovered, 0) == 7
-    assert round(last.susceptible, 2) == 0
-    assert round(last.infected, 2) == 0.18
-    assert round(last.recovered, 2) == 17.82
-
-    assert isinstance(raw_df, pd.DataFrame)
+    assert round(raw["susceptible"][0], 0) == 5
+    assert round(raw["infected"][0], 2) == 6
+    assert round(raw["recovered"][0], 0) == 7
+    assert round(raw["susceptible"][-1], 2) == 0
+    assert round(raw["infected"][-1], 2) == 0.18
+    assert round(raw["recovered"][-1], 2) == 17.82
 
 
 def test_growth_rate():
@@ -131,7 +126,13 @@ def test_model_raw_start(model, param):
 
     assert first.susceptible == 499600.0
     assert round(second.infected, 0) == 449.0
-    assert list(model.dispositions_df.iloc[0, :]) == [
+    assert list(model.dispositions_df.loc[0, [
+        "day",
+        "date",
+        "ever_hospitalized",
+        "ever_icu",
+        "ever_ventilated",
+    ]]) == [
         -43,
         date(year=2020, month=2, day=14),
         1.0,
@@ -140,9 +141,15 @@ def test_model_raw_start(model, param):
     ]
     assert round(raw_df.recovered[30], 0) == 7083.0
 
-    d, dt, s, i, r = list(model.dispositions_df.iloc[60, :])
+    d, dt, hosp, icu, vent = list(model.dispositions_df.loc[60, [
+        "day",
+        "date",
+        "ever_hospitalized",
+        "ever_icu",
+        "ever_ventilated",
+    ]])
     assert dt == date(year=2020, month=4, day=14)
-    assert [round(v, 0) for v in (d, s, i, r)] == [17, 549.0, 220.0, 110.0]
+    assert [round(v, 0) for v in (d, hosp, icu, vent)] == [17, 549.0, 220.0, 110.0]
 
 
 def test_model_conservation(param, model):
@@ -180,9 +187,9 @@ def test_model_cumulative_census(param, model):
     admits_df = model.admits_df
     df = pd.DataFrame(
         {
-            "hospitalized": admits_df.hospitalized,
-            "icu": admits_df.icu,
-            "ventilated": admits_df.ventilated,
+            "hospitalized": admits_df.admits_hospitalized,
+            "icu": admits_df.admits_icu,
+            "ventilated": admits_df.admits_ventilated,
         }
     )
     admits = df.cumsum()