Add sim_sir_df and eliminate comparison float casts

Author: jlubken

penn_chime/models.py

@@ -1,49 +1,93 @@
-from typing import Tuple
+from typing import Generator, Tuple
 
 import numpy as np
+import pandas as pd
 import streamlit as st
 
 
-# The SIR model, one time step
 @st.cache
-def sir(y, beta, gamma, N):
-    S, I, R = y
-    Sn = (-beta * S * I) + S
-    In = (beta * S * I - gamma * I) + I
-    Rn = gamma * I + R
-    if Sn < 0:
-        Sn = 0
-    if In < 0:
-        In = 0
-    if Rn < 0:
-        Rn = 0
-
-    scale = N / (Sn + In + Rn)
-    return Sn * scale, In * scale, Rn * scale
-
-
-# Run the SIR model forward in time
+def sir(
+    s: float, i: float, r: float,
+    beta: float, gamma: float, n: float
+) -> Tuple[float, float, float]:
+    """The SIR model, one time step."""
+    s_n = (-beta * s * i) + s
+    i_n = (beta * s * i - gamma * i) + i
+    r_n = gamma * i + r
+    if s_n < 0.0:
+        s_n = 0.0
+    if i_n < 0.0:
+        i_n = 0.0
+    if r_n < 0.0:
+        r_n = 0.0
+
+    scale = n / (s_n + i_n + r_n)
+    return s_n * scale, i_n * scale, r_n * scale
+
+
+def gen_sir(
+    s: float, i: float, r: float,
+    beta: float, gamma: float, n_days: int, beta_decay: float = 0.0
+) -> Generator:
+    """Simulate SIR model forward in time yielding tuples."""
+    s, i, r, beta_decay = (float(v) for v in (s, i, r, beta_decay))
+    n = s + i + r
+    for _ in range(n_days + 1):
+        yield s, i, r
+        s, i, r = sir(s, i, r, beta, gamma, n)
+        # okay even if beta_decay is 0.0
+        beta = beta * (1.0 - beta_decay)
+
+
 @st.cache
 def sim_sir(
-    S, I, R, beta, gamma, n_days, beta_decay=0
+    s: float, i: float, r: float,
+    beta: float, gamma: float, n_days: int, beta_decay: float = 0.0
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    N = S + I + R
-    s, i, r = [S], [I], [R]
+    """Simulate the SIR model forward in time."""
+    s, i, r, beta_decay = (float(v) for v in (s, i, r, beta_decay))
+    n = s + i + r
+    s_v, i_v, r_v = [s], [i], [r]
     for day in range(n_days):
-        y = S, I, R
-        S, I, R = sir(y, beta, gamma, N)
-        beta = beta * (1 - beta_decay)  # okay even if beta_decay is 0
-        s.append(S)
-        i.append(I)
-        r.append(R)
+        s, i, r = sir(s, i, r, beta, gamma, n)
+        # okay even if beta_decay is 0.0
+        beta = beta * (1.0 - beta_decay)
+        s_v.append(s)
+        i_v.append(i)
+        r_v.append(r)
+
+    return (
+        np.array(s_v),
+        np.array(i_v),
+        np.array(r_v),
+    )
+
+
+@st.cache
+def sim_sir_df(
+    s: float, i: float, r: float,
+    beta: float, gamma: float, n_days: int, beta_decay: float = 0.0
+) -> pd.DataFrame:
+    """Simulate the SIR model forward in time."""
+    return pd.DataFrame(
+        data=gen_sir(s, i, r, beta, gamma, n_days, beta_decay),
+        columns=("S", "I", "R"),
+    )
+
+
+@st.cache
+def get_hospitalizations2(
+    infected: np.ndarray, rates: Tuple[float, ...], market_share: float = 1.0
+) -> Tuple[np.ndarray, ...]:
+    """Get hopitalizations adjusted by rate and market_share."""
+    return (*(infected * rate * market_share for rate in rates),)
 
-    s, i, r = np.array(s), np.array(i), np.array(r)
-    return s, i, r
 
 @st.cache
 def get_hospitalizations(
     infected: np.ndarray, rates: Tuple[float, float, float], market_share: float
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Get hopitalizations adjusted by rate and market_share."""
     hosp_rate, icu_rate, vent_rate = rates
 
     hosp = infected * hosp_rate * market_share

test_app.py

@@ -5,7 +5,7 @@
                  hosp_rate, icu_rate, vent_rate, hosp_los, icu_los, vent_los, market_share, S, initial_infections,
                  detection_prob, hospitalization_rates, I, R, beta, gamma, n_days, beta_decay,
                  projection_admits, alt)
-from penn_chime.models import sir, sim_sir
+from penn_chime.models import sir, sim_sir, sim_sir_df
 from penn_chime.presentation import display_header, new_admissions_chart
 
 
@@ -70,12 +70,11 @@ def test_header_fail():
 
 # Test the math
 
-
 def test_sir():
     """
     Someone who is good at testing, help
     """
-    assert sir((100, 1, 0), 0.2, 0.5, 1) == (
+    assert sir(100, 1, 0, 0.2, 0.5, 1) == (
         0.7920792079207921,
         0.20297029702970298,
         0.0049504950495049506,
@@ -88,13 +87,28 @@ def test_sim_sir():
     """
     s, i, r = sim_sir(S, I, R, beta, gamma, n_days, beta_decay=beta_decay)
     assert round(s[0], 0) == 4119405
-    assert round(s[-1], 2) == 3421436.31
     assert round(i[0], 2) == 533.33
+    assert round(r[0], 0) == 0.0
+    assert round(s[-1], 2) == 3421436.31
     assert round(i[-1], 2) == 418157.62
-    assert round(r[0], 0) == 0
     assert round(r[-1], 2) == 280344.40
 
 
+def test_sim_sir_df():
+    """
+    Rounding to move fast past decimal place issues
+    """
+    df = sim_sir_df(S, I, R, beta, gamma, n_days, beta_decay=beta_decay)
+    first = df.iloc[0]
+    last = df.iloc[-1]
+    assert round(first[0], 0) == 4119405
+    assert round(first[1], 2) == 533.33
+    assert round(first[2], 0) == 0.0
+    assert round(last[0], 2) == 3421436.31
+    assert round(last[1], 2) == 418157.62
+    assert round(last[2], 2) == 280344.40
+
+
 def test_initial_conditions():
     """
     Note: For the rates (ie hosp_rate) - just change the value, leave the "100" alone.