def sweep_gamma_offset(gammas,
offsets,
tau,
sir_model,
strategy="maintain-suppress-time",
t_sim_normed=400 / 14):
results = np.zeros(len(gammas) * len(offsets)).reshape(len(offsets), -1)
if not sir.check_gamma(np.max(gammas),
sir_model,
np.min(offsets),
strategy,
tau=tau,
verbose=True):
raise ValueError("t_i + most negative offset " "is negative!")
sir_model.b_func = sir.Intervention(strategy=strategy, tau=tau)
for col, gamma in enumerate(gammas):
sir_model.gamma = gamma
S_i_expected = 0
if strategy == "maintain-suppress-time":
S_i_expected, f = oi.calc_Sf_opt(sir_model.R0,
sir_model.gamma * tau)
I_i_expected = sir_model.I_of_S(S_i_expected)
sir_model.b_func.S_i_expected = S_i_expected
sir_model.b_func.I_i_expected = I_i_expected
sir_model.b_func.f = f
elif strategy == "fixed":
S_i_expected, sigma = oi.calc_Sb_opt(sir_model.R0, sir_model.gamma,
tau)
sir_model.b_func.sigma = sigma
elif strategy == "full-suppression":
S_i_expected = oi.calc_S_var_opt(sir_model.R0,
sir_model.gamma * tau, 0)
sir_model.b_func.sigma = 0
## calculate intended intervention time
t_i_opt = sir_model.t_of_S(S_i_expected)[0]
for row, offset_val in enumerate(offsets):
sir_model.b_func.t_i = t_i_opt + offset_val
## solve the sir model
## with our intervention b_func
sir_model.reset()
sir_model.integrate(t_sim_normed / gamma)
results[row, col] = sir_model.get_I_max()
return results
def sweep_tau_early_late(taus,
R0s,
offset,
sir_model,
strategy="maintain-suppress-time",
t_sim_max=1000):
sir_model.inits = params.inits_sweep_small
results = np.zeros(len(taus) * len(R0s)).reshape(len(R0s), -1)
sir_model.b_func = sir.Intervention(strategy=strategy)
for col, tau in enumerate(taus):
print(col)
for row, R0 in enumerate(R0s):
sir_model.b_func.tau = tau
sir_model.R0 = R0
S_i_expected = 0
if strategy == "maintain-suppress-time":
S_i_expected, f = oi.calc_Sf_opt(sir_model.R0,
sir_model.gamma * tau)
I_i_expected = sir_model.I_of_S(S_i_expected)
sir_model.b_func.S_i_expected = S_i_expected
sir_model.b_func.I_i_expected = I_i_expected
sir_model.b_func.f = f
elif strategy == "fixed":
S_i_expected, sigma = oi.calc_Sb_opt(sir_model.R0,
sir_model.gamma, tau)
sir_model.b_func.sigma = sigma
elif strategy == "full-suppression":
S_i_expected = oi.calc_S_var_opt(sir_model.R0,
sir_model.gamma * tau, 0)
sir_model.b_func.sigma = 0
t_i_opt = sir_model.t_of_S(S_i_expected)[0]
## numerically solve the sir
## with our b function
## starting too early
sir_model.b_func.t_i = t_i_opt - offset
sir_model.reset()
sir_model.integrate(t_sim_max)
early = sir_model.get_I_max()
## numerically solve the sir
## with our b function
## starting too late
sir_model.b_func.t_i = t_i_opt + offset
sir_model.reset()
sir_model.integrate(t_sim_max)
late = sir_model.get_I_max()
results[row, col] = late - early
pass
return results
def sweep_tau_offset(taus,
offsets,
sir_model,
strategy="maintain-suppress-time",
t_sim_max=1000):
results = np.zeros(len(taus) * len(offsets)).reshape(len(offsets), -1)
sir_model.b_func = sir.Intervention(strategy=strategy)
for col, tau in enumerate(taus):
sir_model.b_func.tau = tau
S_i_expected = 0
if strategy == "maintain-suppress-time":
S_i_expected, f = oi.calc_Sf_opt(sir_model.R0,
sir_model.gamma * tau)
I_i_expected = sir_model.I_of_S(S_i_expected)
sir_model.b_func.S_i_expected = S_i_expected
sir_model.b_func.I_i_expected = I_i_expected
sir_model.b_func.f = f
elif strategy == "fixed":
S_i_expected, sigma = oi.calc_Sb_opt(sir_model.R0, sir_model.gamma,
tau)
sir_model.b_func.sigma = sigma
elif strategy == "full-suppression":
S_i_expected = oi.calc_S_var_opt(sir_model.R0,
sir_model.gamma * tau, 0)
sir_model.b_func.sigma = 0
t_i_opt = sir_model.t_of_S(S_i_expected)[0]
for row, offset_val in enumerate(offsets):
sir_model.b_func.t_i = t_i_opt + offset_val
## numerically solve the sir
## with our mistimed b function
sir_model.reset()
sir_model.integrate(t_sim_max)
results[row, col] = sir_model.get_I_max()
return results
def sweep_gamma_early_late(gammas,
R0s,
offset,
tau,
sir_model,
strategy="maintain-suppress-time",
t_sim_normed=500 / 14):
results = np.zeros(len(gammas) * len(R0s)).reshape(len(R0s), -1)
sir_model.inits = params.inits_sweep_small
if not sir.check_gamma(
np.max(gammas), sir_model, -offset, strategy, tau=tau,
verbose=True):
raise ValueError("t_i + most negative offset " "is negative!")
sir_model.b_func = sir.Intervention(strategy=strategy, tau=tau)
for col, gamma in enumerate(gammas):
print(col)
for row, R0 in enumerate(R0s):
sir_model.gamma = gamma
sir_model.R0 = R0
S_i_expected = 0
if strategy == "maintain-suppress-time":
S_i_expected, f = oi.calc_Sf_opt(
sir_model.R0, sir_model.gamma * sir_model.b_func.tau)
I_i_expected = sir_model.I_of_S(S_i_expected)
sir_model.b_func.S_i_expected = S_i_expected
sir_model.b_func.I_i_expected = I_i_expected
sir_model.b_func.f = f
elif strategy == "fixed":
S_i_expected, sigma = oi.calc_Sb_opt(sir_model.R0,
sir_model.gamma,
sir_model.b_func.tau)
sir_model.b_func.sigma = sigma
elif strategy == "full-suppression":
S_i_expected = oi.calc_S_var_opt(
sir_model.R0, sir_model.gamma * sir_model.b_func.tau, 0)
sir_model.b_func.sigma = 0
t_i_opt = sir_model.t_of_S(S_i_expected)[0]
## numerically solve the sir
## with our b function
## starting too early
sir_model.b_func.t_i = t_i_opt - offset
sir_model.reset()
sir_model.integrate(t_sim_max)
early = sir_model.get_I_max()
## numerically solve the sir
## with our b function
## starting too late
sir_model.b_func.t_i = t_i_opt + offset
sir_model.reset()
sir_model.integrate(t_sim_max)
late = sir_model.get_I_max()
results[row, col] = late - early
pass
return results
# filename: early_late_parameter_sweep.py
# author: Dylan Morris <*****@*****.**>
#
# description: calculate parameter sweeps
# for cost of being early vs. late with R0, tau, gamma
####################################################
import sys
import numpy as np
import pandas as pd
import parameters as params
import InterventionSIR as sir
import optimize_interventions as oi
covid_sir = sir.InterventionSIR(b_func=None,
R0=params.R0_default,
gamma=params.gamma_default,
inits=params.inits_sweep_small)
covid_sir.reset()
tau = params.tau_default
t_sim_max = 1000
def sweep_tau_early_late(taus,
R0s,
offset,
sir_model,
strategy="maintain-suppress-time",
t_sim_max=1000):
sir_model.inits = params.inits_sweep_small
def main(path_to_fixed,
outpath):
## read in data
fixed = pd.read_csv(path_to_fixed)
## set up numerics
taus = params.taus_figure_interventions
covid_sir = sir.InterventionSIR(
b_func = sir.Intervention(),
R0 = params.R0_default,
gamma = params.gamma_default,
inits = params.inits_default)
covid_sir.reset()
t_sim_max = 360
t_plot_max = 300
## set up figure
fig, plots = setup_figure()
null_color = "black"
null_style = "dashed"
null_dashes = (3, 3)
timecourse_alpha = 0.9
tau_plot_darkness = 0.8
tau_crash_lw = mpl.rcParams["lines.linewidth"] * 0.7
## set up legend
color_vals = np.linspace(0.8, 0.3, len(taus))
leg_cmap = plt.cm.Greys
handlelines = []
for color_val, tau in zip(color_vals, taus):
handlelines.append(
mlines.Line2D([], [],
color = leg_cmap(color_val),
label = tau))
## calculate and plot non intervention results
null_time, null_result = covid_sir.integrate_null(t_sim_max)
## plot tau crash below plots of tau
tau_crash = oi.tau_crash(covid_sir.R0) / covid_sir.gamma
plots["imax"].axvline(tau_crash,
color = ps.full_suppression_cmap(
tau_plot_darkness),
lw = tau_crash_lw,
linestyle = "dotted")
plots["ti"].axvline(tau_crash,
color = ps.full_suppression_cmap(
tau_plot_darkness),
lw = tau_crash_lw,
linestyle = "dotted")
for plotname in ["maintain-suppress-time-timecourse",
"fixed-timecourse",
"full-suppression-timecourse"]:
plots[plotname].plot(null_time,
null_result[:, 1],
color = null_color,
ls = null_style,
dashes = null_dashes)
## calculate and plot intervention results
for name, cmap in zip(
["maintain-suppress-time", "fixed", "full-suppression"],
[ps.opt_cmap, ps.fixed_cmap, ps.full_suppression_cmap]):
solids = [] # save solid b lines for plotting later
## set intervention strategy
covid_sir.b_func.strategy = name
## iterate over intervention durations
for i_tau, tau in enumerate(taus):
covid_sir.b_func.tau = tau
color_val = color_vals[i_tau]
S_i_expected = 0
print("optimizing strategy for {} "
"with tau = {}".format(name, tau))
if name == "maintain-suppress-time":
S_i_expected, f = oi.calc_Sf_opt(
covid_sir.R0,
covid_sir.gamma * tau)
I_i_expected = covid_sir.I_of_S(S_i_expected)
covid_sir.b_func.S_i_expected = S_i_expected
covid_sir.b_func.I_i_expected = I_i_expected
covid_sir.b_func.f = f
elif name == "fixed":
S_i_expected, sigma = oi.calc_Sb_opt(
covid_sir.R0,
covid_sir.gamma,
tau)
covid_sir.b_func.sigma = sigma
elif name == "full-suppression":
S_i_expected = oi.calc_S_var_opt(
covid_sir.R0,
covid_sir.gamma * tau,
0)
covid_sir.b_func.sigma = 0
t_i_opt = covid_sir.t_of_S(S_i_expected)[0]
covid_sir.b_func.t_i = t_i_opt
covid_sir.reset()
covid_sir.integrate(t_sim_max)
plotname_a = name + "-timecourse"
plots[plotname_a].plot(
covid_sir.time_ts,
covid_sir.state_ts[:, 1],
color = cmap(color_val),
alpha = timecourse_alpha)
plotname_b = name + "-b"
times = np.array(covid_sir.time_ts)
b_of_t = np.array([
covid_sir.b_func(time,
covid_sir.R0 * covid_sir.gamma,
covid_sir.gamma,
S,
I)
for time, S, I in zip(
times,
covid_sir.state_ts[:, 0],
covid_sir.state_ts[:, 1])]).astype("float")
disc = ((np.abs(np.diff(b_of_t, append = [1])) > 0.01) |
(np.abs(np.diff(b_of_t, prepend = [1])) > 0.01))
not_disc = np.logical_not(disc)
sol_y = np.array(b_of_t)
dash_y = np.array(b_of_t)
dash_y[not_disc] = np.nan
sol_y[disc] = np.nan
solids.append(sol_y)
# save to plot later,
# so that solid lines pass over
# dashed lines
plots[plotname_b].plot(
times,
dash_y,
color = cmap(color_val),
linestyle = "dashed",
lw = height / 4,
dashes = (2, 2),
alpha = timecourse_alpha)
for i_tau, tau in enumerate(taus):
color_val = color_vals[i_tau]
sol_y = solids[i_tau]
plotname_b = name + "-b"
times = np.array(covid_sir.time_ts)
plots[plotname_b].plot(
times,
sol_y,
color = cmap(color_val),
linestyle = "solid",
solid_capstyle = "butt",
alpha = timecourse_alpha)
## plot optima as a function of tau
taus_sweep = params.taus_sweep
print("calculating and plotting t_i and I_max "
"for strategy {}...\n".format(name))
if name == "maintain-suppress-time":
Sfs = [oi.calc_Sf_opt(
covid_sir.R0,
covid_sir.gamma * tau)
for tau in taus_sweep]
t_is = [covid_sir.t_of_S(Sf[0])
for Sf in Sfs]
Imaxes = [oi.I_of_S(S, covid_sir.R0) for S, f in Sfs]
taus_plot = taus_sweep
elif name == "fixed":
t_is = fixed["t_i"]
Imaxes = fixed["Imax"]
taus_plot = fixed["tau"]
elif name == "full-suppression":
S_is = [oi.calc_S_var_opt(
covid_sir.R0,
covid_sir.gamma * tau,
0)
for tau in taus_sweep]
t_is = [covid_sir.t_of_S(S)
for S in S_is]
Imaxes = [oi.I_max_opt_of_S_i(S, 0,
covid_sir.R0,
covid_sir.gamma * tau)
for S, tau in zip(S_is, taus_sweep)]
taus_plot = taus_sweep
plotname_Imax = "imax"
plotname_t_i = "ti"
plots[plotname_Imax].plot(taus_plot,
Imaxes,
color = cmap(tau_plot_darkness),
alpha = timecourse_alpha)
plots[plotname_t_i].plot(taus_plot,
t_is,
color = cmap(tau_plot_darkness),
alpha = timecourse_alpha)
## labeling and styling
# ylim for 0 to 1 plots
ymin = -0.1
ymax = 1.1
plots["maintain-suppress-time-timecourse"].legend(
title = "$\\tau$ (days)",
handles = handlelines)
plots["maintain-suppress-time-timecourse"].set_ylabel(
"prevalence $I(t)$")
plots["maintain-suppress-time-b"].set_ylabel(
"intervention $b(t)$")
plots["imax"].set_ylabel("peak prevalence $I^{\max}$",
rotation = 270,
labelpad = 40)
plots["ti"].set_xlabel("duration $\\tau$ (days)")
plots["ti"].set_ylabel("optimal time $t^{\mathrm{opt}}_i$",
rotation = 270,
labelpad = 45)
plots["fixed-b"].set_xlabel("time (days)")
plots["maintain-suppress-time-timecourse"].set_xlim(0, 300)
plots["maintain-suppress-time-timecourse"].set_ylim(ps.ymin_timecourse,
ps.ymax_timecourse)
plots["maintain-suppress-time-timecourse"].set_yticks(np.arange(0, 0.4, 0.1))
plots["maintain-suppress-time-timecourse"].set_xticks(np.arange(0, 280,
90))
plots["maintain-suppress-time-b"].set_yticks(np.arange(0, 1.1, 0.25))
plots["maintain-suppress-time-b"].set_ylim(ymin, ymax)
plots["imax"].set_xticks(np.arange(0, 120, 30))
plots["imax"].set_xlim(0, 90)
plots["ti"].set_yticks(np.arange(75, 100, 5))
plots["ti"].set_ylim(73, 97)
## add some titles
plots["maintain-suppress-time-timecourse"].set_title("Optimal")
plots["fixed-timecourse"].set_title("Fixed control")
plots["full-suppression-timecourse"].set_title("Full suppression")
plots["full-suppression-timecourse"].set_title("Full suppression")
plots["imax"].set_title("Effect of $\\tau$")
for plotname, plot in plots.items():
if not plot.is_first_col() and not plotname in ["ti", "imax"]:
plt.setp(plot.get_yticklabels(), visible = False)
else:
plot.yaxis.set_major_formatter(FormatStrFormatter('$%g$'))
if plotname in ["ti", "imax"]:
plot.yaxis.set_label_position("right")
plot.yaxis.tick_right()
fig.tight_layout()
fig.align_ylabels()
fig.savefig(outpath)
def main(fixed_path,
outpath):
fixed_dat = pd.read_csv(fixed_path)
fig, plots = setup_figure()
vline_darkness = 0.8
vline_lw = mpl.rcParams["lines.linewidth"] * 0.7
Sfs = [oi.calc_Sf_opt(
params.R0_default,
params.gamma_default * tau)
for tau in params.taus_sweep]
R0s = np.linspace(1.001, 20, 100)
tau_crashes = [oi.tau_crash(R0) / params.gamma_default
for R0 in R0s]
tau_crit = fixed_dat[fixed_dat["R_e_i"] > 1]["tau"].min()
taus = [tau_crit, 80, 160]
tau_crit_lab = "$\\tau_1 \\approx$ {:.0f}".format(tau_crit)
print(tau_crit_lab)
tau_labels = [tau_crit_lab] + taus[1:]
plots["f-of-tau"].axvline(
oi.tau_crash(params.R0_default) / params.gamma_default,
color = ps.full_suppression_cmap(
vline_darkness),
lw = vline_lw,
linestyle = "dotted")
plots["f-of-tau"].axvline(
tau_crit,
color = ps.fixed_cmap(
vline_darkness),
lw = vline_lw,
linestyle = "dotted")
plots["f-of-tau"].plot(
params.taus_sweep,
[item[1] for item in Sfs],
color = ps.opt_cmap(0.8))
plots["f-of-tau"].plot(
fixed_dat["tau"],
fixed_dat["sigma"],
color = ps.fixed_cmap(0.8))
plots["tau-crash-of-R0"].plot(
R0s,
tau_crashes,
color = ps.opt_cmap(0.8))
covid_sir = sir.InterventionSIR(
b_func = sir.Intervention(),
R0 = params.R0_default,
gamma = params.gamma_default,
inits = params.inits_default)
covid_sir.reset()
t_sim_max = 360
t_plot_max = 300
null_color = "black"
null_style = "dashed"
null_dashes = (3, 3)
timecourse_alpha = 0.9
color_vals = np.linspace(0.8, 0.3, len(taus))
leg_cmap = plt.cm.Greys
handlelines = []
for color_val, tau, tau_label in zip(color_vals, taus, tau_labels):
handlelines.append(
mpl.lines.Line2D([], [],
color = leg_cmap(color_val),
label = tau_label))
## calculate and plot non intervention results
null_time, null_result = covid_sir.integrate_null(t_sim_max)
## calculate and plot intervention results
for name, cmap in zip(
["maintain-suppress-time", "fixed"],
[ps.opt_cmap, ps.fixed_cmap]):
## set intervention strategy
covid_sir.b_func.strategy = name
plotname = name + "-timecourse"
plots[plotname].plot(null_time,
null_result[:, 1],
color = null_color,
ls = null_style,
dashes = null_dashes)
## iterate over intervention durations
for i_tau, tau in enumerate(taus):
covid_sir.b_func.tau = tau
color_val = color_vals[i_tau]
S_i_expected = 0
print("optimizing strategy for {} "
"with tau = {}".format(name, tau))
if name == "maintain-suppress-time":
S_i_expected, f = oi.calc_Sf_opt(
covid_sir.R0,
covid_sir.gamma * tau)
I_i_expected = covid_sir.I_of_S(S_i_expected)
covid_sir.b_func.S_i_expected = S_i_expected
covid_sir.b_func.I_i_expected = I_i_expected
covid_sir.b_func.f = f
elif name == "fixed":
S_i_expected, sigma = oi.calc_Sb_opt(
covid_sir.R0,
covid_sir.gamma,
tau)
covid_sir.b_func.sigma = sigma
pass
t_i_opt = covid_sir.t_of_S(S_i_expected)[0]
covid_sir.b_func.t_i = t_i_opt
covid_sir.reset()
covid_sir.integrate(t_sim_max)
plots[plotname].plot(
covid_sir.time_ts,
covid_sir.state_ts[:, 1],
color = cmap(color_val),
alpha = timecourse_alpha)
pass
pass
pass
plots["maintain-suppress-time-timecourse"].set_xlim(0, 300)
plots["maintain-suppress-time-timecourse"].set_ylim(ps.ymin_timecourse,
ps.ymax_timecourse)
plots["maintain-suppress-time-timecourse"].set_yticks(np.arange(0, 0.4, 0.1))
plots["maintain-suppress-time-timecourse"].set_xticks(np.arange(0, 280,
90))
plots["f-of-tau"].set_ylim(bottom = -0.1, top = 1.1)
plots["f-of-tau"].set_yticks([0, 0.25, 0.5, 0.75, 1])
plots["f-of-tau"].set_xticks(np.arange(0, 90, 20))
plots["f-of-tau"].set_ylabel("$f, \sigma$")
plots["f-of-tau"].set_xlabel("$\\tau$ (days)")
plots["tau-crash-of-R0"].set_xlabel("$\mathcal{R}_0$")
plots["tau-crash-of-R0"].set_ylabel("$\\tau_{\mathrm{crit}}$")
plots["tau-crash-of-R0"].set_ylim(bottom = -5, top = 55)
plots["tau-crash-of-R0"].set_yticks(np.arange(0, 60, 15))
plots["tau-crash-of-R0"].set_xticks(np.arange(0, 22, 5))
for plot in plots.values():
plot.yaxis.set_major_formatter(FormatStrFormatter('$%g$'))
plots["maintain-suppress-time-timecourse"].legend(title = "$\\tau$ (days)",
handles = handlelines)
plots["maintain-suppress-time-timecourse"].set_ylabel("prevalence $I(t)$")
plots["maintain-suppress-time-timecourse"].set_title("Optimal")
plots["fixed-timecourse"].set_title("Fixed control")
plots["maintain-suppress-time-timecourse"].set_xlabel("time (days)")
plots["fixed-timecourse"].set_xlabel("time (days)")
fig.tight_layout()
fig.align_ylabels()
fig.savefig(outpath)