Merge pull request #182 from CodeForPhilly/112-model

112 model

Author: jlubken

src/app.py

@@ -1,28 +1,25 @@
 """App."""
 
-import altair as alt
-import streamlit as st
+import altair as alt  # type: ignore
+import streamlit as st  # type: ignore
 
 from penn_chime.presentation import (
-    additional_projections_chart,
-    admitted_patients_chart,
+    #additional_projections_chart,
     display_header,
     display_sidebar,
     display_n_days_slider,
     draw_census_table,
     draw_projected_admissions_table,
     draw_raw_sir_simulation_table,
     hide_menu_style,
-    new_admissions_chart,
     show_additional_projections,
     show_more_info_about_this_tool,
     write_definitions,
     write_footer,
 )
-from penn_chime.utils import build_admissions_df, build_census_df
-from penn_chime.settings import DEFAULTS
-
-
+from penn_chime.settings import DEFAULTS 
+from penn_chime.models import sim_sir_df, build_admissions_df, build_census_df
+from penn_chime.charts import additional_projections_chart, admitted_patients_chart, new_admissions_chart
 # This is somewhat dangerous:
 # Hide the main menu with "Rerun", "run on Save", "clear cache", and "record a screencast"
 # This should not be hidden in prod, but removed
@@ -31,35 +28,15 @@
 
 p = display_sidebar(st, DEFAULTS)
 
-display_header(
-    st,
-    total_infections=p.infected,
-    initial_infections=p.known_infected,
-    detection_prob=p.detection_probability,
-    current_hosp=p.current_hospitalized,
-    hosp_rate=p.hospitalized.rate,
-    S=p.susceptible,
-    market_share=p.market_share,
-    recovery_days=p.recovery_days,
-    r_naught=p.r_naught,
-    doubling_time=p.doubling_time,
-    relative_contact_rate=p.relative_contact_rate,
-    r_t=p.r_t,
-    doubling_time_t=p.doubling_time_t,
-)
+display_header(st, p)
+
 if st.checkbox("Show more info about this tool"):
     notes = "The total size of the susceptible population will be the entire catchment area for Penn Medicine entities (HUP, PAH, PMC, CCH)"
     show_more_info_about_this_tool(
         st=st,
-        recovery_days=p.recovery_days,
-        doubling_time=p.doubling_time,
-        r_naught=p.r_naught,
-        relative_contact_rate=p.relative_contact_rate,
-        doubling_time_t=p.doubling_time_t,
-        r_t=p.r_t,
+        parameters=p,
         inputs=DEFAULTS,
         notes=notes
-
     )
 
 # PRESENTATION
@@ -68,30 +45,36 @@
 display_n_days_slider(st, p, DEFAULTS)
 
 # begin format data
-projection_admits = build_admissions_df(p.n_days, *p.dispositions)
-census_df = build_census_df(projection_admits, *p.lengths_of_stay)
+admissions_df = build_admissions_df(p=p) # p.n_days, *p.dispositions)
+census_df = build_census_df(admissions_df, parameters=p)
 # end format data
 
 st.subheader("New Admissions")
 st.markdown("Projected number of **daily** COVID-19 admissions at Penn hospitals")
 st.altair_chart(
-    new_admissions_chart(alt, projection_admits, p.n_days - 10, as_date=as_date, max_y_axis=p.max_y_axis), use_container_width=True
+    new_admissions_chart(alt, admissions_df, parameters=p, as_date=as_date), use_container_width=True
 )
 if st.checkbox("Show Projected Admissions in tabular form"):
-    draw_projected_admissions_table(st, projection_admits, as_date=as_date)
+    draw_projected_admissions_table(st, admissions_df, as_date=as_date)
 st.subheader("Admitted Patients (Census)")
 st.markdown(
     "Projected **census** of COVID-19 patients, accounting for arrivals and discharges at Penn hospitals"
 )
