def test_linear_rates(self): epi = EpiModel(list("SEIR")) epi.add_transition_processes([ ("E", 1.0, "I"), ("I", 1.0, "R"), ]) linear_rates = [ConstantLinearRate(1.0, 1), ConstantLinearRate(1.0, 2)] linear_events = [np.array([0, -1, +1, 0]), np.array([0, 0, -1, +1.])] for r0, r1 in zip(linear_rates, epi.linear_rate_functions): assert (r0(0, [0.1, 0.2, 0.3, 0.4, 0.5]) == r1( 0, [0.1, 0.2, 0.3, 0.4, 0.5])) for e0, e1 in zip(linear_events, epi.linear_event_updates): assert (all([_e0 == _e1 for _e0, _e1 in zip(e0, e1)])) epi = EpiModel(list("SEIR")) _r0 = lambda t, y: 2 + np.cos(t) _r1 = lambda t, y: 2 + np.sin(t) epi.add_transition_processes([ ("E", _r0, "I"), ("I", _r1, "R"), ]) linear_rates = [DynamicLinearRate(_r0, 1), DynamicLinearRate(_r1, 2)] linear_events = [np.array([0, -1, +1, 0]), np.array([0, 0, -1, +1.])] for r0, r1 in zip(linear_rates, epi.linear_rate_functions): assert (r0(0, [0.1, 0.2, 0.3, 0.4, 0.5]) == r1( 0, [0.1, 0.2, 0.3, 0.4, 0.5])) for e0, e1 in zip(linear_events, epi.linear_event_updates): assert (all([_e0 == _e1 for _e0, _e1 in zip(e0, e1)]))
def test_correcting_for_declining_pop_size(self): A, B = list("AB") epi = EpiModel([A, B], 10, correct_for_dynamical_population_size=True) epi.add_transition_processes([ #(None, 0.1, A), ]) epi.add_fusion_processes([ (A, B, 1, B), ]) epi.set_initial_conditions({B: 4, A: 6}) tt = np.linspace(0, 30) result = epi.integrate(tt) #from matplotlib import pyplot as pl #pl.plot(tt, result[A], label=A) #pl.plot(tt, result[B], label=B) epi.correct_for_dynamical_population_size = False result = epi.integrate(tt)