Skip to content

API reference

ChainBinomial

Bases: ABC

General chain binomial model

Source code in src/reedfrost/__init__.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class ChainBinomial(abc.ABC):
    """General chain binomial model"""

    def __init__(self, s0: int, params: dict[str, Any], i0: int = 1):
        assert s0 >= 0
        assert i0 >= 0
        self._validate_params(params)

        self.s0 = s0
        self.i0 = i0
        self.params = params

    @abc.abstractmethod
    def _validate_params(self, params: dict[str, Any]) -> None:
        """Validate parameters for the model"""
        pass

    def prob_final_s(self, s_inf: int) -> float:
        """Probability of a certain final number of susceptibles"""
        assert s_inf >= 0
        return self.prob_state(s=s_inf, i=0, t=self.s0 + 1)

    def prob_final_i_cum_extra(self, k: int) -> float:
        """Probability of a certain number of infection beyond the initial infection(s)"""
        assert k >= 0
        return self.prob_final_s(s_inf=self.s0 - k)

    def prob_final_i_cum(self, i_cum: int) -> float:
        """Probability of a certain number of total infections, including the initial infection(s)"""
        assert i_cum >= 0
        if i_cum < self.i0:
            return 0.0
        else:
            return self.prob_final_i_cum_extra(k=i_cum - self.i0)

    @functools.cache
    def prob_state(self, s: int, i: int, t: int) -> float:
        """
        Probability of being in state (s, i) at time t

        Args:
            s (int): number of susceptibles
            i (int): number of infected
            t (int): generation

        Returns:
            float: probability mass
        """
        assert t >= 0
        assert i >= 0
        assert s >= 0

        if t == 0 and s == self.s0 and i == self.i0:
            return 1.0
        elif t == 0:
            return 0.0
        elif s + i > self.s0:
            return 0.0
        elif t > 0:
            return sum(
                [
                    self._tp(i, s + i, ip) * self.prob_state(s + i, ip, t - 1)
                    for ip in range((self.s0 + self.i0) - (s + i) + 1)
                    if self.prob_state(s + i, ip, t - 1) > 0.0
                ]
            )
        else:
            raise ValueError(f"Negative generation {t}")

    @functools.cache
    def _tp(self, i, si, ip) -> float:
        """Transition probability"""
        return self._pmf_binom(k=i, n=si, p=self._pi(ip))

    @abc.abstractmethod
    def _pi(self, i: int) -> float:
        """Probability of infection given i infected individuals"""
        pass

    @staticmethod
    def _pmf_binom(k: int, n: int, p: float) -> float:
        """
        Binomial distribution pmf

        This implementation is substantially faster than scipy.stats.binom.pmf

        Args:
            k (int): number of successes
            n (int): number of trials
            p (float): success probability

        Returns:
            float: probability mass
        """
        return scipy.special.binom(n, k) * p**k * (1.0 - p) ** (n - k)

    def simulate(
        self,
        rng: np.random.Generator = np.random.default_rng(),
    ) -> NDArray[np.integer]:
        """
        Simulate a Reed-Frost outbreak

        Args:
            rng (np.random.Generator): random number generator

        Returns:
            NDArray[np.integer]: number of infected in each generation
        """
        # time series of infections, starting with the initial infected,
        # is at most of length s + 1
        it = np.zeros(self.s0 + 1, dtype=np.int64)

        if self.i0 == 0:
            return it

        it[0] = self.i0
        s = self.s0

        for t in range(self.s0):
            if it[t] == 0:
                break

            next_i = rng.binomial(n=s, p=self._pi(it[t]))
            it[t + 1] = next_i
            s = s - next_i

        return it

_pi(i) abstractmethod

Probability of infection given i infected individuals

Source code in src/reedfrost/__init__.py 
84
85
86
87
@abc.abstractmethod
def _pi(self, i: int) -> float:
    """Probability of infection given i infected individuals"""
    pass

_pmf_binom(k, n, p) staticmethod

Binomial distribution pmf

This implementation is substantially faster than scipy.stats.binom.pmf

Parameters:

Name Type Description Default
k int

number of successes

 required
n int

number of trials

 required
p float

success probability

 required

Returns:

Name Type Description
float float

probability mass
Source code in src/reedfrost/__init__.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@staticmethod
def _pmf_binom(k: int, n: int, p: float) -> float:
    """
    Binomial distribution pmf

    This implementation is substantially faster than scipy.stats.binom.pmf

    Args:
        k (int): number of successes
        n (int): number of trials
        p (float): success probability

    Returns:
        float: probability mass
    """
    return scipy.special.binom(n, k) * p**k * (1.0 - p) ** (n - k)

_tp(i, si, ip) cached

Transition probability

Source code in src/reedfrost/__init__.py 
79
80
81
82
@functools.cache
def _tp(self, i, si, ip) -> float:
    """Transition probability"""
    return self._pmf_binom(k=i, n=si, p=self._pi(ip))

_validate_params(params) abstractmethod

Validate parameters for the model

Source code in src/reedfrost/__init__.py 
22
23
24
25
@abc.abstractmethod
def _validate_params(self, params: dict[str, Any]) -> None:
    """Validate parameters for the model"""
    pass

prob_final_i_cum(i_cum)

Probability of a certain number of total infections, including the initial infection(s)

Source code in src/reedfrost/__init__.py 
37
38
39
40
41
42
43
def prob_final_i_cum(self, i_cum: int) -> float:
    """Probability of a certain number of total infections, including the initial infection(s)"""
    assert i_cum >= 0
    if i_cum < self.i0:
        return 0.0
    else:
        return self.prob_final_i_cum_extra(k=i_cum - self.i0)

