Explicitly output new infections, rather than calculating it post-hoc

Author: PhilMiller

src/penn_chime/models.py

@@ -100,7 +100,7 @@ def __init__(self, p: Parameters) -> SimSirModel:
 
 def sir(
     s: float, i: float, r: float, beta: float, gamma: float, n: float
-) -> Tuple[float, float, float]:
+) -> Tuple[float, float, float, float]:
     """The SIR model, one time step."""
 
     potential_new_infections = beta * s * i
@@ -119,18 +119,18 @@ def sir(
         r_n = 0.0
 
     scale = n / (s_n + i_n + r_n)
-    return s_n * scale, i_n * scale, r_n * scale
+    return s_n * scale, i_n * scale, r_n * scale, new_infections * scale
 
 
 def gen_sir(
     s: float, i: float, r: float, beta: float, gamma: float, n_days: int
-) -> Generator[Tuple[float, float, float], None, None]:
+) -> Generator[Tuple[float, float, float, float], None, None]:
     """Simulate SIR model forward in time yielding tuples."""
-    s, i, r = (float(v) for v in (s, i, r))
+    s, i, r, new_infections = (float(v) for v in (s, i, r, 0.0))
     n = s + i + r
     for d in range(n_days + 1):
-        yield d, s, i, r
-        s, i, r = sir(s, i, r, beta, gamma, n)
+        yield d, s, i, r, new_infections
+        s, i, r, new_infections = sir(s, i, r, beta, gamma, n)
 
 
 def sim_sir_df(
@@ -139,7 +139,7 @@ def sim_sir_df(
     """Simulate the SIR model forward in time."""
     return pd.DataFrame(
         data=gen_sir(s, i, r, beta, gamma, n_days),
-        columns=("day", "susceptible", "infected", "recovered"),
+        columns=("day", "susceptible", "infected", "recovered", "new_infections"),
     )
 
 def build_dispositions_df(

tests/test_app.py

@@ -137,14 +137,12 @@ def test_sir():
     Someone who is good at testing, help
     """
     sir_test = sir(100, 1, 0, 0.2, 0.5, 1)
-    assert sir_test == (
-        0.7920792079207921,
-        0.20297029702970298,
-        0.0049504950495049506,
-    ), "This contrived example should work"
-
-    assert isinstance(sir_test, tuple)
-    for v in sir_test:
+    s, i, r, i_n = sir_test
+    assert s == 0.7920792079207921, "This contrived example should work"
+    assert i == 0.20297029702970298, "This contrived example should work"
+    assert r == 0.0049504950495049506, "This contrived example should work"
+
+    for v in (sir_test):
         assert isinstance(v, float)
         assert v >= 0