Merge pull request #208 from CodeForPhilly/move_computation

Move computation

Author: mdbecker

src/app.py

@@ -17,11 +17,13 @@
     write_footer,
 )
 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,
-                               chart_descriptions)
+from penn_chime.models import SimSirModel
+from penn_chime.charts import (
+    additional_projections_chart,
+    admitted_patients_chart,
+    new_admissions_chart,
+    chart_descriptions
+)
 
 # This is somewhat dangerous:
 # Hide the main menu with "Rerun", "run on Save", "clear cache", and "record a screencast"
@@ -30,55 +32,52 @@
 st.markdown(hide_menu_style, unsafe_allow_html=True)
 
 p = display_sidebar(st, DEFAULTS)
+m = SimSirModel(p)
 
-display_header(st, p)
+display_header(st, m, 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, parameters=p, inputs=DEFAULTS, notes=notes)
-
-
-# begin format data
-admissions_df = build_admissions_df(p=p)  # p.n_days, *p.dispositions)
-census_df = build_census_df(admissions_df, parameters=p)
-# end format data
+    show_more_info_about_this_tool(st=st, model=m, parameters=p, defaults=DEFAULTS, notes=notes)
 
 st.subheader("New Admissions")
 st.markdown("Projected number of **daily** COVID-19 admissions at Penn hospitals")
-new_admit_chart = new_admissions_chart(alt, admissions_df, parameters=p)
+new_admit_chart = new_admissions_chart(alt, m.admits_df, parameters=p)
 st.altair_chart(
-    new_admit_chart, use_container_width=True
+    new_admissions_chart(alt, m.admits_df, parameters=p),
+    use_container_width=True,
 )
 
-st.markdown(chart_descriptions(new_admit_chart))
+st.markdown(chart_descriptions(new_admit_chart, p.labels))
 
 if st.checkbox("Show Projected Admissions in tabular form"):
     if st.checkbox("Show Daily Counts"):
-        draw_projected_admissions_table(st, admissions_df, as_date=p.as_date, daily_count=True)
+        draw_projected_admissions_table(st, m.admits_df, p.labels, as_date=p.as_date, daily_count=True)
     else:
-        draw_projected_admissions_table(st, admissions_df, as_date=p.as_date, daily_count=False)
+        draw_projected_admissions_table(st, m.admits_df, p.labels, as_date=p.as_date, daily_count=False)
     build_download_link(st,
         filename="projected_admissions.csv",
-        df=admissions_df,
+        df=m.admits_df,
         parameters=p
     )
 st.subheader("Admitted Patients (Census)")
 st.markdown(
     "Projected **census** of COVID-19 patients, accounting for arrivals and discharges at Penn hospitals"
 )
-census_chart = admitted_patients_chart(alt=alt, census=census_df, parameters=p)
+census_chart = admitted_patients_chart(alt=alt, census=m.census_df, parameters=p)
 st.altair_chart(
-    census_chart, use_container_width=True
+    admitted_patients_chart(alt=alt, census=m.census_df, parameters=p),
+    use_container_width=True,
 )
-st.markdown(chart_descriptions(census_chart, suffix=" Census"))
+st.markdown(chart_descriptions(census_chart, p.labels, suffix=" Census"))
 if st.checkbox("Show Projected Census in tabular form"):
     if st.checkbox("Show Daily Census Counts"):
-        draw_census_table(st, admissions_df, as_date=p.as_date, daily_count=True)
+        draw_census_table(st, m.census_df, p.labels, as_date=p.as_date, daily_count=True)
     else:
-        draw_census_table(st, census_df, as_date=p.as_date, daily_count=False)
+        draw_census_table(st, m.census_df, p.labels, as_date=p.as_date, daily_count=False)
     build_download_link(st,
         filename="projected_census.csv",
-        df=census_df,
+        df=m.census_df,
         parameters=p
     )
 
@@ -87,9 +86,9 @@
 )
 if st.checkbox("Show Additional Projections"):
     show_additional_projections(
-        st, alt, additional_projections_chart, parameters=p
+        st, alt, additional_projections_chart, model=m, parameters=p
     )
     if st.checkbox("Show Raw SIR Simulation Data"):
-        draw_raw_sir_simulation_table(st, parameters=p)
+        draw_raw_sir_simulation_table(st, model=m, parameters=p)
 write_definitions(st)
 write_footer(st)

src/cli.py

@@ -9,8 +9,8 @@
 from pandas import DataFrame
 
 from penn_chime.parameters import Parameters
+from penn_chime.models import SimSirModel
 from penn_chime.utils import RateLos
-from penn_chime.models import build_admissions_df, build_census_df
 
 class FromFile(Action):
     """From File."""
@@ -108,30 +108,22 @@ def main():
         doubling_time=a.doubling_time,
         known_infected=a.known_infected,
         market_share=a.market_share,
+        n_days=a.n_days,
         relative_contact_rate=a.relative_contact_rate,
         susceptible=a.susceptible,
