Math

Helper functions for doing analytical and/or numerical calculations.

get_asymptotic_growth_rate

get_asymptotic_growth_rate(
    R: float, generation_interval_pmf: ArrayLike
) -> float

Get the asymptotic per timestep growth rate for a renewal process with a given value of \(\mathcal{R}\) and a given discrete generation interval probability mass vector.

This function computes that growth rate finding the dominant eigenvalue of the renewal process's Leslie matrix.

Parameters:

Name	Type	Description	Default
R	float	The reproduction number of the renewal process	required
generation_interval_pmf	ArrayLike	The discrete generation interval probability mass vector of the renewal process	required

Returns:

Type	Description
float	The asymptotic growth rate of the renewal process, as a jax float.

Source code in pyrenew/math.py

def get_asymptotic_growth_rate(R: float, generation_interval_pmf: ArrayLike) -> float:
    """
    Get the asymptotic per timestep growth rate
    for a renewal process with a given value of
    $\\mathcal{R}$ and a given discrete
    generation interval probability mass vector.

    This function computes that growth rate
    finding the dominant eigenvalue of the
    renewal process's Leslie matrix.

    Parameters
    ----------
    R
        The reproduction number of the renewal process

    generation_interval_pmf
        The discrete generation interval probability
        mass vector of the renewal process

    Returns
    -------
    float
        The asymptotic growth rate of the renewal process,
        as a jax float.
    """
    return get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf)[0]

get_asymptotic_growth_rate_and_age_dist

get_asymptotic_growth_rate_and_age_dist(
    R: float, generation_interval_pmf: ArrayLike
) -> tuple[float, ArrayLike]

Get the asymptotic per-timestep growth rate of the renewal process (the dominant eigenvalue of its Leslie matrix) and the associated stable age distribution (a normalized eigenvector associated to that eigenvalue).

Parameters:

Name	Type	Description	Default
R	float	The reproduction number of the renewal process	required
generation_interval_pmf	ArrayLike	The discrete generation interval probability mass vector of the renewal process	required

Returns:

Type	Description
tuple[float, ArrayLike]	A tuple consisting of the asymptotic growth rate of the process, as jax float, and the stable age distribution of the process, as a jax array probability vector of the same shape as the generation interval probability vector.

Raises:

Type	Description
ValueError	If an age distribution vector with non-zero imaginary part is produced.

Source code in pyrenew/math.py

def get_asymptotic_growth_rate_and_age_dist(
    R: float, generation_interval_pmf: ArrayLike
) -> tuple[float, ArrayLike]:
    """
    Get the asymptotic per-timestep growth
    rate of the renewal process (the dominant
    eigenvalue of its Leslie matrix) and the
    associated stable age distribution
    (a normalized eigenvector associated to
    that eigenvalue).

    Parameters
    ----------
    R
        The reproduction number of the renewal process

    generation_interval_pmf
        The discrete generation interval probability
        mass vector of the renewal process

    Returns
    -------
    tuple[float, ArrayLike]
        A tuple consisting of the asymptotic growth rate of
        the process, as jax float, and the stable age distribution
        of the process, as a jax array probability vector of the
        same shape as the generation interval probability vector.

    Raises
    ------
    ValueError
        If an age distribution vector with non-zero
        imaginary part is produced.
    """
    L = get_leslie_matrix(R, generation_interval_pmf)
    eigenvals, eigenvecs = jnp.linalg.eig(L)
    d = jnp.argmax(jnp.abs(eigenvals))  # index of dominant eigenvalue
    d_vec, d_val = eigenvecs[:, d], eigenvals[d]
    d_vec_real, d_val_real = jnp.real(d_vec), jnp.real(d_val)
    if not all(d_vec_real == d_vec):
        raise ValueError(
            "get_asymptotic_growth_rate_and_age_dist() "
            "produced an age distribution vector with "
            "non-zero imaginary part. "
            "Check your generation interval distribution "
            "vector and R value"
        )
    if not d_val_real == d_val:
        raise ValueError(
            "get_asymptotic_growth_rate_and_age_dist() "
            "produced an asymptotic growth rate with "
            "non-zero imaginary part. "
            "Check your generation interval distribution "
            "vector and R value"
        )
    d_vec_norm = d_vec_real / jnp.sum(d_vec_real)
    return d_val_real, d_vec_norm

get_leslie_matrix

get_leslie_matrix(R: float, generation_interval_pmf: ArrayLike) -> ArrayLike

Create the Leslie matrix corresponding to a basic renewal process with the given \(\mathcal{R}\) value and discrete generation interval pmf vector.

Parameters:

Name	Type	Description	Default
R	float	The reproduction number of the renewal process	required
generation_interval_pmf	ArrayLike	The discrete generation interval probability mass vector of the renewal process	required

Returns:

Type	Description
ArrayLike	The Leslie matrix for the renewal process, as a jax array.

Source code in pyrenew/math.py

