def sirt_params():
params = SIRtParams(
gamma=GAMMA,
S0=S0,
I0=I0,
R0=R0
)
return params
def test_fit_model(sirt_model):
n = 50
xatol = 1e-3
t_eval = range(n)
y_obs = sirt_model.simulate(t_eval=t_eval)
start_params = SIRtParams.from_random(seed=42)
start_params.gamma = GAMMA
start_params.S0 = S0
new_model = SIRt(params=start_params, beta_t=BETA_T, t=T)
new_model.fit(y_obs, t_eval, options={'xatol': xatol})
assert (absolute_error(new_model.params.values, sirt_model.params.values) < 0.01).all()
def test_fit_model_and_beta(sirt_model):
n = 50
xatol = 1e-3
t_eval = range(n)
y_obs = sirt_model.simulate(t_eval=t_eval)
start_params = SIRtParams.from_random(seed=42)
start_params.gamma = GAMMA
start_params.S0 = S0
new_model = SIRt(params=start_params)
new_model.fit(y_obs, t_eval, lowess_frac=0.01, options={'xatol': xatol})
assert mean_absolute_error(sirt_model.beta(t_eval), new_model.beta(t_eval)) < 0.2
assert mean_absolute_error(new_model.params.values, sirt_model.params.values) < 0.015
def test_params_from_random():
new_params = SIRtParams.from_random(seed=123)
assert new_params.gamma > 0.0
assert new_params.gamma < 1.0
assert new_params.S0 > 0.9
assert new_params.S0 < 1.0
assert new_params.I0 > 0.0
assert new_params.I0 < 0.1
assert new_params.R0 > 0.0
assert new_params.R0 < 0.1
def test_params_from_values(sirt_params):
vals = sirt_params.values
new_params = SIRtParams.from_values(vals)
assert new_params == sirt_params
def test_not_equal_params():
params1 = SIRtParams(gamma=GAMMA, S0=S0, I0=I0, R0=R0)
params2 = SIRtParams(gamma=GAMMA - 1, S0=S0, I0=I0, R0=R0)
assert params1 != params2
def sirt_model():
params = SIRtParams(gamma=GAMMA, S0=S0, I0=I0, R0=R0)
return SIRt(params=params, beta_t=BETA_T, t=T)