prob_final_i_cum_extra(k)

Probability of a certain number of infection beyond the initial infection(s)

Source code in src/reedfrost/__init__.py 
32
33
34
35
def prob_final_i_cum_extra(self, k: int) -> float:
    """Probability of a certain number of infection beyond the initial infection(s)"""
    assert k >= 0
    return self.prob_final_s(s_inf=self.s0 - k)

prob_final_s(s_inf)

Probability of a certain final number of susceptibles

Source code in src/reedfrost/__init__.py 
27
28
29
30
def prob_final_s(self, s_inf: int) -> float:
    """Probability of a certain final number of susceptibles"""
    assert s_inf >= 0
    return self.prob_state(s=s_inf, i=0, t=self.s0 + 1)

prob_state(s, i, t) cached

Probability of being in state (s, i) at time t

Parameters:

Name Type Description Default
s int

number of susceptibles

 required
i int

number of infected

 required
t int

generation

 required

Returns:

Name Type Description
float float

probability mass
Source code in src/reedfrost/__init__.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
@functools.cache
def prob_state(self, s: int, i: int, t: int) -> float:
    """
    Probability of being in state (s, i) at time t

    Args:
        s (int): number of susceptibles
        i (int): number of infected
        t (int): generation

    Returns:
        float: probability mass
    """
    assert t >= 0
    assert i >= 0
    assert s >= 0

    if t == 0 and s == self.s0 and i == self.i0:
        return 1.0
    elif t == 0:
        return 0.0
    elif s + i > self.s0:
        return 0.0
    elif t > 0:
        return sum(
            [
                self._tp(i, s + i, ip) * self.prob_state(s + i, ip, t - 1)
                for ip in range((self.s0 + self.i0) - (s + i) + 1)
                if self.prob_state(s + i, ip, t - 1) > 0.0
            ]
        )
    else:
        raise ValueError(f"Negative generation {t}")

simulate(rng=np.random.default_rng())

Simulate a Reed-Frost outbreak

Parameters:

Name Type Description Default
rng Generator

random number generator

 default_rng()

Returns:

Type Description
NDArray[integer]

NDArray[np.integer]: number of infected in each generation
Source code in src/reedfrost/__init__.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def simulate(
    self,
    rng: np.random.Generator = np.random.default_rng(),
) -> NDArray[np.integer]:
    """
    Simulate a Reed-Frost outbreak

    Args:
        rng (np.random.Generator): random number generator

    Returns:
        NDArray[np.integer]: number of infected in each generation
    """
    # time series of infections, starting with the initial infected,
    # is at most of length s + 1
    it = np.zeros(self.s0 + 1, dtype=np.int64)

    if self.i0 == 0:
        return it

    it[0] = self.i0
    s = self.s0

    for t in range(self.s0):
        if it[t] == 0:
            break

        next_i = rng.binomial(n=s, p=self._pi(it[t]))
        it[t + 1] = next_i
        s = s - next_i

    return it

Enko

Bases: ChainBinomial

Enko model

Source code in src/reedfrost/__init__.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
class Enko(ChainBinomial):
    """Enko model"""

    @staticmethod
    def _validate_params(params: dict[str, Any]) -> None:
        """Validate parameters for the Enko model"""
        assert "k" in params
        assert params["k"] >= 0.0
        assert "n" in params
        assert params["n"] > 0.0

    def _pi(self, i: int) -> float:
        return 1.0 - (1.0 - i / (self.params["n"] - 1.0)) ** self.params["k"]

_validate_params(params) staticmethod

Validate parameters for the Enko model

Source code in src/reedfrost/__init__.py 
169
170
171
172
173
174
175
@staticmethod
def _validate_params(params: dict[str, Any]) -> None:
    """Validate parameters for the Enko model"""
    assert "k" in params
    assert params["k"] >= 0.0
    assert "n" in params
    assert params["n"] > 0.0

Greenwood

Bases: ChainBinomial

Greenwood model

Source code in src/reedfrost/__init__.py 
153
154
155
156
157
158
159
160
161
162
163
class Greenwood(ChainBinomial):
    """Greenwood model"""

    @staticmethod
    def _validate_params(params: dict[str, Any]) -> None:
        """Validate parameters for the Greenwood model"""
        assert "p" in params
        assert 0.0 <= params["p"] <= 1.0

    def _pi(self, i: int) -> float:
        return self.params["p"] if i > 0 else 0.0

_validate_params(params) staticmethod

Validate parameters for the Greenwood model

Source code in src/reedfrost/__init__.py 
156
157
158
159
160
@staticmethod
def _validate_params(params: dict[str, Any]) -> None:
    """Validate parameters for the Greenwood model"""
    assert "p" in params
    assert 0.0 <= params["p"] <= 1.0

ReedFrost

Bases: ChainBinomial

Reed-Frost model

Source code in src/reedfrost/__init__.py 
140
141
142
143
144
145
146
147
148
149
150
class ReedFrost(ChainBinomial):
    """Reed-Frost model"""

    @staticmethod
    def _validate_params(params: dict[str, Any]) -> None:
        """Validate parameters for the Reed-Frost model"""
        assert "p" in params
        assert 0.0 <= params["p"] <= 1.0

    def _pi(self, i: int) -> float:
        return 1.0 - (1.0 - self.params["p"]) ** i

_validate_params(params) staticmethod

Validate parameters for the Reed-Frost model

Source code in src/reedfrost/__init__.py 
143
144
145
146
147
@staticmethod
def _validate_params(params: dict[str, Any]) -> None:
    """Validate parameters for the Reed-Frost model"""
    assert "p" in params
    assert 0.0 <= params["p"] <= 1.0