def get_leslie_matrix(R: float, generation_interval_pmf: ArrayLike) -> ArrayLike:
    """
    Create the Leslie matrix
    corresponding to a basic
    renewal process with the
    given $\\mathcal{R}$
    value and discrete
    generation interval pmf
    vector.

    Parameters
    ----------
    R
        The reproduction number of the renewal process

    generation_interval_pmf
        The discrete generation interval probability
        mass vector of the renewal process

    Returns
    -------
    ArrayLike
        The Leslie matrix for the
        renewal process, as a jax array.
    """
    validate_discrete_dist_vector(generation_interval_pmf)
    gen_int_len = generation_interval_pmf.size
    aging_matrix = jnp.hstack(
        [
            jnp.identity(gen_int_len - 1),
            jnp.zeros(gen_int_len - 1)[..., jnp.newaxis],
        ]
    )

    return jnp.vstack([R * generation_interval_pmf, aging_matrix])

get_stable_age_distribution

get_stable_age_distribution(
    R: float, generation_interval_pmf: ArrayLike
) -> ArrayLike

Get the stable age distribution for a renewal process with a given value of R and a given discrete generation interval probability mass vector.

This function computes that stable age distribution by finding and then normalizing an eigenvector associated to the dominant eigenvalue of the renewal process's Leslie matrix.

Parameters:

Name	Type	Description	Default
R	float	The reproduction number of the renewal process	required
generation_interval_pmf	ArrayLike	The discrete generation interval probability mass vector of the renewal process	required

Returns:

Type	Description
ArrayLike	The stable age distribution for the process, as a jax array probability vector of the same shape as the generation interval probability vector.

Source code in pyrenew/math.py

def get_stable_age_distribution(
    R: float, generation_interval_pmf: ArrayLike
) -> ArrayLike:
    """
    Get the stable age distribution for a
    renewal process with a given value of
    R and a given discrete generation
    interval probability mass vector.

    This function computes that stable age
    distribution by finding and then normalizing
    an eigenvector associated to the dominant
    eigenvalue of the renewal process's
    Leslie matrix.

    Parameters
    ----------
    R
        The reproduction number of the renewal process

    generation_interval_pmf
        The discrete generation interval probability
        mass vector of the renewal process

    Returns
    -------
    ArrayLike
        The stable age distribution for the
        process, as a jax array probability vector of
        the same shape as the generation interval
        probability vector.
    """
    return get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf)[1]

integrate_discrete

integrate_discrete(
    init_diff_vals: ArrayLike, highest_order_diff_vals: ArrayLike
) -> ArrayLike

Integrate (de-difference) the differenced process, obtaining the process values \(X(t=0), X(t=1), ..., X(t)\) from the \(n^{th}\) differences and a set of initial process / difference values \(X(t=0), X^1(t=1), X^2(t=2), ..., X^{(n-1)}(t=n-1)\), where \(X^k(t)\) is the value of the \(n^{th}\) difference at index \(t\) of the process, obtaining a sequence of length equal to the length of the provided highest_order_diff_vals vector plus the order of the process.

Parameters:

Name	Type	Description	Default
init_diff_vals	ArrayLike	Values of \(X(t=0), X^1(t=1), X^2(t=2), ..., X^{(n-1)}(t=n-1)\).	required
highest_order_diff_vals	ArrayLike	Array of differences at the highest order of differencing, i.e. the order of the overall process, starting with \(X^{n}(t=n)\)	required

Returns:

Type	Description
ArrayLike	The integrated (de-differenced) sequence of values, of length n_diffs + order, where n_diffs is the number of highest_order_diff_vals and order is the order of the process.

Source code in pyrenew/math.py

def integrate_discrete(
    init_diff_vals: ArrayLike, highest_order_diff_vals: ArrayLike
) -> ArrayLike:
    """
    Integrate (de-difference) the differenced process,
    obtaining the process values $X(t=0), X(t=1), ..., X(t)$
    from the $n^{th}$ differences and a set of
    initial process / difference values
    $X(t=0), X^1(t=1), X^2(t=2), ..., X^{(n-1)}(t=n-1)$,
    where $X^k(t)$ is the value of the $n^{th}$
    difference at index $t$ of the process,
    obtaining a sequence of length equal to the length of
    the provided `highest_order_diff_vals` vector plus
    the order of the process.

    Parameters
    ----------
    init_diff_vals
        Values of
        $X(t=0), X^1(t=1), X^2(t=2), ..., X^{(n-1)}(t=n-1)$.

    highest_order_diff_vals
        Array of differences at the highest order of
        differencing, i.e. the order of the overall process,
        starting with $X^{n}(t=n)$

    Returns
    -------
    ArrayLike
        The integrated (de-differenced) sequence of values,
        of length n_diffs + order, where n_diffs is the
        number of highest_order_diff_vals and order is the
        order of the process.
    """
    inits_by_order = jnp.atleast_1d(init_diff_vals)
    highest_diffs = jnp.atleast_1d(highest_order_diff_vals)
    order = inits_by_order.shape[0]
    n_diffs = highest_diffs.shape[0]

    try:
        batch_shape = broadcast_shapes(
            highest_diffs.shape[1:], inits_by_order.shape[1:]
        )
    except Exception as e:
        raise ValueError(
            "Non-time dimensions "
            "(i.e. dimensions after the first) "
            "for highest_order_diff_vals and init_diff_vals "
            "must be broadcastable together. "
            "Got highest_order_diff_vals of shape "
            f"{highest_diffs.shape} and "
            "init_diff_vals of shape "
            f"{inits_by_order.shape}"
        ) from e

    highest_diffs = jnp.broadcast_to(highest_diffs, (n_diffs,) + batch_shape)
    inits_by_order = jnp.broadcast_to(inits_by_order, (order,) + batch_shape)

    highest_diffs = jnp.concatenate(
        [jnp.zeros_like(inits_by_order), highest_diffs],
        axis=0,
    )

    scan_arrays = (
        jnp.arange(start=order - 1, stop=-1, step=-1),
        jnp.flip(inits_by_order, axis=0),
    )

    integrated, _ = scan(f=_integrate_one_step, init=highest_diffs, xs=scan_arrays)

    return integrated