-        n_days=a.n_days,
+
         hospitalized=RateLos(a.hospitalized_rate, a.hospitalized_los),
         icu=RateLos(a.icu_rate, a.icu_los),
         ventilated=RateLos(a.ventilated_rate, a.ventilated_los),
     )
 
-    raw_df = DataFrame(
-        {
-            "Susceptible": p.susceptible_v,
-            "Infected": p.infected_v,
-            "Recovered": p.recovered_v,
-        }
-    )
-
-    admits_df = build_admissions_df(p)
-    census_df = build_census_df(admits_df, p)
+    m = SimSirModel(p)
 
     prefix = a.prefix
     for df, name in (
-        (raw_df, "raw"),
-        (admits_df, "admits"),
-        (census_df, "census"),
+        (m.raw_df, "raw"),
+        (m.admits_df, "admits"),
+        (m.census_df, "census"),
     ):
         df.to_csv(prefix + name + ".csv")
 

src/penn_chime/charts.py

@@ -1,9 +1,9 @@
 
 from math import ceil
 import datetime
+
 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
@@ -31,9 +31,10 @@ def new_admissions_chart(
     else:
         x_kwargs = {"shorthand": "day", "title": "Days from today"}
 
+    # TODO fix the fold to allow any number of dispositions
     return (
         alt.Chart(projection_admits.head(plot_projection_days))
-        .transform_fold(fold=["Hospitalized", "ICU", "Ventilated"])
+        .transform_fold(fold=["hospitalized", "icu", "ventilated"])
         .mark_line(point=True)
         .encode(
             x=alt.X(**x_kwargs),
@@ -71,9 +72,10 @@ def admitted_patients_chart(
         y_scale.domain = (0, max_y_axis)
         y_scale.clamp = True
 
+    # TODO fix the fold to allow any number of dispositions
     return (
         alt.Chart(census.head(plot_projection_days))
-        .transform_fold(fold=["Hospitalized", "ICU", "Ventilated"])
+        .transform_fold(fold=["hospitalized", "icu", "ventilated"])
         .mark_line(point=True)
         .encode(
             x=alt.X(**x_kwargs),
@@ -90,11 +92,15 @@ def admitted_patients_chart(
 
 
 def additional_projections_chart(
-    alt, parameters: Parameters
+    alt, model, parameters
 ) -> Chart:
-    i = parameters.infected_v
-    r = parameters.recovered_v
-    dat = pd.DataFrame({"Infected": i, "Recovered": r})
+
+    # TODO use subselect of df_raw instead of creating a new df
+    raw_df = model.raw_df
+    dat = pd.DataFrame({
+        "infected": raw_df.infected,
+        "recovered": raw_df.recovered
+    })
     dat["day"] = dat.index
 
     as_date = parameters.as_date
@@ -114,7 +120,7 @@ def additional_projections_chart(
 
     return (
         alt.Chart(dat)
-        .transform_fold(fold=["Infected", "Recovered"])
+        .transform_fold(fold=["infected", "recovered"])
         .mark_line()
         .encode(
             x=alt.X(**x_kwargs),
@@ -126,7 +132,7 @@ def additional_projections_chart(
     )
 
 
-def chart_descriptions(chart: Chart, suffix: str = ""):
+def chart_descriptions(chart: Chart, labels, suffix: str = ""):
     """
 
     :param chart: Chart: The alt chart to be used in finding max points
@@ -137,7 +143,7 @@ def chart_descriptions(chart: Chart, suffix: str = ""):
     """
     messages = []
 
-    cols = ["Hospitalized", "ICU", "Ventilated"]
+    cols = ["hospitalized", "icu", "ventilated"]
     asterisk = False
     day = "date" if "date" in chart.data.columns else "day"
 
@@ -153,7 +159,7 @@ def chart_descriptions(chart: Chart, suffix: str = ""):
 
         messages.append(
             "{}{} peaks at {:,} on day {}{}".format(
-                col,
+                labels[col],
                 suffix,
                 ceil(chart.data[col].max()),
                 on,

src/penn_chime/defaults.py

@@ -25,25 +25,34 @@ def __init__(
         current_hospitalized: int,
         doubling_time: int,
         known_infected: int,
-        n_days: int,
         relative_contact_rate: int,
         region: Regions,
+
         hospitalized: RateLos,
         icu: RateLos,
         ventilated: RateLos,
+
+        as_date: bool = False,
         market_share: float = 1.0,
+        max_y_axis: int = None,
+        n_days: int = 60,
+        recovery_days: int = 14,
     ):
         self.region = region
-        self.known_infected = known_infected
         self.current_hospitalized = current_hospitalized
+        self.known_infected = known_infected
         self.doubling_time = doubling_time
-        self.market_share = market_share
         self.relative_contact_rate = relative_contact_rate
 
         self.hospitalized = hospitalized
         self.icu = icu
         self.ventilated = ventilated
+
+        self.as_date = as_date
+        self.market_share = market_share
+        self.max_y_axis = max_y_axis
         self.n_days = n_days
+        self.recovery_days = recovery_days
 
     def __repr__(self) -> str:
         return f"Constants(susceptible_default: {self.region.susceptible}, known_infected: {self.known_infected})"