def plot_parameter_traces(param_values_by_chain, max_n_iter=2500):
param_names = list(param_values_by_chain["belgium"].keys())
n_rows = len(param_names) + 1
n_cols = len(OPTI_REGIONS) + 1
w, h = 6, 2
title_w, title_h = 4, 2
fig = plt.figure(
constrained_layout=True,
figsize=(title_w + w * len(OPTI_REGIONS),
title_h + h * len(param_names)),
) # (w, h)
widths = [title_w] + [w] * len(OPTI_REGIONS)
heights = [title_h] + [h] * len(param_names)
spec = fig.add_gridspec(ncols=n_cols,
nrows=n_rows,
width_ratios=widths,
height_ratios=heights)
# load priors to set y range
prior_list = get_prior_distributions_for_opti()
for i_country, country in enumerate(OPTI_REGIONS):
ax = fig.add_subplot(spec[0, i_country + 1])
ax.text(
0.5,
0.2,
COUNTRY_TITLES[country],
fontsize=23,
horizontalalignment="center",
verticalalignment="center",
)
ax.axis("off")
for i_param, param_name in enumerate(param_names):
if param_name in param_values_by_chain[country]:
ax = fig.add_subplot(spec[i_param + 1, i_country + 1])
if i_param < len(param_names) - 1:
ax.axes.get_xaxis().set_visible(False)
param_values = param_values_by_chain[country][param_name]
n_iterations = min(max_n_iter, len(param_values))
param_values = param_values[:n_iterations]
ax.plot(range(n_iterations),
param_values,
"-",
color="royalblue")
ax.grid(False, axis="x")
y_range = param_info[param_name]["range"]
_prior = [
p for p in prior_list if p["param_name"] == param_name
]
if len(_prior) > 0:
prior = _prior[0]
if prior["distribution"] == "uniform":
y_range = prior["distri_params"]
range_w = y_range[1] - y_range[0]
buffer = 0.1 * range_w
y_range = [y_range[0] - buffer, y_range[1] + buffer]
ax.set_ylim(y_range)
if i_country == 0:
ax = fig.add_subplot(spec[i_param + 1, 0])
ax.text(
0.5,
0.2,
param_info[param_name]["name"],
fontsize=20,
horizontalalignment="center",
verticalalignment="center",
)
ax.axis("off")
plt.tight_layout()
plt.savefig("figures/param_traces.pdf")
from apps.covid_19.calibration import base
from autumn.constants import Region
from apps.covid_19.mixing_optimisation.utils import (
get_prior_distributions_for_opti,
get_target_outputs_for_opti,
get_weekly_summed_targets,
add_dispersion_param_prior_for_gaussian,
)
from autumn.tool_kit.utils import print_target_to_plots_from_calibration
country = Region.ITALY
PAR_PRIORS = get_prior_distributions_for_opti()
for i, par in enumerate(PAR_PRIORS):
if par["param_name"] == "contact_rate":
PAR_PRIORS[i]["distri_params"] = [0.025, 0.06]
TARGET_OUTPUTS = get_target_outputs_for_opti(country,
source='who',
data_start_time=61,
data_end_time=182)
# Use weekly counts
for target in TARGET_OUTPUTS:
target["years"], target["values"] = get_weekly_summed_targets(
target["years"], target["values"])
MULTIPLIERS = {}
PAR_PRIORS = add_dispersion_param_prior_for_gaussian(PAR_PRIORS,
def make_posterior_ranges_figure(param_values):
n_panels = len(param_values["belgium"])
country_list = OPTI_REGIONS[::-1]
n_col = 4
n_row = int(n_panels // n_col)
if n_col * n_row < n_panels:
n_row += 1
fig, axs = plt.subplots(n_row, n_col, figsize=(11, 12))
plt.subplots_adjust(left=None,
bottom=None,
right=None,
top=None,
wspace=None,
hspace=0.4)
# load priors to set x range
prior_list = get_prior_distributions_for_opti()
i_col = -1
i_row = 0
for param_name in list(param_values["belgium"].keys()):
i_col += 1
if i_col >= n_col:
i_col = 0
i_row += 1
h = 0
for country in country_list:
h += 1
# find mean and CI
if param_name in param_values[country]:
values = param_values[country][param_name]
point_estimate = mean(values)
low_95, low_50, med, up_50, up_95 = quantile(values,
q=(0.025, 0.25,
0.5, 0.75,
0.975))
axs[i_row, i_col].plot([low_95, up_95], [h, h],
linewidth=1,
color="black")
axs[i_row, i_col].plot([low_50, up_50], [h, h],
linewidth=3,
color="steelblue")
axs[i_row, i_col].plot([point_estimate], [h],
marker="o",
color="crimson",
markersize=5)
axs[i_row, i_col].plot([0], [0])
axs[i_row, i_col].set_ylim((0.5, 6.5))
axs[i_row, i_col].set_title(param_info[param_name]["name"],
fontsize=10.5)
x_range = param_info[param_name]["range"]
_prior = [p for p in prior_list if p["param_name"] == param_name]
if len(_prior) > 0:
prior = _prior[0]
if prior["distribution"] == "uniform":
x_range = prior["distri_params"]
range_w = x_range[1] - x_range[0]
buffer = 0.1 * range_w
x_range = [x_range[0] - buffer, x_range[1] + buffer]
axs[i_row, i_col].set_xlim(x_range)
# Format x-ticks if requested
if "xticks" in param_info[param_name]:
axs[i_row, i_col].set_xticks(param_info[param_name]["xticks"])
axs[i_row,
i_col].set_xticklabels(param_info[param_name]["xlabels"])
if "multiplier" in param_info[param_name]:
axs[i_row, i_col].set_xlabel(param_info[param_name]["multiplier"],
labelpad=-7.5)
# Set y-ticks and yticks-labels
if i_col == 0:
axs[i_row, i_col].set_yticks([1, 2, 3, 4, 5, 6])
axs[i_row, i_col].set_yticklabels([
c.title().replace("United-Kingdom", "UK") for c in country_list
])
else:
axs[i_row, i_col].set_yticks([])
axs[i_row, i_col].grid(False, axis="y")
# Leave blank axis for remaining panels
for i_col_blank in range(i_col + 1, n_col):
axs[i_row, i_col_blank].axis("off")
plt.tight_layout()
plt.savefig("figures/param_posteriors.pdf")