neg_MGF

neg_MGF(r: float, w: ArrayLike) -> float

Compute the negative moment generating function (MGF) for a given rate r and weights w.

Parameters:

Name	Type	Description	Default
r	float	The rate parameter.	required
w	ArrayLike	An array of weights.	required

Returns:

Type	Description
float	The value of the negative MGF evaluated at r and w.

Notes

For a finite discrete random variable \(X\) supported on the first \(n\) positive integers (\(\{1, 2, ..., n \}\)), the moment generating function (MGF) \(M_+(r)\) is defined as the expected value of \(\exp(rX)\). Similarly, the negative moment generating function \(M_-(r)\) is the expected value of \(\exp(-rX)\). So if we represent the PMF of \(X\) as a "weights" vector \(w\) of length \(n\), the negative MGF \(M_-(r)\) is given by:

\[ M_-(r) = \sum_{t = 1}^{n} w_i \exp(-rt) \]

Source code in pyrenew/math.py

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
        The rate parameter.

    w
        An array of weights.

    Returns
    -------
    float
        The value of the negative MGF evaluated at `r`
        and `w`.

    Notes
    -----
    For a finite discrete random variable $X$ supported on
    the first $n$ positive integers ($\\{1, 2, ..., n \\}$),
    the moment generating function (MGF) $M_+(r)$ is defined
    as the expected value of $\\exp(rX)$. Similarly, the negative
    moment generating function $M_-(r)$ is the expected value of
    $\\exp(-rX)$. So if we represent the PMF of $X$ as a
    "weights" vector $w$ of length $n$, the negative MGF
    $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)))

neg_MGF_del_r

neg_MGF_del_r(r: float, w: ArrayLike) -> float

Compute the value of the partial deriative of pyrenew.math.neg_MGF with respect to r evaluated at a particular r and w pair.

Parameters:

Name	Type	Description	Default
r	float	The rate parameter.	required
w	ArrayLike	An array of weights.	required

Returns:

Type	Description
float	The value of the partial derivative evaluated at r and w.

Source code in pyrenew/math.py

def neg_MGF_del_r(r: float, w: ArrayLike) -> float:
    """
    Compute the value of the partial deriative of
    [`pyrenew.math.neg_MGF`][] with respect to `r`
    evaluated at a particular `r` and `w` pair.

    Parameters
    ----------
    r
        The rate parameter.

    w
        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))

r_approx_from_R

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:

Name	Type	Description	Default
R	float	The reproduction number	required
g	ArrayLike	The probability mass function of the generation interval.	required
n_newton_steps	int	Number of steps to take when performing Newton's method.	required

Returns:

Type	Description
float	The approximate value of r.

Notes

For a fixed value of \(\mathcal{R}\), a renewal process has an asymptotic geometric growth rate \(r\) that satisfies

\[ M_{-}(r) - \frac{1}{\mathcal{R}} = 0 \]

where \(M_-(r)\) is the negative moment generating function for a random variable \(\tau\) representing the (discrete) generation interval. See pyrenew.math.neg_MGF for details.

We obtain a value for \(r\) via approximate numerical solution of this implicit equation.

We first make an initial guess based on the mean generation interval \(\bar{\tau} = \mathbb{E}(\tau)\):

\[ 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.

Source code in pyrenew/math.py

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
        The reproduction number

    g
        The probability mass function of the generation
        interval.

    n_newton_steps
        Number of steps to take when performing Newton's method.

    Returns
    -------
    float
        The approximate value of `r`.

    Notes
    -----
    For a fixed value of $\\mathcal{R}$, a renewal process
    has an asymptotic geometric growth rate $r$ that satisfies

    ```math
    M_{-}(r) - \\frac{1}{\\mathcal{R}} = 0
    ```

    where $M_-(r)$ is the negative moment generating function
    for a random variable $\\tau$ representing the (discrete)
    generation interval. See [`pyrenew.math.neg_MGF`][] for details.

    We obtain a value for $r$ via approximate numerical solution
    of this implicit equation.

    We first make an initial guess based on the mean generation interval
    $\\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