def create_model(change_points, params_model):
with cov19.Cov19Model(**params_model) as model:
lambda_t_log = cov19.lambda_t_with_sigmoids(
pr_median_lambda_0=0.4,
pr_sigma_lambda_0=0.5,
change_points_list=change_points,
)
mu = pm.Lognormal(name="mu", mu=np.log(1 / 8), sigma=0.2)
pr_median_delay = 10
prior_I = cov19.make_prior_I(lambda_t_log,
mu,
pr_median_delay=pr_median_delay)
new_I_t = cov19.SIR(lambda_t_log, mu, pr_I_begin=prior_I)
new_cases_inferred_raw = cov19.delay_cases(
new_I_t,
pr_median_delay=pr_median_delay,
pr_median_scale_delay=0.3)
new_cases_inferred = cov19.week_modulation(new_cases_inferred_raw)
cov19.student_t_likelihood(new_cases_inferred)
return model
# It is not necessary to use it, one can simply remove it and use the default argument
# for pr_I_begin in cov19.SIR
prior_I = cov19.make_prior_I(lambda_t_log, mu, pr_median_delay = pr_median_delay)
# Use lambda_t_log and mu to run the SIR model
new_I_t = cov19.SIR(lambda_t_log,mu, prob_test = prob_test, pr_I_begin = prior_I)
# Delay the cases by a lognormal reporting delay
new_cases_inferred_raw = cov19.delay_cases(new_I_t, pr_median_delay=pr_median_delay,
pr_median_scale_delay=0.3)
# Modulate the inferred cases by a abs(sin(x)) function, to account for weekend effects
new_cases_inferred = cov19.week_modulation(new_cases_inferred_raw)
# Define the likelihood, uses the new_cases_obs set as model parameter
cov19.student_t_likelihood(new_cases_inferred)
# %% [markdown]
# ## MCMC sampling
# %%
trace = pm.sample(model=model, tune=500, draws=1000, init='advi+adapt_diag')
# %%
varnames = cov19.plotting.get_all_free_RVs_names(model)
num_cols = 5
num_rows = int(np.ceil(len(varnames)/num_cols))
x_size = num_cols * 2.5
y_size = num_rows * 2.5
fig, axes = plt.subplots(num_rows, num_cols, figsize = (x_size, y_size),squeeze=False)