def _diff(self, dt: float):
state_diff = {comp: 0 for comp in self.model.compartments}
for (src, edges) in self.model.transition_rates.items():
if len(edges) == 1:
[(dest, rate)] = edges
flow = binomial(self.model[src], rate * dt)
# remove flow from src component
state_diff[src] -= flow
# add flow to dest component
state_diff[dest] += flow
else:
assert src == "E"
# add the last compartment as the "catchall" i.e. the probability that they stay exposed.
infectious_rates = np.array([rate
for (_, rate) in edges] + [0])
flow = multinomial(self.model[src], infectious_rates * dt)
state_diff[src] -= sum(flow[:-1])
for index, (dest, _) in enumerate(edges):
state_diff[dest] += flow[index]
return state_diff
def binomial(trials, p, shape=[]):
"""binomial(trials, p) or binomial(trials, p, [n, m, ...]) returns array of binomially distributed random integers.
trials is the number of trials in the binomial distribution.
p is the probability of an event in each trial of the binomial distribution."""
if shape == []:
shape = None
return mt.binomial(trials, p, shape)
def test_n_zero_stream(self):
# Regression test for gh-14522. Ensure that future versions
# generate the same variates as version 1.16.
np.random.seed(8675309)
expected = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[3, 4, 2, 3, 3, 1, 5, 3, 1, 3]])
assert_array_equal(random.binomial([[0], [10]], 0.25, size=(2, 10)),
expected)
def binomial(trials, p, shape=[]):
"""binomial(trials, p) or binomial(trials, p, [n, m, ...]) returns array of binomially distributed random integers.
trials is the number of trials in the binomial distribution.
p is the probability of an event in each trial of the binomial distribution."""
if shape == []:
shape = None
return mt.binomial(trials, p, shape)
def test_p_zero_stream(self):
# Regression test for gh-14522. Ensure that future versions
# generate the same variates as version 1.16.
np.random.seed(12345)
assert_array_equal(random.binomial(1, [0, 0.25, 0.5, 0.75, 1]),
[0, 0, 0, 1, 1])
