Merge pull request #164 from nickcanz/set_max_y_axis

Add option to set a max y-axis value

Author: quinn-dougherty

app.py

@@ -73,22 +73,22 @@
 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), use_container_width=True
+    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
 )
 if st.checkbox("Show Projected Admissions in tabular form"):
     draw_projected_admissions_table(st, projection_admits, 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), use_container_width=True)
+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)
 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)
+    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)
     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)
 write_definitions(st)

penn_chime/models.py

@@ -21,6 +21,8 @@ def __init__(
         hospitalized: RateLos,
         icu: RateLos,
         ventilated: RateLos,
+
+        max_y_axis: int
     ):
         self.current_hospitalized = current_hospitalized
         self.doubling_time = doubling_time
@@ -34,6 +36,8 @@ def __init__(
         self.icu = icu
         self.ventilated = ventilated
 
+        self.max_y_axis = max_y_axis
+
         self.rates = tuple(
             each.rate
             for each in (hospitalized, icu, ventilated)

penn_chime/presentation.py

@@ -211,6 +211,16 @@ def display_sidebar(st, d: Constants) -> Parameters:
         format="%i",
     )
 
+    max_y_axis_set = st.sidebar.checkbox("Set the Y-axis on graphs to a static value")
+    max_y_axis = None
+    if max_y_axis_set:
+        max_y_axis = st.sidebar.number_input(
+            "Y-axis static value",
+            value=500,
+            format="%i",
+            step=25,
+        )
+
     return Parameters(
         current_hospitalized=current_hospitalized,
         doubling_time=doubling_time,
@@ -221,7 +231,8 @@ def display_sidebar(st, d: Constants) -> Parameters:
 
         hospitalized=RateLos(hospitalized_rate, hospitalized_los),
         icu=RateLos(icu_rate, icu_los),
-        ventilated=RateLos(ventilated_rate, ventilated_los)
+        ventilated=RateLos(ventilated_rate, ventilated_los),
+        max_y_axis=max_y_axis
     )
 
 
@@ -375,12 +386,23 @@ def write_footer(st):
 
 
 def new_admissions_chart(
-    alt, projection_admits: pd.DataFrame, plot_projection_days: int, as_date: bool = False
+    alt,
+    projection_admits: pd.DataFrame,
+    plot_projection_days: int,
+    as_date: bool = False,
+    max_y_axis: int = None
 ) -> alt.Chart:
     """docstring"""
     projection_admits = projection_admits.rename(
         columns={"hosp": "Hospitalized", "icu": "ICU", "vent": "Ventilated"}
     )
+
+    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)
@@ -394,7 +416,7 @@ def new_admissions_chart(
         .mark_line(point=True)
         .encode(
             x=alt.X(**x_kwargs),
-            y=alt.Y("value:Q", title="Daily admissions"),
+            y=alt.Y("value:Q", title="Daily admissions", scale=y_scale),
             color="key:N",
             tooltip=[
                 tooltip_dict[as_date],
@@ -410,7 +432,8 @@ def admitted_patients_chart(
     alt,
     census: pd.DataFrame,
     plot_projection_days: int,
-    as_date: bool = False
+    as_date: bool = False,
+    max_y_axis: int = None
 ) -> alt.Chart:
     """docstring"""
     census = census.rename(
@@ -427,13 +450,19 @@ def admitted_patients_chart(
     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(census.head(plot_projection_days))
         .transform_fold(fold=["Hospital Census", "ICU Census", "Ventilated Census"])
         .mark_line(point=True)
         .encode(
             x=alt.X(**x_kwargs),
-            y=alt.Y("value:Q", title="Census"),
+            y=alt.Y("value:Q", title="Census", scale=y_scale),
             color="key:N",
             tooltip=[
                 tooltip_dict[as_date],
@@ -449,7 +478,8 @@ def additional_projections_chart(
     alt,
     i: np.ndarray,
     r: np.ndarray,
-    as_date: bool = False
+    as_date: bool = False,
+    max_y_axis: int = None
 ) -> alt.Chart:
     dat = pd.DataFrame({"Infected": i, "Recovered": r})
     dat["day"] = dat.index
@@ -459,26 +489,40 @@ def additional_projections_chart(
     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"),
+            y=alt.Y("value:Q", title="Case Volume", scale=y_scale),
             tooltip=["key:N", "value:Q"],
             color="key:N",
         )
         .interactive()
     )
 
 
-def show_additional_projections(st, alt, charting_func, i, r, as_date: bool = False):
+def show_additional_projections(
+    st,
+    alt,
+    charting_func,
+    i,
+    r,
+    as_date: bool = False,
+    max_y_axis: int = None
+):
     st.subheader(
         "The number of infected and recovered individuals in the hospital catchment region at any given moment"
     )
 
-    st.altair_chart(charting_func(alt, i, r, as_date=as_date), use_container_width=True)
+    st.altair_chart(charting_func(alt, i, r, as_date=as_date, max_y_axis=max_y_axis), use_container_width=True)
 
 
 ##########