-st.altair_chart(admitted_patients_chart(alt, census_df, p.n_days - 10, as_date=as_date, max_y_axis=p.max_y_axis), use_container_width=True)
+st.altair_chart(
+    admitted_patients_chart(alt=alt, census=census_df, parameters=p, as_date=as_date), use_container_width=True
+)
 if st.checkbox("Show Projected Census in tabular form"):
     draw_census_table(st, census_df, as_date=as_date)
 st.markdown(
     """**Click the checkbox below to view additional data generated by this simulation**"""
 )
 if st.checkbox("Show Additional Projections"):
-    show_additional_projections(st, alt, additional_projections_chart, p.infected_v, p.recovered_v, as_date=as_date, max_y_axis=p.max_y_axis)
+    show_additional_projections(
+        st, alt,
+        additional_projections_chart,
+        parameters=p,
+        as_date=as_date)
     if st.checkbox("Show Raw SIR Simulation Data"):
-        draw_raw_sir_simulation_table(st, p.n_days, p.susceptible_v, p.infected_v, p.recovered_v, as_date=as_date)
+        draw_raw_sir_simulation_table(st, parameters=p, as_date=as_date)
 write_definitions(st)
 write_footer(st)

src/penn_chime/charts.py

@@ -0,0 +1,125 @@
+
+from altair import Chart  # type: ignore
+import pandas as pd  # type: ignore
+import numpy as np  # type: ignore
+
+from .parameters import Parameters
+from .utils import add_date_column
+
+
+def new_admissions_chart(
+    alt,
+    projection_admits: pd.DataFrame,
+    parameters: Parameters,
+    as_date: bool = False,
+) -> Chart:
+    """docstring"""
+    plot_projection_days = parameters.n_days - 10
+    max_y_axis = parameters.max_y_axis
+
+    y_scale = alt.Scale()
+
+    if max_y_axis is not None:
+        y_scale.domain = (0, max_y_axis)
+        y_scale.clamp = True
+
+    tooltip_dict = {False: "day", True: "date:T"}
+    if as_date:
+        projection_admits = add_date_column(projection_admits)
+        x_kwargs = {"shorthand": "date:T", "title": "Date"}
+    else:
+        x_kwargs = {"shorthand": "day", "title": "Days from today"}
+
+    return (
+        alt.Chart(projection_admits.head(plot_projection_days))
+        .transform_fold(fold=["Hospitalized", "ICU", "Ventilated"])
+        .mark_line(point=True)
+        .encode(
+            x=alt.X(**x_kwargs),
+            y=alt.Y("value:Q", title="Daily admissions", scale=y_scale),
+            color="key:N",
+            tooltip=[
+                tooltip_dict[as_date],
+                alt.Tooltip("value:Q", format=".0f", title="Admissions"),
+                "key:N",
+            ],
+        )
+        .interactive()
+    )
+
+
+def admitted_patients_chart(
+    alt,
+    census: pd.DataFrame,
+    parameters: Parameters,
+    as_date: bool = False
+) -> Chart:
+    """docstring"""
+
+    plot_projection_days = parameters.n_days - 10
+    max_y_axis = parameters.max_y_axis
+    if as_date:
+        census = add_date_column(census)
+        x_kwargs = {"shorthand": "date:T", "title": "Date"}
+        idx = "date:T"
+    else:
+        x_kwargs ={"shorthand": "day", "title": "Days from today"}
+        idx = "day"
+
+    y_scale = alt.Scale()
+
+    if max_y_axis:
+        y_scale.domain = (0, max_y_axis)
+        y_scale.clamp = True
+
+    return (
+        alt.Chart(census.head(plot_projection_days))
+        .transform_fold(fold=["Hospitalized Census", "ICU Census", "Ventilated Census"])
+        .mark_line(point=True)
+        .encode(
+            x=alt.X(**x_kwargs),
+            y=alt.Y("value:Q", title="Census", scale=y_scale),
+            color="key:N",
+            tooltip=[
+                idx,
+                alt.Tooltip("value:Q", format=".0f", title="Census"),
+                "key:N",
+            ],
+        )
+        .interactive()
+    )
+
+
+def additional_projections_chart(
+    alt,
+    i: np.ndarray,
+    r: np.ndarray,
+    as_date: bool = False,
+    max_y_axis: int = None
+) -> Chart:
+    dat = pd.DataFrame({"Infected": i, "Recovered": r})
+    dat["day"] = dat.index
+    if as_date:
+        dat = add_date_column(dat)
+        x_kwargs = {"shorthand": "date:T", "title": "Date"}
+    else:
+        x_kwargs = {"shorthand": "day", "title": "Days from today"}
+
+    y_scale = alt.Scale()
+
+    if max_y_axis is not None:
+        y_scale.domain = (0, max_y_axis)
+        y_scale.clamp = True
+
+    return (
+        alt.Chart(dat)
+        .transform_fold(fold=["Infected", "Recovered"])
+        .mark_line()
+        .encode(
+            x=alt.X(**x_kwargs),
+            y=alt.Y("value:Q", title="Case Volume", scale=y_scale),
+            tooltip=["key:N", "value:Q"],
+            color="key:N",
+        )
+        .interactive()
+    )

