Make an explicit expression for new infections in the model implementation

Author: PhilMiller

src/penn_chime/models.py

@@ -102,11 +102,17 @@ 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
+
+    potential_new_infections = beta * s * i
+    # Can't infect more people than there are left to become infected
+    new_infections = min(s, potential_new_infections)
+
+    s_n = s - new_infections
+    i_n = (new_infections - gamma * i) + i
     r_n = gamma * i + r
-    if s_n < 0.0:
-        s_n = 0.0
+
+    assert s_n >= 0.0
+
     if i_n < 0.0:
         i_n = 0.0
     if r_n < 0.0:

tests/test_app.py

@@ -163,7 +163,7 @@ def test_sir():
 
     with pytest.raises(TypeError) as error:
         sir(100, 1, 0, "beta", 0.5, 1)
-    assert str(error.value) == "bad operand type for unary -: 'str'"
+    assert str(error.value) == "'<' not supported between instances of 'str' and 'int'"
 
     with pytest.raises(TypeError) as error:
         sir(100, 1, 0, 0.2, "gamma", 1)