Repository: CDCgov/PyRenew

Issue:
Add infection initialization at stable age distribution given R

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov) Report
Attention: Patch coverage is `35.71429%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 95.69%. Comparing base [(`bbaf1b9`)](https://app.codecov.io/gh/CDCgov/PyRenew/commit/bbaf1b9b629d861ba534dd5b4c5693f84676bfe8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov) to head [(`6c2bd88`)](https://app.codecov.io/gh/CDCgov/PyRenew/commit/6c2bd88ffbb62b87d3134ccdf6b8471a81c01cad?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov).

| [Files with missing lines](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov) | Patch % | Lines |
|---|---|---|
| [pyrenew/math.py](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482?src=pr&el=tree&filepath=pyrenew%2Fmath.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov#diff-cHlyZW5ldy9tYXRoLnB5) | 35.71% | [9 Missing :warning: ](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #482      +/-   ##
==========================================
- Coverage   96.56%   95.69%   -0.88%     
==========================================
  Files          41       41              
  Lines         962      975      +13     
==========================================
+ Hits          929      933       +4     
- Misses         33       42       +9     
```

| [Flag](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov) | `95.69% <35.71%> (-0.88%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/CDCgov/PyRenew/pull/482?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=CDCgov).


Patch:
diff --git a/pyrenew/math.py b/pyrenew/math.py
index db94200..b8bfd74 100755
--- a/pyrenew/math.py
+++ b/pyrenew/math.py
@@ -12,6 +12,146 @@ from jax.typing import ArrayLike
 from pyrenew.distutil import validate_discrete_dist_vector
 
 
+def _positive_ints_like(vec: ArrayLike) -> jnp.ndarray:
+    """
+    Given an array of size n, return the 1D Jax array
+    ``[1, ... n]``.
+
+    Parameters
+    ----------
+    vec: ArrayLike
+        The template array
+
+    Returns
+    -------
+    jnp.ndarray
+        The resulting array ``[1, ..., n]``.
+    """
+    return jnp.arange(1, jnp.size(vec) + 1)
+
+
+def neg_MGF(r: float, w: ArrayLike) -> float:
+    """
+    Compute the negative moment generating function (MGF)
+    for a given rate ``r`` and weights ``w``.
+
+    Parameters
+    ----------
+    r: float
+       The rate parameter.
+
+    w: ArrayLike
+       An array of weights.
+
+    Returns
+    -------
+    float
+        The value of the negative MGF evaluated at ``r``
+        and ``w``.
+
+    Notes
+    -----
+    For a finite discrete random variable :math:`X` supported on
+    the first :math:`n` positive integers (:math:`\\{1, 2, ..., n \\}`),
+    the moment generating function (MGF) :math:`M_+(r)` is defined
+    as the expected value of :math:`\\exp(rX)`. Similarly, the negative
+    moment generating function :math:`M_-(r)` is the expected value of
+    :math:`\\exp(-rX)`. So if we represent the PMF of :math:`X` as a
+    "weights" vector :math:`w` of length :math:`n`, the negative MGF
+    :math:`M_-(r)` is given by:
+
+    .. math::
+        M_-(r) = \\sum_{t = 1}^{n} w_i \\exp(-rt)
+    """
+    return jnp.sum(w * jnp.exp(-r * _positive_ints_like(w)))
+
+
+def neg_MGF_del_r(r: float, w: ArrayLike) -> float:
+    """
+    Compute the value of the partial deriative of
+    :func:`neg_MGF` with respect to ``r``
+    evaluated at a particular ``r`` and ``w`` pair.
+
+    Parameters
+    ----------
+    r: float
+       The rate parameter.
+
+    w: ArrayLike
+       An array of weights.
+
+    Returns
+    -------
+    float
+        The value of the partial derivative evaluated at ``r``
+        and ``w``.
+    """
+    t_vec = _positive_ints_like(w)
+    return -jnp.sum(w * t_vec * jnp.exp(-r * t_vec))
+
+
+def r_approx_from_R(R: float, g: ArrayLike, n_newton_steps: int) -> ArrayLike:
+    """
+    Get the approximate asymptotic geometric growth rate ``r``
+    for a renewal process with a fixed reproduction number ``R``
+    and discrete generation interval PMF ``g``.
+
+    Uses Newton's method with a fixed number of steps.
+
+    Parameters
+    ----------
+    R: float
+        The reproduction number
+
+    g: ArrayLike
+        The probability mass function of the generation
+        interval.
+
+    n_newton_steps: int
+        Number of steps to take when performing Newton's method.
+
+    Returns
+    -------
+    float
+        The approximate value of ``r``.
+
+    Notes
+    -----
+    For a fixed value of :math:`\\mathcal{R}`, a renewal process
+    has an asymptotic geometric growth rate :math:`r` that satisfies
+
+    .. math::
+        M_{-}(r) - \\frac{1}{\\mathcal{R}} = 0
+
+    where :math:`M_-(r)` is the negative moment generating function
+    for a random variable :math:`\\tau` representing the (discrete)
+    generation interval. See :func:`neg_MGF` for details.
+
+    We obtain a value for :math:`r` via approximate numerical solution
+    of this implicit equation.
+
+    We first make an initial guess based on the mean generation interval
+    :math:`\\bar{\\tau} = \\mathbb{E}(\\tau)`:
+
+    .. math::
+        r \\approx \\frac{\\mathcal{R} - 1}{\\mathcal{R} \\bar{\\tau}}
+
+    We then refine this approximation by applying Newton's method for
+    a fixed number of steps.
+    """
+    mean_gi = jnp.dot(g, _positive_ints_like(g))
+    init_r = (R - 1) / (R * mean_gi)
+
+    def _r_next(r, _) -> tuple[ArrayLike, None]:  # numpydoc ignore=GL08
+        return (
+            r - ((R * neg_MGF(r, g) - 1) / (R * neg_MGF_del_r(r, g))),
+            None,
+        )
+
+    result, _ = scan(f=_r_next, init=init_r, xs=None, length=n_newton_steps)
+    return result
+
+
 def get_leslie_matrix(
     R: float, generation_interval_pmf: ArrayLike
 ) -> ArrayLike:


Test patch:
diff --git a/test/test_leslie_matrix.py b/test/test_leslie_matrix.py
deleted file mode 100755
index 3610685..0000000
--- a/test/test_leslie_matrix.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-from numpy.testing import assert_array_almost_equal
-
-import pyrenew.math as pmath
-
-
-def test_get_leslie():
-    """
-    Test that get_leslie matrix
-    returns expected Leslie matrices
-    """
-
-    gi = jnp.array([0.4, 0.2, 0.2, 0.1, 0.1])
-    R_a = 0.4
-    R_b = 3.0
-    expected_a = jnp.array(
-        [
-            [0.16, 0.08, 0.08, 0.04, 0.04],
-            [1, 0, 0, 0, 0],
-            [0, 1, 0, 0, 0],
-            [0, 0, 1, 0, 0],
-            [0, 0, 0, 1, 0],
-        ]
-    )
-    expected_b = jnp.array(
-        [
-            [1.2, 0.6, 0.6, 0.3, 0.3],
-            [1, 0, 0, 0, 0],
-            [0, 1, 0, 0, 0],
-            [0, 0, 1, 0, 0],
-            [0, 0, 0, 1, 0],
-        ]
-    )
-
-    assert_array_almost_equal(pmath.get_leslie_matrix(R_a, gi), expected_a)
-    assert_array_almost_equal(pmath.get_leslie_matrix(R_b, gi), expected_b)
diff --git a/test/test_math.py b/test/test_math.py
new file mode 100644
index 0000000..bace5d4
--- /dev/null
+++ b/test/test_math.py
@@ -0,0 +1,142 @@
+"""
+Unit tests for the pyrenew.math module.
+"""
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from numpy.random import RandomState
+from numpy.testing import (
+    assert_almost_equal,
+    assert_array_almost_equal,
+    assert_array_equal,
+)
+
+import pyrenew.math as pmath
+
+rng = RandomState(5)
+
+
+@pytest.mark.parametrize(
+    "arr, arr_len",
+    [
+        ([3, 1, 2], 3),
+        (np.ones(50), 50),
+        ((jnp.nan * jnp.ones(250)).reshape((50, -1)), 250),
+    ],
+)
+def test_positive_ints_like(arr, arr_len):
+    """
+    Test the _positive_ints_like helper function.
+    """
+    result = pmath._positive_ints_like(arr)
+    expected = jnp.arange(1, arr_len + 1)
+    assert_array_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "R, G",
+    [
+        (5, rng.dirichlet(np.ones(2))),
+        (0.2, rng.dirichlet(np.ones(50))),
+        (1, rng.dirichlet(np.ones(10))),
+        (1.01, rng.dirichlet(np.ones(4))),
+        (0.99, rng.dirichlet(np.ones(6))),
+    ],
+)
+def test_r_approx(R, G):
+    """
+    Test that r_approx_from_R gives answers
+    consistent with those gained from a Leslie
+    matrix approach.
+    """
+    r_val = pmath.r_approx_from_R(R, G, n_newton_steps=5)
+    e_val, stable_dist = pmath.get_asymptotic_growth_rate_and_age_dist(R, G)
+
+    unnormed = r_val * stable_dist
+    if r_val != 0:
+        assert_array_almost_equal(unnormed / np.sum(unnormed), stable_dist)
+    else:
+        assert_almost_equal(e_val, 1, decimal=5)
+
+
+def test_asymptotic_properties():
+    """
+    Check that the calculated
+    asymptotic growth rate and
+    age distribution given by
+    get_asymptotic_growth_rate()
+    and get_stable_age_distribution()
+    agree with simulated ones from
+    just running a process for a
+    while.
+    """
+    R = 1.2
+    gi = np.array([0.2, 0.1, 0.2, 0.15, 0.05, 0.025, 0.025, 0.25])
+    A = pmath.get_leslie_matrix(R, gi)
+
+    # check via Leslie matrix multiplication
+    x = np.array([1, 0, 0, 0, 0, 0, 0, 0])
+    for i in range(1000):
+        x_new = A @ x
+        rat_x = np.sum(x_new) / np.sum(x)
+        x = x_new
+
+    assert_almost_equal(
+        rat_x, pmath.get_asymptotic_growth_rate(R, gi), decimal=5
+    )
+    assert_array_almost_equal(
+        x / np.sum(x), pmath.get_stable_age_distribution(R, gi)
+    )
+
+    # check via backward-looking convolution
+    y = np.array([1, 0, 0, 0, 0, 0, 0, 0])
+    for j in range(1000):
+        new_pop = np.dot(y, R * gi)
+        rat_y = new_pop / y[0]
+        y = np.hstack([new_pop, y[:-1]])
+    assert_almost_equal(
+        rat_y, pmath.get_asymptotic_growth_rate(R, gi), decimal=5
+    )
+    assert_array_almost_equal(
+        y / np.sum(x), pmath.get_stable_age_distribution(R, gi)
+    )
+
+
+@pytest.mark.parametrize(
+    "R, gi, expected",
+    [
+        (
+            0.4,
+            np.array([0.4, 0.2, 0.2, 0.1, 0.1]),
+            np.array(
+                [
+                    [0.16, 0.08, 0.08, 0.04, 0.04],
+                    [1, 0, 0, 0, 0],
+                    [0, 1, 0, 0, 0],
+                    [0, 0, 1, 0, 0],
+                    [0, 0, 0, 1, 0],
+                ]
+            ),
+        ),
+        (
+            3,
+            np.array([0.4, 0.2, 0.2, 0.1, 0.1]),
+            np.array(
+                [
+                    [1.2, 0.6, 0.6, 0.3, 0.3],
+                    [1, 0, 0, 0, 0],
+                    [0, 1, 0, 0, 0],
+                    [0, 0, 1, 0, 0],
+                    [0, 0, 0, 1, 0],
+                ]
+            ),
+        ),
+    ],
+)
+def test_get_leslie(R, gi, expected):
+    """
+    Test that get_leslie matrix
+    returns expected Leslie matrices
+    """
+    assert_array_almost_equal(pmath.get_leslie_matrix(R, gi), expected)
diff --git a/test/test_process_asymptotics.py b/test/test_process_asymptotics.py
deleted file mode 100755
index 25079c9..0000000
--- a/test/test_process_asymptotics.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# numpydoc ignore=GL08
-
-import jax.numpy as jnp
-from numpy.testing import assert_almost_equal, assert_array_almost_equal
-
-import pyrenew.math as pmath
-
-
-def test_asymptotic_properties():
-    """
-    Check that the calculated
-    asymptotic growth rate and
-    age distribution given by
-    get_asymptotic_growth_rate()
-    and get_stable_age_distribution()
-    agree with simulated ones from
-    just running a process for a
-    while.
-    """
-    R = 1.2
-    gi = jnp.array([0.2, 0.1, 0.2, 0.15, 0.05, 0.025, 0.025, 0.25])
-    A = pmath.get_leslie_matrix(R, gi)
-
-    # check via Leslie matrix multiplication
-    x = jnp.array([1, 0, 0, 0, 0, 0, 0, 0])
-    for i in range(1000):
-        x_new = A @ x
-        rat_x = jnp.sum(x_new) / jnp.sum(x)
-        x = x_new
-
-    assert_almost_equal(
-        rat_x, pmath.get_asymptotic_growth_rate(R, gi), decimal=5
-    )
-    assert_array_almost_equal(
-        x / jnp.sum(x), pmath.get_stable_age_distribution(R, gi)
-    )
-
-    # check via backward-looking convolution
-    y = jnp.array([1, 0, 0, 0, 0, 0, 0, 0])
-    for j in range(1000):
-        new_pop = jnp.dot(y, R * gi)
-        rat_y = new_pop / y[0]
-        y = jnp.hstack([new_pop, y[:-1]])
-    assert_almost_equal(
-        rat_y, pmath.get_asymptotic_growth_rate(R, gi), decimal=5
-    )
-    assert_array_almost_equal(
-        y / jnp.sum(x), pmath.get_stable_age_distribution(R, gi)
-    )