diff --git a/src/reedfrost/__init__.py b/src/reedfrost/__init__.py index fcd80d1..a5d52b1 100755 --- a/src/reedfrost/__init__.py +++ b/src/reedfrost/__init__.py @@ -3,12 +3,13 @@ import numpy as np import scipy.stats +from numpy.typing import NDArray from scipy.optimize import brentq from scipy.special import factorial @functools.cache -def _gontcharoff(k: int, q: float, m: int) -> float: +def _gontcharoff1(k: int, q: float, m: int) -> float: """ Gontcharoff polynomials, specific to the Lefevre & Picard formulations of Reed-Frost final outbreak size pmf calculations @@ -28,26 +29,51 @@ def _gontcharoff(k: int, q: float, m: int) -> float: else: return 1.0 / math.factorial(k) - sum( [ - q ** ((m + i) * (k - i)) / factorial(k - i) * _gontcharoff(i, q, m) + q ** ((m + i) * (k - i)) / factorial(k - i) * _gontcharoff1(i, q, m) for i in range(0, k) ] ) -def pmf(k: int, n: int, p: float, m: int = 1) -> float: +def _gontcharoff( + k: int | NDArray[np.int64], q: float, m: int +) -> float | NDArray[np.float64]: + """ + Gontcharoff polynomials, specific to the Lefevre & Picard + formulations of Reed-Frost final outbreak size pmf calculations + + See Lefevre & Picard 1990 (doi:10.2307/1427595) equation 2.1 + + Args: + k (int, or int array): degree + q (float): 1 - probability of "effective" contact + m (int): number of initial infected + + Returns: + float, or float array: value of the polynomial + """ + if isinstance(k, int): + return _gontcharoff1(k=k, q=q, m=m) + else: + return np.array([_gontcharoff1(k=kk, q=q, m=m) for kk in k]) + + +def pmf( + k: int | NDArray[np.int64], n: int, p: float, m: int = 1 +) -> float | NDArray[np.float64]: """ Probability mass function for final size of a Reed-Frost outbreak See Lefevre & Picard 1990 (doi:10.2307/1427595) equation 3.10 Args: - k (int): number of total infections + k (int, or int array): number of total infections n (int): initial number susceptible m (int): initial number infected p (float): probability of "effective contact" (i.e., infection) Returns: - float: pmf of the total infection distribution + float, or float array: pmf of the total infection distribution """ q = 1.0 - p return ( @@ -82,21 +108,21 @@ def f(t): return result -def dist_large(n: int, lambda_: float, i_n: int = 1): +def pmf_large( + k: int | NDArray[np.int64], n: int, lambda_: float, i_n: int = 1 +) -> float | NDArray[np.float64]: """Distribution of outbreak sizes, given a large outbreak See Barbour & Sergey 2004 (doi:10.1016/j.spa.2004.03.013) corollary 3.4 Args: + k (int, or int array): number of total infections n (int): initial number of susceptibles lambda_ (float): reproduction number i_n (int, optional): initial number of susceptibles. Defaults to 1. Returns: - scipy.stats.norm: RV object - - Examples: - dist_large(100, 1.5, 1).pdf(np.linspace(0, 100)) + float, or float array: pmf of the total infection distribution """ if not lambda_ > 1.0: raise RuntimeWarning( @@ -106,4 +132,4 @@ def dist_large(n: int, lambda_: float, i_n: int = 1): sigma = np.sqrt(theta * (1.0 - theta) / (1 - lambda_ * theta) ** 2) sd = np.sqrt(n) * sigma mean = n * (1.0 - theta) - return scipy.stats.norm(loc=mean, scale=sd) + return scipy.stats.norm.pdf(x=k, loc=mean, scale=sd) diff --git a/tests/test_reed_frost.py b/tests/test_reed_frost.py index efbd4ca..d0912c4 100755 --- a/tests/test_reed_frost.py +++ b/tests/test_reed_frost.py @@ -1,3 +1,4 @@ +import numpy as np import pytest from pytest import approx @@ -13,6 +14,13 @@ def test_pmf_2(): assert rf.pmf(2, 2, p, m=1) == approx(p**2 + 2 * p**2 * (1 - p), abs=1e-6) +def test_pmf_vector(): + """PMF can take vectors""" + result = rf.pmf(k=np.array([0, 1, 2]), n=2, p=0.1, m=1) + assert isinstance(result, np.ndarray) + assert len(result) == 3 + + def test_large_dist_warning(): with pytest.raises(RuntimeWarning): - rf.dist_large(n=10, lambda_=0.5) + rf.pmf_large(k=1, n=10, lambda_=0.5)