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