def model(districts, populations, cases, seed) -> NetworkedSIR:
max_ts = cases["date_reported"].max()
units = [
SIR(district,
populations[i],
dT0=len(cases[(cases.geo_reported == district)
& (cases.date_reported == max_ts)]),
mobility=0.000001) for (i, district) in enumerate(districts)
]
return NetworkedSIR(units, random_seed=seed)
def models(seed, simulated_school_Rt=school_Rt):
school = SIR("school",
population=school_proportion * total_pop,
infectious_period=6.25,
I0=age_ts.sum(level=0)["school"],
dT0=age_ts.loc["school"][-1],
lower_CI=school_T_lb,
upper_CI=school_T_ub,
Rt0=simulated_school_Rt,
random_seed=seed)
nonschool = SIR("nonschool",
population=(1 - school_proportion) * total_pop,
infectious_period=5.0,
I0=age_ts.sum(level=0)["nonschool"],
dT0=age_ts.loc["nonschool"][-1],
lower_CI=school_T_lb,
upper_CI=school_T_ub,
Rt0=nonschool_Rt,
random_seed=seed)
return (school, nonschool)
def model(seed=0):
return NetworkedSIR([
SIR(name="SULSEL",
population=8_819_500,
dT0=historical.iloc[-1][0],
Rt0=Rt_pred[-1],
upper_CI=T_CI_upper[-1],
lower_CI=T_CI_lower[-1],
mobility=0,
I0=historical.sum()[0])
],
random_seed=seed)
def run_simulation(school: SIR, nonschool: SIR):
for _ in range(30):
school.run(1)
nonschool.forward_epi_step(dB=2 * school.dT[-1] //
3) # doesn't preserve population
dT_school = np.array(school.dT)
dT_nonschool = np.array(nonschool.dT)
dT = dT_school + dT_nonschool
var_up_school = np.array(school.upper_CI) - dT_school
var_dn_school = np.array(school.lower_CI) - dT_school
var_up_nonschool = np.array(nonschool.upper_CI) - dT_nonschool
var_dn_nonschool = np.array(nonschool.lower_CI) - dT_nonschool
dT_CI_l = dT - np.sqrt(var_dn_school**2 + var_dn_nonschool**2)
dT_CI_u = dT + np.sqrt(var_up_school**2 + var_up_nonschool**2)
return (dT[1:], dT_CI_l[1:], dT_CI_u[1:])
logger.info("running province-level Rt estimate")
(dates, RR_pred, RR_CI_upper, RR_CI_lower, T_pred, T_CI_upper, T_CI_lower, total_cases, new_cases_ts, anomalies, anomaly_dates)\
= analytical_MPVS(jakarta_cases, CI = CI, smoothing = smoothing, totals=False)
plt.Rt(dates, RR_pred[1:], RR_CI_upper[1:], RR_CI_lower[1:], CI)\
.title("\nDKI Jakarta: Reproductive Number Estimate")\
.xlabel("\ndate")\
.ylabel("$R_t$\n", rotation=0, labelpad=30)\
.annotate(f"\n{window}-day smoothing window, gamma-prior Bayesian estimation method")\
.show()
logger.info("running case-forward prediction")
prediction_period = 14*days
IDN = SIR(name = "IDN", population = 267.7e6, dT0 = T_pred[-1], Rt0 = RR_pred[-1], upper_CI = T_CI_upper[-1], lower_CI = T_CI_lower[-1], mobility = 0, random_seed = 0)\
.run(prediction_period)
plt.daily_cases(dates, T_pred[1:], T_CI_upper[1:], T_CI_lower[1:], new_cases_ts[1:], anomaly_dates, anomalies, CI,
prediction_ts = [
(IDN.dT[:-1], IDN.lower_CI[1:], IDN.upper_CI[1:], None, "predicted cases")
])\
.title("\nDKI Jakarta: Daily Cases")\
.xlabel("\ndate")\
.ylabel("cases\n")\
.annotate("\nBayesian training process on empirical data, with anomalies identified")\
.show()
logger.info("district-level projections")
district_cases = dkij.groupby(["district", "date_positiveresult"])["id"].count().sort_index()
districts = dkij.district.unique()
migration = np.zeros((len(districts), len(districts)))
novax_districts.add(district)
else:
param_tag = "_".join([
vaccination_policy.name(),
f"ve{int(100*vax_effectiveness)}",
f"annualgoal{int(100 * vax_pct_annual_goal)}",
f"Rt_threshold{Rt_threshold}"
])
tag = geo_tag + param_tag
model = SIR(
name=district,
population=N_district,
dT0=np.ones(num_sims) *
(dT_conf_district_smooth[simulation_start] *
T_ratio).astype(int),
Rt0=Rt[simulation_start],
I0=np.ones(num_sims) * I0,
R0=np.ones(num_sims) * R,
D0=np.ones(num_sims) * D,
mortality=(ts.loc[district].dD.cumsum()[simulation_start] /
T_scaled if I0 == 0 else dD0 / (gamma * I0)),
random_seed=0)
model.dD[0] = np.ones(num_sims) * dD0
t = 0
dVx_adm = [np.zeros(len(IN_age_structure))] # administered doses
dVx_eff = [np.zeros(len(IN_age_structure))] # effective doses
dVx_imm = [np.zeros(len(IN_age_structure))] # immunizing doses
for t in range(1 * 365):
adm, eff, imm = vaccination_policy.distribute_doses(
if vax_effectiveness != 1.00:
continue
policies = [RandomVaccineAssignment(daily_vax_doses, IN_age_ratios)]
else:
policies = [
RandomVaccineAssignment(daily_vax_doses, IN_age_ratios),
PrioritizedAssignment(daily_vax_doses, split_by_age(S), [6, 5, 4, 3, 2, 1, 0], "mortality"),
PrioritizedAssignment(daily_vax_doses, split_by_age(S), [1, 2, 3, 4, 0, 5, 6], "contactrate")
]
for vaccination_policy in policies:
model = SIR(
name = state,
population = N,
dT0 = np.ones(num_sims) * (dT_conf_smooth[simulation_start] * T_ratio).astype(int),
Rt0 = Rt[simulation_start] * N/(N - T_sero),
I0 = np.ones(num_sims) * S,
R0 = np.ones(num_sims) * R,
D0 = np.ones(num_sims) * D,
random_seed = 0
)
t = 0
dVx = [np.zeros(len(IN_age_structure))]
# run vax rate forward 1 year, then until 75% of pop is recovered or vax
while (t <= 365) or ((t > 365) and (model.R[-1].mean() + (daily_vax_doses * t))/N < immunity_threshold):
dVx.append(vaccination_policy.distribute_doses(model))
print("::::", state, vax_pct_annual_goal, vax_effectiveness, vaccination_policy.name(), t, np.mean(model.dT[-1]), np.std(model.dT[-1]))
t += 1
if vax_pct_annual_goal == 0 and t > 365 * 5:
break
# first, look at state level predictions
(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower, T_pred, T_CI_upper, T_CI_lower,
total_cases, new_cases_ts, anomalies, anomaly_dates) = analytical_MPVS(
state_ts,
CI=CI,
smoothing=notched_smoothing(window=smoothing),
totals=False)
plt.Rt(dates, Rt_pred[1:], Rt_CI_upper[1:], Rt_CI_lower[1:], CI, ymin=0, ymax=3)\
.title("\nBihar: Reproductive Number Estimate (Covid19India Data)")\
.annotate(f"public data from {str(dates[0]).split()[0]} to {str(dates[-1]).split()[0]}")\
.xlabel("\ndate")\
.ylabel("$R_t$", rotation=0, labelpad=20)\
.show()
Bihar = SIR(name="Bihar",
population=99_000_000,
dT0=T_pred[-1],
Rt0=Rt_pred[-1],
mobility=0,
random_seed=0)
Bihar.run(14)
plt.daily_cases(dates, T_pred[1:], T_CI_upper[1:], T_CI_lower[1:], new_cases_ts[1:], anomaly_dates, anomalies, CI,
prediction_ts = [(Bihar.dT[:-1], Bihar.lower_CI[1:], Bihar.upper_CI[1:], None, "predicted cases")])\
.title("\nBihar: Daily Cases")\
.xlabel("\ndate")\
.ylabel("cases\n")\
.annotate("\nBayesian training process on empirical data, with anomalies identified")\
.show()
# now, do district-level estimation
I_bins=np.array([[0, 0, 0, 5, 6, 7, 10], [0, 0, 0, 5, 6, 7, 45],
[0, 0, 0, 5, 6, 7, 70]]),
age_ratios=np.array([0.2, 0.2, 0.25, 0.1, 0.1, 0.1, 0.05]),
IFRs=np.array([0.01, 0.01, 0.01, 0.02, 0.02, 0.03, 0.04]),
prioritization=[1, 0, 5, 2, 4, 3, 6],
label="test-contact")
result = np.array([
[10, 20, 30, 40, 50, 10, 0],
[10, 20, 30, 40, 45, 0, 0],
[10, 20, 30, 40, 0, 0, 0],
])
model = SIR("testmodel",
100000,
I0=np.array([28, 63, 88]),
D0=np.array([0, 0, 0]),
R0=np.array([0, 0, 0]))
mp.distribute_doses(model, num_sims=3)
print(mp.S_bins)
model = SIR("testmodel",
100000,
I0=np.array([28, 63, 88]),
D0=np.array([0, 0, 0]),
R0=np.array([0, 0, 0]))
cr.distribute_doses(model, num_sims=3)
print(cr.S_bins)
# p = vp.S_bins.copy()
import epimargin.plots as plt
from epimargin.models import SIR
from epimargin.estimators import analytical_MPVS, parametric_scheme_mcmc, branching_random_walk
from epimargin.smoothing import convolution
import numpy as np
import pandas as pd
import pymc3 as pm
sir_model = SIR("test",
population=500000,
I0=100,
dT0=20,
Rt0=1.01,
random_seed=0)
total_t = 0
schedule = [(1.01, 75), (1.4, 75), (0.9, 75)]
R0_timeseries = []
for (R0, t) in schedule:
R0_timeseries += [R0] * t
sir_model.Rt0 = R0
sir_model.run(t)
total_t += t
plt.plot(sir_model.dT)
plt.show()
plt.plot(R0_timeseries, "-", color="black", label="$R_0$")
plt.plot(sir_model.Rt, "-", color="dodgerblue", label="$R_t$")
plt.legend(framealpha=1, handlelength=1, loc="best")
plt.PlotDevice().xlabel("time").ylabel("reproductive rate").adjust(left=0.10,
bottom=0.15,
def setup(district) -> Tuple[Callable[[str], SIR], pd.DataFrame]:
demographics = simulation_initial_conditions.loc[district]
dR_conf = ts.loc[district].dR
dR_conf = dR_conf.reindex(pd.date_range(dR_conf.index.min(), dR_conf.index.max()), fill_value = 0)
dR_conf_smooth = pd.Series(smooth(dR_conf), index = dR_conf.index).clip(0).astype(int)
R_conf_smooth = dR_conf_smooth.cumsum().astype(int)
R0 = R_conf_smooth[data_recency]
dD_conf = ts.loc[district].dD
dD_conf = dD_conf.reindex(pd.date_range(dD_conf.index.min(), dD_conf.index.max()), fill_value = 0)
dD_conf_smooth = pd.Series(smooth(dD_conf), index = dD_conf.index).clip(0).astype(int)
D_conf_smooth = dD_conf_smooth.cumsum().astype(int)
D0 = D_conf_smooth[data_recency]
dT_conf = ts.loc[district].dT
dT_conf = dT_conf.reindex(pd.date_range(dT_conf.index.min(), dT_conf.index.max()), fill_value = 0)
(
dates,
Rt_pred, Rt_CI_upper, Rt_CI_lower,
T_pred, T_CI_upper, T_CI_lower,
total_cases, new_cases_ts,
*_
) = analytical_MPVS(ts.loc[district].dT, CI = CI, smoothing = notched_smoothing(window = smoothing), totals = False)
Rt_estimates = pd.DataFrame(data = {
"dates" : dates,
"Rt_pred" : Rt_pred,
"Rt_CI_upper" : Rt_CI_upper,
"Rt_CI_lower" : Rt_CI_lower,
"T_pred" : T_pred,
"T_CI_upper" : T_CI_upper,
"T_CI_lower" : T_CI_lower,
"total_cases" : total_cases[2:],
"new_cases_ts": new_cases_ts,
})
dT_conf_smooth = pd.Series(smooth(dT_conf), index = dT_conf.index).clip(0).astype(int)
T_conf_smooth = dT_conf_smooth.cumsum().astype(int)
T0 = T_conf_smooth[data_recency]
dT0 = dT_conf_smooth[data_recency]
S0 = max(0, demographics.N_tot - T0)
I0 = max(0, T0 - R0 - D0)
return (
lambda seed = 0: SIR(
name = district,
mortality = demographics[[f"N_{i}" for i in range(7)]] @ np.array(list(TN_IFRs.values()))/demographics.N_tot,
population = demographics.N_tot,
random_seed = seed,
infectious_period = 10,
S0 = S0,
I0 = I0,
R0 = R0,
D0 = D0,
dT0 = dT0,
Rt0 = Rt_estimates.set_index("dates").loc[data_recency].Rt_pred * demographics.N_tot/S0),
Rt_estimates
)
logger.info("running province-level Rt estimate")
(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower, T_pred, T_CI_upper, T_CI_lower, total_cases, new_cases_ts, anomalies, anomaly_dates)\
= analytical_MPVS(new_cases, CI = CI, smoothing = smoothing, totals = False)
plt.Rt(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower, CI)\
.title("\nSouth Sulawesi: Reproductive Number Estimate")\
.xlabel("\ndate")\
.ylabel("$R_t$\n", rotation=0, labelpad=30)\
.annotate(f"\n{window}-day smoothing window, gamma-prior Bayesian estimation method")\
.show()
logger.info("running case-forward prediction")
prediction_period = 14 * days
I0 = (~cases.confirmed.isna()).sum() - (~cases.recovered.isna()).sum() - (
~cases.died.isna()).sum()
IDN = SIR(name = "IDN", population = 8_819_500, dT0 = T_pred[-1], Rt0 = Rt_pred[-1], upper_CI = T_CI_upper[-1], lower_CI = T_CI_lower[-1], mobility = 0, random_seed = 0, I0 = I0)\
.run(prediction_period)
plt.daily_cases(dates, T_pred, T_CI_upper, T_CI_lower, new_cases_ts, anomaly_dates, anomalies, CI,
prediction_ts = [
(IDN.dT[:-1], IDN.lower_CI[1:], IDN.upper_CI[1:], plt.PRED_PURPLE, "predicted cases")
])\
.title("\nSouth Sulawesi: Daily Cases")\
.xlabel("\ndate")\
.ylabel("cases\n")\
.annotate("\nBayesian training process on empirical data, with anomalies identified")\
.show()
# makassar estimates
mak_cases = cases[cases.regency == "Makassar"]
mak_new_cases = mak_cases.confirmed.value_counts().sort_index()
# set up model
from epimargin.models import SIR
num_sims = 100
N0 = 12.48e6
R0, D0 = daily_reports.loc[end][["recovered", "deceased"]]
I0 = smoothed_cases[:end].sum()
dT0 = smoothed_cases[end]
S0 = N0 - I0 - R0 - D0
Rt0 = Rt[-1] * N0 / S0
no_lockdown = SIR(name="no lockdown",
population=N0,
dT0=np.ones(num_sims) * dT0,
Rt0=np.ones(num_sims) * Rt0,
I0=np.ones(num_sims) * I0,
R0=np.ones(num_sims) * R0,
D0=np.ones(num_sims) * D0,
S0=np.ones(num_sims) * S0,
infectious_period=10)
lockdown = SIR(name="partial lockdown",
population=N0,
dT0=np.ones(num_sims) * dT0,
Rt0=np.ones(num_sims) * 0.75 * Rt0,
I0=np.ones(num_sims) * I0,
R0=np.ones(num_sims) * R0,
D0=np.ones(num_sims) * D0,
S0=np.ones(num_sims) * S0,
infectious_period=10)
# run models forward
T_scaled = dT_conf_smooth.cumsum()[simulation_start] * T_ratio
D, R = df.loc[simulation_start, "TT"]["total"][["deceased", "recovered"]]
S = N_natl - T_scaled - D - R
(Rt_dates, Rt_est, *_) = analytical_MPVS(T_ratio * dT_conf_smooth,
CI=CI,
smoothing=lambda _: _,
totals=False)
Rt = dict(zip(Rt_dates, Rt_est))
model = SIR(name=state,
population=N_natl,
dT0=np.ones(num_sims) *
(dT_conf_smooth[simulation_start] * T_ratio).astype(int),
Rt0=Rt[simulation_start],
I0=np.ones(num_sims) * (T_scaled - R - D),
R0=np.ones(num_sims) * R,
D0=np.ones(num_sims) * D,
mortality=mu_TN,
random_seed=0)
while np.mean(model.dT[-1]) > 0:
model.parallel_forward_epi_step()
pd.DataFrame(data = {
"date": [simulation_start + pd.Timedelta(n, "days") for n in range(len(model.dT))],
"dT" : [np.mean(_)/N_natl for _ in model.dT],
"dD" : [np.mean(_)/N_natl for _ in model.dD],
})\
.assign(month = lambda _: _.date.dt.month.astype(str) + "_" + _.date.dt.year.astype(str))\
.groupby("month")\
logger.info("running national-level Rt estimate")
(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower, T_pred, T_CI_upper, T_CI_lower, total_cases, new_cases_ts, anomalies, anomaly_dates)\
= analytical_MPVS(natl_cases, CI = CI, smoothing = smoothing)
plt.Rt(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower, CI, ymin=0, ymax=4)\
.title("\nIndonesia: Reproductive Number Estimate")\
.xlabel("\ndate")\
.ylabel("$R_t$", rotation=0, labelpad=30)\
.annotate(f"\n{window}-day smoothing window, gamma-prior Bayesian estimation method")\
.show()
logger.info("running case-forward prediction")
IDN = SIR("IDN",
267.7e6,
dT0=T_pred[-1],
Rt0=Rt_pred[-1],
mobility=0,
random_seed=0).run(14)
logger.info("province-level projections")
migration = np.zeros((len(provinces), len(provinces)))
estimates = []
max_len = 1 + max(map(len, provinces))
with tqdm(provinces) as progress:
for (province, cases) in province_cases.items():
progress.set_description(f"{province :<{max_len}}")
(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower,
*_) = analytical_MPVS(cases, CI=CI, smoothing=smoothing)
apr_idx = np.argmax(dates > "31 Mar, 2020")
may_idx = np.argmax(dates >= "01 May, 2020")
max_idx = np.argmax(Rt_pred[apr_idx:may_idx])
def model(districts, populations, cases, seed) -> NetworkedSIR:
units = [
SIR(district, populations[i], dT0=cases[i], mobility=0.000001)
for (i, district) in enumerate(districts)
]
return NetworkedSIR(units, random_seed=seed)
(Rt_dates, Rt_est, *_) = analytical_MPVS(T_ratio * dT_conf_district_smooth,
CI=CI,
smoothing=lambda _: _,
totals=False)
Rt = dict(zip(Rt_dates, Rt_est))
daily_rate = vax_pct_annual_goal / 365
daily_vax_doses = int(vax_effectiveness * daily_rate * N_district)
T_scaled = dT_conf_district_smooth.cumsum()[simulation_start] * T_ratio
model = SIR(
name=state,
population=N_district,
dT0=np.ones(num_sims) *
(dT_conf_district_smooth[simulation_start] * T_ratio).astype(int),
Rt0=Rt[simulation_start] * N_district / (N_district - T_scaled),
I0=np.ones(num_sims) * (T_scaled - R - D),
R0=np.ones(num_sims) * R,
D0=np.ones(num_sims) * D,
random_seed=0)
t = 0
dVx = [np.zeros(len(IN_age_structure))]
# run vax rate forward 1 year, then until 75% of pop is recovered or vax
while (t <= 365) or (
(t > 365) and
(model.R[-1].mean() +
(daily_vax_doses * t)) / N_district < immunity_threshold):
dVx.append(Multinomial.rvs(daily_vax_doses, IN_age_ratios))
model.S[-1] -= daily_vax_doses
return (dT[1:], dT_CI_l[1:], dT_CI_u[1:])
rt1_10x = run_simulation(*models(0, 1.10 * school_Rt))
rt1_25x = run_simulation(*models(0, 1.25 * school_Rt))
rt1_50x = run_simulation(*models(0, 1.50 * school_Rt))
(dates, Rt_pred, Rt_CI_upper, Rt_CI_lower, T_pred, T_CI_upper, T_CI_lower, total_cases, new_cases_ts, anomalies, anomaly_dates)\
= analytical_MPVS(age_ts.sum(level = 1), CI = CI, smoothing = smoothing, totals = False)
MAK = SIR(name="MAK",
population=1.339e6,
dT0=T_pred[-1],
Rt0=Rt_pred[-1],
upper_CI=T_CI_upper[-1],
lower_CI=T_CI_lower[-1],
mobility=0,
random_seed=0,
I0=age_ts.sum(level=1).sum()).run(30)
plt.daily_cases(dates, T_pred, T_CI_upper, T_CI_lower, new_cases_ts, anomaly_dates, anomalies, CI,
prediction_ts = [
(MAK.dT[:-1], MAK.lower_CI[1:], MAK.upper_CI[1:], plt.PRED_PURPLE, "current social distancing"),
(*rt1_10x, "orange", "10% increase in school-age $R_t$"),
(*rt1_25x, "mediumseagreen", "25% increase in school-age $R_t$"),
(*rt1_50x, "hotpink", "50% increase in school-age $R_t$"),
])\
.xlabel("\ndate")\
.ylabel("cases\n")
# .title("\nMakassar Daily Cases - school reopening scenarios")\