Fit a stochastic SIR model to synthetic data using Bayesian methods¶
Building on the deterministic SIR example, this notebook demonstrates the calibration of a simple stochastic SIR model to synthetic data. Here, we use a Bayesian approach to capture uncertainty in the parameters and latent trajectories.
Notable features of this example include:
- Parameter selection, including re-identifying known parameters and defining parameters
- Iteration and visualization
To open an interactive version of this notebook, click 
. Note that it may take a few minutes to create the environment and edits will not be saved.
Background¶
The complexity of model calibration warrants supplemental check to ensure workflows and embedded algorithms are functioning as expected. One simple check that you can perform is to generate synthetic data from a known set of parameters and then attempt to re-identify those "true" parameters using the calibration process. This notebook demonstrates this process using a simple stochastic SIR model.
Project goal¶
The motivation for this project is to provide a simple illustration of the process of stochastic model calibration. In this exercise, we create synthetic data as part of the exercise and then re-identify those known values using the same model to validate the calibration workflow. Because the model is stochastic, we will not be able to exactly re-identify the true parameters, but we expect the posterior distribution over parameters to contain the true parameters.
Technical details¶
Model¶
This notebook defines a simple stochastic SIR model, generates synthetic data using known values for the transmission (β) and recovery (γ) parameters, and then uses Bayesian sampling with importance weighting (and resampling) to calibrate the model. We validate that the true parameters lie within the posterior mean and that the trajectories look like the data. A better test would be to hold out some data for validation.
Resource requirements¶
This model can easily be run locally if desired or cloud-hosted versions are available on Google Colab or Binder.
In [1]:
Copied!
# Configure notebook autoreloading and inline plotting
%load_ext autoreload
%autoreload 2
%matplotlib inline
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import sciris as sc
import seaborn as sns
from IPython.display import display
from scipy.special import betaln, gammaln
# Configure notebook autoreloading and inline plotting
%load_ext autoreload
%autoreload 2
%matplotlib inline
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import sciris as sc
import seaborn as sns
from IPython.display import display
from scipy.special import betaln, gammaln
Define stochastic SIR model¶
The following function defines a simple stochastic SIR model that evolves in discrete time. We will use this model to create some synthetic data and then calibrate the model using a Bayesian approach.
In [2]:
Copied!
def run_sir(pars, rand_seed):
"""
Run a simple SIR model with the given parameters and return the results as a DataFrame.
Parameters
----------
pars : dict
The parameters dictionary with Optuna values stored in "value".
Returns
-------
DataFrame
The results of the SIR model with columns of S, I, and R and index of Time.
"""
# Set random seed for reproducibility
np.random.seed(rand_seed)
β = pars['beta']['value']
γ = pars['gamma']['value']
dt = 1 # Day
n_steps = 100
S, I, R = np.zeros(n_steps), np.zeros(n_steps), np.zeros(n_steps)
S[0], I[0], R[0], N = 990, 10, 0, 1000 # Initial conditions
timevec = np.arange(n_steps) * dt
for step in range(n_steps-1):
di = np.random.binomial(S[step], 1 - np.exp(-β*I[step]/N*dt))
dr = np.random.binomial(I[step], 1 - np.exp(-γ*dt))
S[step+1] = S[step] - di
I[step+1] = I[step] + di - dr
R[step+1] = R[step] + dr
results = pd.DataFrame({
'S': S,
'I': I,
'R': R,
}, index=pd.Index(timevec, name='Time'))
results['rand_seed'] = rand_seed # Inefficient use of RAM, but useful for tracking
return results
def run_sir(pars, rand_seed):
"""
Run a simple SIR model with the given parameters and return the results as a DataFrame.
Parameters
----------
pars : dict
The parameters dictionary with Optuna values stored in "value".
Returns
-------
DataFrame
The results of the SIR model with columns of S, I, and R and index of Time.
"""
# Set random seed for reproducibility
np.random.seed(rand_seed)
β = pars['beta']['value']
γ = pars['gamma']['value']
dt = 1 # Day
n_steps = 100
S, I, R = np.zeros(n_steps), np.zeros(n_steps), np.zeros(n_steps)
S[0], I[0], R[0], N = 990, 10, 0, 1000 # Initial conditions
timevec = np.arange(n_steps) * dt
for step in range(n_steps-1):
di = np.random.binomial(S[step], 1 - np.exp(-β*I[step]/N*dt))
dr = np.random.binomial(I[step], 1 - np.exp(-γ*dt))
S[step+1] = S[step] - di
I[step+1] = I[step] + di - dr
R[step+1] = R[step] + dr
results = pd.DataFrame({
'S': S,
'I': I,
'R': R,
}, index=pd.Index(timevec, name='Time'))
results['rand_seed'] = rand_seed # Inefficient use of RAM, but useful for tracking
return results
Generate synthetic data¶
Next, generate the data that we will later use as calibration targets. Because this simple model is deterministic, we only need to run it once for each unique set of input parameters.
In [3]:
Copied!
# These are the true parameter values will we later re-identify via Optuna.
true_pars = {
'beta': {'value': 0.2},
'gamma': {'value': 0.1},
}
# Run the model with the true parameters to get the results
results = run_sir(true_pars, rand_seed=123)
observation_times = np.array([20, 40, 80]) # Observe prevalence at these time points
sir_data = pd.DataFrame({
'x': results.loc[observation_times, 'I'],
'n': results.loc[observation_times, ['S', 'I', 'R']].sum(axis=1)
}, index=pd.Index(observation_times, name='t'))
print('Here is the data we will be fitting to:')
display(sir_data)
# And now lets plot the results
fig, ax = plt.subplots(2,1)
ax[0].plot(results[['S', 'I', 'R']])
ax[0].scatter(sir_data.index, sir_data['x'], marker='o',\
color='black', label='Observed data')
ax[0].legend(['S', 'I', 'R', 'Observed data'])
ax[0].set_ylabel('Count')
ax[1].plot(
results['I'] / results[['S', 'I', 'R']].sum(axis=1),
color='black',
label='Prevalence'
)
ax[1].legend()
ax[1].set_ylabel('Prevalence')
ax[1].yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x:.0%}'))
ax[1].set_xlabel('Time (days)');
# These are the true parameter values will we later re-identify via Optuna.
true_pars = {
'beta': {'value': 0.2},
'gamma': {'value': 0.1},
}
# Run the model with the true parameters to get the results
results = run_sir(true_pars, rand_seed=123)
observation_times = np.array([20, 40, 80]) # Observe prevalence at these time points
sir_data = pd.DataFrame({
'x': results.loc[observation_times, 'I'],
'n': results.loc[observation_times, ['S', 'I', 'R']].sum(axis=1)
}, index=pd.Index(observation_times, name='t'))
print('Here is the data we will be fitting to:')
display(sir_data)
# And now lets plot the results
fig, ax = plt.subplots(2,1)
ax[0].plot(results[['S', 'I', 'R']])
ax[0].scatter(sir_data.index, sir_data['x'], marker='o',\
color='black', label='Observed data')
ax[0].legend(['S', 'I', 'R', 'Observed data'])
ax[0].set_ylabel('Count')
ax[1].plot(
results['I'] / results[['S', 'I', 'R']].sum(axis=1),
color='black',
label='Prevalence'
)
ax[1].legend()
ax[1].set_ylabel('Prevalence')
ax[1].yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x:.0%}'))
ax[1].set_xlabel('Time (days)');
Here is the data we will be fitting to:
| x | n | |
|---|---|---|
| t | ||
| 20 | 72.0 | 1000.0 |
| 40 | 160.0 | 1000.0 |
| 80 | 47.0 | 1000.0 |
Re-identify known parameters¶
Now that we have generated synthetic data (the black dots in the figure above), we can use model calibration to try re-identifying the parameters. The code below calibrates two parameters, β and γ, each between 0 and 1, using the Optuna optimizer.
In [4]:
Copied!
# We will calibrate two parameters, beta and gamma, each allowed to take on
# values between 0 and 1.
calib_pars = {
'beta': {'low': 0, 'high': 0.5},
'gamma': {'low': 0, 'high': 0.3},
}
def sample_from_prior(size=1):
"""
Sample beta and gamma from a uniform prior.
Parameters
----------
size : int
Number of samples to draw.
Returns
-------
Pandas DataFrame
DataFrame of shape (size, 2) with columns [beta, gamma].
"""
beta = np.random.uniform(calib_pars['beta']['low'], calib_pars['beta']['high'], size)
gamma = np.random.uniform(calib_pars['gamma']['low'], calib_pars['gamma']['high'], size)
return pd.DataFrame({'beta': beta, 'gamma': gamma})
def beta_binomial_likelihood(results, calib_data,
kappa = 10.0,
prior_alpha = 1.0,
prior_beta = 1.0
):
"""
Likelihood for Beta-Binomial observation model:
s_obs ~ BetaBinomial(n_obs, alpha=kappa*p_hat, beta=kappa*(1-p_hat))
p_hat is computed as (sim_x + prior_alpha) / (sim_n + prior_alpha + prior_beta).
Parameters
----------
results: DataFrame
The results of a simulation.
calib_data: DataFrame
The calibration data to compare against.
kappa: float
Concentration parameter: larger values imply less over-dispersion.
prior_alpha: float
Prior alpha used to smooth the estimate of p_hat.
prior_beta: float
Prior beta used to smooth the estimate of p_hat.
Returns
-------
float
The likelihood of the observed data given the simulation results.
"""
obs_x = calib_data['x'].values
obs_n = calib_data['n'].values
sim_x = results['I']
sim_n = results.sum(axis=1)
denom = np.maximum(sim_n + prior_alpha + prior_beta, 1e-12) # guard
p_hat = (sim_x + prior_alpha) / denom
p_hat = np.clip(p_hat, 1e-9, 1 - 1e-9)
alpha = kappa * p_hat
beta = kappa * (1.0 - p_hat)
# log binomial coefficient: log C(n, x)
logC = gammaln(obs_n + 1) - gammaln(obs_x + 1) - gammaln(obs_n - obs_x + 1)
loglik = logC + betaln(obs_x + alpha, obs_n - obs_x + beta) - betaln(alpha, beta)
return np.exp(np.sum(loglik)) # Exponentiated sum of logs
# We will calibrate two parameters, beta and gamma, each allowed to take on
# values between 0 and 1.
calib_pars = {
'beta': {'low': 0, 'high': 0.5},
'gamma': {'low': 0, 'high': 0.3},
}
def sample_from_prior(size=1):
"""
Sample beta and gamma from a uniform prior.
Parameters
----------
size : int
Number of samples to draw.
Returns
-------
Pandas DataFrame
DataFrame of shape (size, 2) with columns [beta, gamma].
"""
beta = np.random.uniform(calib_pars['beta']['low'], calib_pars['beta']['high'], size)
gamma = np.random.uniform(calib_pars['gamma']['low'], calib_pars['gamma']['high'], size)
return pd.DataFrame({'beta': beta, 'gamma': gamma})
def beta_binomial_likelihood(results, calib_data,
kappa = 10.0,
prior_alpha = 1.0,
prior_beta = 1.0
):
"""
Likelihood for Beta-Binomial observation model:
s_obs ~ BetaBinomial(n_obs, alpha=kappa*p_hat, beta=kappa*(1-p_hat))
p_hat is computed as (sim_x + prior_alpha) / (sim_n + prior_alpha + prior_beta).
Parameters
----------
results: DataFrame
The results of a simulation.
calib_data: DataFrame
The calibration data to compare against.
kappa: float
Concentration parameter: larger values imply less over-dispersion.
prior_alpha: float
Prior alpha used to smooth the estimate of p_hat.
prior_beta: float
Prior beta used to smooth the estimate of p_hat.
Returns
-------
float
The likelihood of the observed data given the simulation results.
"""
obs_x = calib_data['x'].values
obs_n = calib_data['n'].values
sim_x = results['I']
sim_n = results.sum(axis=1)
denom = np.maximum(sim_n + prior_alpha + prior_beta, 1e-12) # guard
p_hat = (sim_x + prior_alpha) / denom
p_hat = np.clip(p_hat, 1e-9, 1 - 1e-9)
alpha = kappa * p_hat
beta = kappa * (1.0 - p_hat)
# log binomial coefficient: log C(n, x)
logC = gammaln(obs_n + 1) - gammaln(obs_x + 1) - gammaln(obs_n - obs_x + 1)
loglik = logC + betaln(obs_x + alpha, obs_n - obs_x + beta) - betaln(alpha, beta)
return np.exp(np.sum(loglik)) # Exponentiated sum of logs
Perform sampling and importance weighting¶
In [5]:
Copied!
N = 1_000 # Number of prior samples (more is better, but slower)
# For Beta-Binomial, the concentration parameter. Lower kappa --> broader
# posterior and higher effective sample size (ESS) for the same N. But not a
# free parameter, kappa comes from the over-dispersion in the observed data.
kappa = 10.0
# Sample from the prior
prior_samples = sample_from_prior(size=N)
rand_seeds = np.random.randint(0, 1e6, size=2*N)
rand_seeds = np.unique(rand_seeds)[:N]
# Prepare parameter dicts for each sample
sample_pars_list = [
{
'pars': {'beta': {'value': row['beta']}, 'gamma': {'value': row['gamma']}},
'rand_seed': rand_seeds[idx] # Random seed for each simulation
}
for idx, row in prior_samples.iterrows()
]
# Run simulations and collect trajectories. Each run is fast, so we run them
# serially; sc.parallelize keeps the same interface if you later want to
# parallelize a more expensive model.
sim_results_list = sc.parallelize(run_sir, iterkwargs=sample_pars_list, serial=True)
# Store latent state trajectories for each sample
trajectories = pd.concat(sim_results_list) \
.reset_index() \
.set_index(['rand_seed', 'Time'])
N = 1_000 # Number of prior samples (more is better, but slower)
# For Beta-Binomial, the concentration parameter. Lower kappa --> broader
# posterior and higher effective sample size (ESS) for the same N. But not a
# free parameter, kappa comes from the over-dispersion in the observed data.
kappa = 10.0
# Sample from the prior
prior_samples = sample_from_prior(size=N)
rand_seeds = np.random.randint(0, 1e6, size=2*N)
rand_seeds = np.unique(rand_seeds)[:N]
# Prepare parameter dicts for each sample
sample_pars_list = [
{
'pars': {'beta': {'value': row['beta']}, 'gamma': {'value': row['gamma']}},
'rand_seed': rand_seeds[idx] # Random seed for each simulation
}
for idx, row in prior_samples.iterrows()
]
# Run simulations and collect trajectories. Each run is fast, so we run them
# serially; sc.parallelize keeps the same interface if you later want to
# parallelize a more expensive model.
sim_results_list = sc.parallelize(run_sir, iterkwargs=sample_pars_list, serial=True)
# Store latent state trajectories for each sample
trajectories = pd.concat(sim_results_list) \
.reset_index() \
.set_index(['rand_seed', 'Time'])
In [6]:
Copied!
# Compute likelihoods for each sample
likelihoods = []
for rand_seed, sim_result in trajectories.groupby('rand_seed'):
sim_sir = sim_result.loc[rand_seed].loc[observation_times, ['S', 'I', 'R']]
likelihood = beta_binomial_likelihood(sim_sir, sir_data, kappa=kappa)
likelihoods.append((rand_seed, likelihood))
results = pd.DataFrame(likelihoods, columns=['rand_seed', 'likelihood'])
results = pd.concat([prior_samples, results], axis=1).set_index('rand_seed')
# Compute likelihoods for each sample
likelihoods = []
for rand_seed, sim_result in trajectories.groupby('rand_seed'):
sim_sir = sim_result.loc[rand_seed].loc[observation_times, ['S', 'I', 'R']]
likelihood = beta_binomial_likelihood(sim_sir, sir_data, kappa=kappa)
likelihoods.append((rand_seed, likelihood))
results = pd.DataFrame(likelihoods, columns=['rand_seed', 'likelihood'])
results = pd.concat([prior_samples, results], axis=1).set_index('rand_seed')
Within a Bayesian workflow, there are two ways to proceed from this point.
- Use all samples, weighted by their likelihood. For this approach, there is no resampling step.
- Resample 𝐾 ≤ 𝑁 samples with probability proportional to their likelihood (with replacement). This is sampling-importance resampling (SIR), not to be confused with the susceptible-infected-recovered (SIR) model. Ideally, we would choose 𝐾 = 𝑁 resamples, but if using these samples for further analysis, it may be useful to choose 𝐾 < 𝑁 to reduce computational cost. It's generally recommended to choose 𝐾 as one to two times the effective sample size (ESS).
In [7]:
Copied!
results['weight'] = results['likelihood'] / np.sum(results['likelihood'])
# Importance resampling
ESS = 1 / np.sum(results['weight']**2)
print('='*60)
print(f'Effective Sample Size (ESS) = {ESS:.1f} out of {N}')
if ESS < 30:
print('WARNING: ESS is below 30, consider increasing N')
print('='*60)
K = np.ceil(1.5 * ESS).astype(int) # Number of samples to draw
print('Resampling K =', K, 'samples from the weighted ensemble of N =', N, 'samples')
resample_seeds = np.random.choice(results.index, size=K, replace=True, p=results['weight'])
# Merge results into combined (all) and selected (K<N subset)
combined = trajectories.reset_index().merge(results, on='rand_seed').set_index(['rand_seed', 'Time'])
selected = combined.loc[resample_seeds]
# Show posterior samples as a DataFrame
posterior_pars = results.loc[resample_seeds, ['beta', 'gamma']]
display(posterior_pars.head())
results['weight'] = results['likelihood'] / np.sum(results['likelihood'])
# Importance resampling
ESS = 1 / np.sum(results['weight']**2)
print('='*60)
print(f'Effective Sample Size (ESS) = {ESS:.1f} out of {N}')
if ESS < 30:
print('WARNING: ESS is below 30, consider increasing N')
print('='*60)
K = np.ceil(1.5 * ESS).astype(int) # Number of samples to draw
print('Resampling K =', K, 'samples from the weighted ensemble of N =', N, 'samples')
resample_seeds = np.random.choice(results.index, size=K, replace=True, p=results['weight'])
# Merge results into combined (all) and selected (K<N subset)
combined = trajectories.reset_index().merge(results, on='rand_seed').set_index(['rand_seed', 'Time'])
selected = combined.loc[resample_seeds]
# Show posterior samples as a DataFrame
posterior_pars = results.loc[resample_seeds, ['beta', 'gamma']]
display(posterior_pars.head())
============================================================ Effective Sample Size (ESS) = 107.9 out of 1000 ============================================================ Resampling K = 162 samples from the weighted ensemble of N = 1000 samples
| beta | gamma | |
|---|---|---|
| rand_seed | ||
| 414863 | 0.221336 | 0.070608 |
| 320053 | 0.191597 | 0.084288 |
| 71144 | 0.194937 | 0.082401 |
| 396132 | 0.157581 | 0.100613 |
| 371701 | 0.229411 | 0.122212 |
Let's look at the parameters, full (weighted) posterior mean, and posterior mean from the resample.
In [8]:
Copied!
fig, ax = plt.subplots(figsize=(7, 6))
# Scatter all prior samples, colored by likelihood
scat = ax.scatter(
results['beta'], results['gamma'],
c=results['weight'], cmap='viridis', s=15, edgecolor='none', alpha=0.8
)
# Overlay red rings for the K resampled posterior samples
ax.scatter(
results.loc[resample_seeds, 'beta'], results.loc[resample_seeds, 'gamma'],
facecolors='none', edgecolors='red', s=25, linewidths=1, label='Resampled'
)
# Overlay the true parameter values and best posterior sample
ax.axvline(true_pars['beta']['value'], ls='--', color='black', label='True parameters')
ax.axhline(true_pars['gamma']['value'], ls='--', color='black')
# Mark the full and resampled posterior means
full_posterior_mean = np.average(results[['beta', 'gamma']], weights=results['weight'], axis=0)
ax.scatter(full_posterior_mean[0], full_posterior_mean[1], 100, marker='x', color='red', lw=2, zorder=10, label='Posterior mean (full)')
posterior_mean = results.loc[resample_seeds, ['beta', 'gamma']].mean(axis=0)
ax.scatter(posterior_mean['beta'], posterior_mean['gamma'], 100, marker='x', color='blue', lw=2, zorder=11, label='Posterior mean (resampled)')
ax.set_xlabel('beta')
ax.set_ylabel('gamma')
ax.set_title('Posterior parameter samples (red rings) and prior samples (colored by likelihood)')
plt.colorbar(scat, ax=ax, label='Normalized likelihood')
# Move the legend outside the figure to the right
plt.legend()
plt.tight_layout()
plt.show()
fig, ax = plt.subplots(figsize=(7, 6))
# Scatter all prior samples, colored by likelihood
scat = ax.scatter(
results['beta'], results['gamma'],
c=results['weight'], cmap='viridis', s=15, edgecolor='none', alpha=0.8
)
# Overlay red rings for the K resampled posterior samples
ax.scatter(
results.loc[resample_seeds, 'beta'], results.loc[resample_seeds, 'gamma'],
facecolors='none', edgecolors='red', s=25, linewidths=1, label='Resampled'
)
# Overlay the true parameter values and best posterior sample
ax.axvline(true_pars['beta']['value'], ls='--', color='black', label='True parameters')
ax.axhline(true_pars['gamma']['value'], ls='--', color='black')
# Mark the full and resampled posterior means
full_posterior_mean = np.average(results[['beta', 'gamma']], weights=results['weight'], axis=0)
ax.scatter(full_posterior_mean[0], full_posterior_mean[1], 100, marker='x', color='red', lw=2, zorder=10, label='Posterior mean (full)')
posterior_mean = results.loc[resample_seeds, ['beta', 'gamma']].mean(axis=0)
ax.scatter(posterior_mean['beta'], posterior_mean['gamma'], 100, marker='x', color='blue', lw=2, zorder=11, label='Posterior mean (resampled)')
ax.set_xlabel('beta')
ax.set_ylabel('gamma')
ax.set_title('Posterior parameter samples (red rings) and prior samples (colored by likelihood)')
plt.colorbar(scat, ax=ax, label='Normalized likelihood')
# Move the legend outside the figure to the right
plt.legend()
plt.tight_layout()
plt.show()
Great news! The true parameters (dashed lines) are well within the posterior (red circled samples) and teh posterior means (both full and resampled) are very close to the true values.
View latent trajectories based on all 𝑁 samples, weighted by their likelihood. Takes some code to incorporate the weights into the plotting...
In [9]:
Copied!
# Use ALL samples, with weights, to compute mean and quantiles
df = combined \
.reset_index() \
.melt(
id_vars=['Time', 'rand_seed'],
value_vars=['S', 'I', 'R'],
var_name='State',
value_name='Count'
)
df['Weight'] = df['rand_seed'].map(results['weight']) # Add weights to the DataFrame
def weighted_quantile(values, quantiles, weights):
v = np.asarray(values, float)
q = np.atleast_1d(quantiles).astype(float)
w = np.asarray(weights, float)
order = np.argsort(v)
v, w = v[order], w[order]
cw = np.cumsum(w)
cw /= cw[-1] if cw[-1] > 0 else 1.0
return np.interp(q, cw, v)
def summarize(group):
vals = group['Count'].to_numpy()
wts = group['Weight'].to_numpy()
mean = np.average(vals, weights=wts)
q05, q25, q75, q95 = weighted_quantile(vals, [0.05, 0.25, 0.75, 0.95], wts)
return pd.Series({'mean': mean, 'q05': q05, 'q25': q25, 'q75': q75, 'q95': q95})
summary = (
df.groupby(['Time','State'], sort=True, as_index=False)
.apply(summarize, include_groups=False)
.reset_index(drop=True)
)
fig, ax = plt.subplots(figsize=(9,5))
lines = sns.lineplot(data=summary, x='Time', y='mean', hue='State', hue_order=['S', 'I', 'R'],
estimator=None, errorbar=None, ax=ax, zorder=5)
# Get the colors used by seaborn for each state
state_colors = {line.get_label(): line.get_color() for line in ax.lines}
for state, g in summary.groupby('State'):
color = state_colors.get(state, None)
ax.fill_between(g['Time'], g['q05'], g['q95'], alpha=0.12, linewidth=0, color=color) # 90% band
ax.fill_between(g['Time'], g['q25'], g['q75'], alpha=0.25, linewidth=0, color=color) # 50% band
for ot in observation_times:
ax.axvline(ot, ls='--', color='black')
ax.scatter(sir_data.index, sir_data['x'], marker='o', color='black',\
label='Observed data', zorder=10);
ax.set_ylabel('Count')
ax.set_title('Weighted posterior trajectories')
ax.legend(title='State')
ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
plt.tight_layout()
# Use ALL samples, with weights, to compute mean and quantiles
df = combined \
.reset_index() \
.melt(
id_vars=['Time', 'rand_seed'],
value_vars=['S', 'I', 'R'],
var_name='State',
value_name='Count'
)
df['Weight'] = df['rand_seed'].map(results['weight']) # Add weights to the DataFrame
def weighted_quantile(values, quantiles, weights):
v = np.asarray(values, float)
q = np.atleast_1d(quantiles).astype(float)
w = np.asarray(weights, float)
order = np.argsort(v)
v, w = v[order], w[order]
cw = np.cumsum(w)
cw /= cw[-1] if cw[-1] > 0 else 1.0
return np.interp(q, cw, v)
def summarize(group):
vals = group['Count'].to_numpy()
wts = group['Weight'].to_numpy()
mean = np.average(vals, weights=wts)
q05, q25, q75, q95 = weighted_quantile(vals, [0.05, 0.25, 0.75, 0.95], wts)
return pd.Series({'mean': mean, 'q05': q05, 'q25': q25, 'q75': q75, 'q95': q95})
summary = (
df.groupby(['Time','State'], sort=True, as_index=False)
.apply(summarize, include_groups=False)
.reset_index(drop=True)
)
fig, ax = plt.subplots(figsize=(9,5))
lines = sns.lineplot(data=summary, x='Time', y='mean', hue='State', hue_order=['S', 'I', 'R'],
estimator=None, errorbar=None, ax=ax, zorder=5)
# Get the colors used by seaborn for each state
state_colors = {line.get_label(): line.get_color() for line in ax.lines}
for state, g in summary.groupby('State'):
color = state_colors.get(state, None)
ax.fill_between(g['Time'], g['q05'], g['q95'], alpha=0.12, linewidth=0, color=color) # 90% band
ax.fill_between(g['Time'], g['q25'], g['q75'], alpha=0.25, linewidth=0, color=color) # 50% band
for ot in observation_times:
ax.axvline(ot, ls='--', color='black')
ax.scatter(sir_data.index, sir_data['x'], marker='o', color='black',\
label='Observed data', zorder=10);
ax.set_ylabel('Count')
ax.set_title('Weighted posterior trajectories')
ax.legend(title='State')
ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
plt.tight_layout()
In [10]:
Copied!
# Plot the resampled trajectories. Should be similar to the weighted posterior over latent trajectories shown above.
df = selected \
.reset_index() \
.melt(
id_vars=['Time', 'rand_seed'],
value_vars=['S', 'I', 'R'],
var_name='State',
value_name='Count'
)
fig, ax = plt.subplots(figsize=(9,5))
sns.lineplot(data=df, hue='State', x='Time', y='Count', units='rand_seed', estimator=None, alpha=0.5, lw=0.2, legend=False, ax=ax)
sns.lineplot(data=df, hue='State', x='Time', y='Count', errorbar=('pi', 50), alpha=0.5, ax=ax, legend=True, zorder=5)
sns.lineplot(data=df, hue='State', x='Time', y='Count', errorbar=('pi', 90), alpha=0.25, ax=ax, legend=False, zorder=6)
ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
for ot in observation_times:
ax.axvline(ot, ls='--', color='black')
ax.scatter(sir_data.index, sir_data['x'], marker='o', color='black',\
label='Observed data', zorder=10);
ax.legend(title='State')
ax.set_title('Resampled posterior trajectories')
plt.tight_layout()
# Plot the resampled trajectories. Should be similar to the weighted posterior over latent trajectories shown above.
df = selected \
.reset_index() \
.melt(
id_vars=['Time', 'rand_seed'],
value_vars=['S', 'I', 'R'],
var_name='State',
value_name='Count'
)
fig, ax = plt.subplots(figsize=(9,5))
sns.lineplot(data=df, hue='State', x='Time', y='Count', units='rand_seed', estimator=None, alpha=0.5, lw=0.2, legend=False, ax=ax)
sns.lineplot(data=df, hue='State', x='Time', y='Count', errorbar=('pi', 50), alpha=0.5, ax=ax, legend=True, zorder=5)
sns.lineplot(data=df, hue='State', x='Time', y='Count', errorbar=('pi', 90), alpha=0.25, ax=ax, legend=False, zorder=6)
ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
for ot in observation_times:
ax.axvline(ot, ls='--', color='black')
ax.scatter(sir_data.index, sir_data['x'], marker='o', color='black',\
label='Observed data', zorder=10);
ax.legend(title='State')
ax.set_title('Resampled posterior trajectories')
plt.tight_layout()
The next step would be scenario analysis using these 𝐾 resamples parameters and latent trajectories. Each of these 𝐾 samples represents a different possible future trajectory of the epidemic, and the ensemble of these 𝐾 trajectories can be used to quantify uncertainty in future projections. Each of these 𝐾 trajectories should be simulated forward 𝑀 times for each scenario, using common random numbers, to produce final results.