src/penn_chime/models.py

@@ -2,8 +2,8 @@
 
 from typing import Generator, Tuple
 
-import numpy as np
-import pandas as pd
+import numpy as np  # type: ignore
+import pandas as pd  # type: ignore
 
 
 def sir(
@@ -58,13 +58,13 @@ def sim_sir(
     )
 
 
-def sim_sir_df(
-    s: float, i: float, r: float,
-    beta: float, gamma: float, n_days: int
-) -> pd.DataFrame:
-    """Simulate the SIR model forward in time."""
+def sim_sir_df(p) -> pd.DataFrame:
+    """Simulate the SIR model forward in time.
+
+    p is a Parameters instance. for circuluar dependency reasons i can't annotate it.
+    """
     return pd.DataFrame(
-        data=gen_sir(s, i, r, beta, gamma, n_days),
+        data=gen_sir(p.susceptible, p.infected, p.recovered, p.beta, p.gamma, p.n_days),
         columns=("Susceptible", "Infected", "Recovered"),
     )
 
@@ -74,3 +74,51 @@ def get_dispositions(
 ) -> Tuple[np.ndarray, ...]:
     """Get dispositions of infected adjusted by rate and market_share."""
     return (*(infected * rate * market_share for rate in rates),)
+
+
+
+def build_admissions_df(p) -> pd.DataFrame:
+    """Build admissions dataframe from Parameters."""
+    days = np.array(range(0, p.n_days + 1))
+    data_dict = dict(zip(["day", "Hospitalized", "ICU", "Ventilated"],
+                         [days] + [disposition for disposition in p.dispositions]
+    ))
+    projection = pd.DataFrame.from_dict(data_dict)
+    # New cases
+    projection_admits = projection.iloc[:-1, :] - projection.shift(1)
+    projection_admits[projection_admits < 0] = 0
+    projection_admits["day"] = range(projection_admits.shape[0])
+    return projection_admits
+
+
+def build_census_df(
+        projection_admits: pd.DataFrame,
+        parameters
+) -> pd.DataFrame:
+    """ALOS for each category of COVID-19 case (total guesses)"""
+    n_days = np.shape(projection_admits)[0]
+    hosp_los, icu_los, vent_los = parameters.lengths_of_stay
+    los_dict = {
+        "Hospitalized": hosp_los,
+        "ICU": icu_los,
+        "Ventilated": vent_los,
+    }
+
+    census_dict = dict()
+    for k, los in los_dict.items():
+        census = (
+            projection_admits.cumsum().iloc[:-los, :]
+            - projection_admits.cumsum().shift(los).fillna(0)
+        ).apply(np.ceil)
+        census_dict[k] = census[k]
+
+    census_df = pd.DataFrame(census_dict)
+    census_df["day"] = census_df.index
+    census_df = census_df[["day", "Hospitalized", "ICU", "Ventilated"]]
+    census_df = census_df.head(n_days)
+    census_df = census_df.rename(
+        columns={disposition: f"{disposition} Census"
+                 for disposition
+                 in ("Hospitalized", "ICU", "Ventilated")}
+    )
+    return census_df

src/penn_chime/parameters.py

@@ -1,6 +1,6 @@
 """Parameters."""
 
-from numpy import log2
+from numpy import log2  # type: ignore
 
 from .utils import RateLos
 from .models import (