use session state

ringvax/app.py

@@ -3,6 +3,7 @@
 
 import altair as alt
 import graphviz
+import numpy as np
 import polars as pl
 import streamlit as st
 
@@ -72,6 +73,11 @@ def format_duration(x: float, digits=3) -> str:
         return f"{round(x, digits)} seconds"
 
 
+def set_session_default(key, value) -> None:
+    if key not in st.session_state:
+        st.session_state[key] = value
+
+
 def app():
     st.title("Ring vaccination")
 
@@ -86,11 +92,14 @@ def app():
             format="%.1f days",
         )
 
+        default_infectious_duration = 3.0
         max_infectious_duration = 10.0
         min_infectious_duration = 0.1
-        default_infectious_duration = 3.0
-        max_r0 = 10.0
-        default_r0 = 1.5
+
+        default_R0 = 1.5
+        max_R0 = 10.0
+        set_session_default("R0", default_R0)
+        set_session_default("infection_rate", default_R0 / default_infectious_duration)
 
         infectious_duration = st.slider(
             "Infectious duration",
@@ -102,48 +111,52 @@ def app():
         )
 
         st.subheader("Infectiousness")
-        r0_control = st.segmented_control(
+        infectiousness_control = st.segmented_control(
             "Variable infectiousness parameter",
-            options=["R0", "infection rate"],
+            options=["R0", "rate"],
             selection_mode="single",
             default="R0",
             help="R0 = infectious duration * infection rate. Only two of the "
             "three parameters can be varied.",
         )
 
-        if r0_control == "R0":
-            r0 = st.slider(
-                "R0", min_value=0.0, max_value=max_r0, value=default_r0, step=0.1
+        def callback_R0():
+            st.session_state["infection_rate"] = (
+                st.session_state["R0"] / infectious_duration
             )
 
-            infection_rate = st.slider(
-                "Infection rate (mean infections per day)",
-                min_value=0.0,
-                max_value=max_r0 / min_infectious_duration,
-                value=r0 / infectious_duration,
-                step=0.1,
-                disabled=True,
+        def callback_rate():
+            st.session_state["R0"] = (
+                st.session_state["infection_rate"] * infectious_duration
             )
 
-        elif r0_control == "infection rate":
-            infection_rate = st.slider(
-                "Infection rate (mean infections per day)",
-                min_value=0.0,
-                max_value=max_r0 / max_infectious_duration,
-                value=default_r0 / default_infectious_duration,
-                step=0.1,
-            )
+        R0 = st.slider(
+            "R0",
+            key="R0",
+            min_value=0.0,
+            max_value=max_R0,
+            step=0.1,
+            disabled=infectiousness_control != "R0",
+            on_change=callback_R0,
+        )
 
-            r0 = st.slider(
-                "R0",
-                min_value=0.0,
-                max_value=max_r0,
-                value=infectious_duration * infection_rate,
-                step=0.1,
-                disabled=True,
-            )
-        else:
-            raise RuntimeError(f"Unexpected control value: {r0_control=}")
+        infection_rate = st.slider(
+            "Infection rate (mean infections per day)",
+            key="infection_rate",
+            min_value=0.0,
+            max_value=max_R0 / min_infectious_duration,
+            step=0.1,
+            disabled=infectiousness_control != "rate",
+            on_change=callback_rate,
+        )
+
+        # double check that R0, rate, and duration are not in conflict
+        assert np.isclose(R0, infection_rate * infectious_duration)
+
+        # rather than trying to dynamically adjust the ranges on the sliders,
+        # issue a warning if the selected rate leads to an out-of-range R0
+        if R0 > max_R0:
+            st.warning(f"Selected infectious rate yields R0 > {max_R0}")
 
         st.subheader("Detection")
         p_passive_detect = (
@@ -228,7 +241,7 @@ def app():
         toc = time.perf_counter()
 
     st.write(
-        f"Ran {nsim} simulations in {format_duration(toc - tic)} with an $R_0$ of {infectious_duration * infection_rate:.2f} (the product of the average duration of infection and the infectious rate)."
+        f"Ran {nsim} simulations in {format_duration(toc - tic)} with an $R_0$ of {R0:.2f} (the product of the average duration of infection and the infectious rate)."
     )
 
     tab1, tab2 = st.tabs(["Simulation summary", "Per-simulation results"])