Fit a deterministic SIR model to synthetic data using optimization¶
Notable features of this example include:
- Parameter selection, including re-identifying known parameters and defining parameters
- Iteration and visualization
Note that the model used in this notebook is deterministic. Therefore, we expect to be able to exactly re-identify the parameters used to generate the synthetic data. We take an optimization-based approach, and therefore do not capture parameter uncertainty in this example.
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 deterministic SIR model.
Project goal¶
The motivation for this project is to provide a simple illustration of the process of deterministic 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 deterministic, we hope to be able to find the exact parameters that were used to generate the synthetic data.
Technical details¶
Model¶
This notebook defines a simple deterministic SIR model, generates synthetic data using known values for the transmission (β) and recovery (γ) parameters, and then uses Bayesian optimization to calibrate the model. We validate that the parameters identified through the calibration process are a close match to the true values used to generate the synthetic data.
Note¶
In this simple deterministic example (two parameters and three data points), we are able to uniquely re-identify the true parameters used to generate the synthetic data. However, in more complex and stochastic models with more input parameters and sparse data observations, it may not be possible to uniquely re-identify the true parameters, even if the calibration process is functioning correctly. A variety of parameter combinations may yield similar goodness-of-fit metrics or stochastic noise may corrupt the ability to precisely re-identify the true parameters. Regardless, we can compare the goodness-of-fit metric at the true and calibrated parameters to gain trust that the workflow is functioning as expected.
Resource requirements¶
This model can easily be run locally if desired or cloud-hosted versions are available on Google Colab or Binder.
# Configure notebook autoreloading and inline plotting
%load_ext autoreload
%autoreload 2
%matplotlib inline
import optuna
optuna.logging.set_verbosity(optuna.logging.CRITICAL)
import numpy as np
import pandas as pd
import sciris as sc
import matplotlib.pyplot as plt
from IPython.display import display
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=optuna.exceptions.ExperimentalWarning)
Define deterministic SIR model¶
The following function defines a simple deterministic SIR model that evolves in discrete time. We will use this model to create some synthetic data and then try to estimate the parameters of the model using the Optuna algorithm.
def run_sir(pars):
"""
Run a simple SIR model with the given parameters and return the results as a DataFrame.
Parameters
----------
pars : dict
The parameters dictionary with optuna a values stored in "value".
Returns
-------
dataframe
The results of the SIR model with columns of S, I, and R and index of Time.
"""
β = 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):
S[step+1] = S[step] - β*S[step]*I[step]/N*dt
I[step+1] = I[step] + β*S[step]*I[step]/N*dt - γ*I[step]*dt
R[step+1] = R[step] + γ*I[step]*dt
results = pd.DataFrame(dict(
S = S,
I = I,
R = R,
), index=pd.Index(timevec, name='Time'))
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.
# These are the true parameter values will we later re-identify via Optuna.
true_pars = dict(
beta = dict(value=0.2),
gamma = dict(value=0.1),
)
results = run_sir(true_pars) # Run the model with the true parameters to get the results.
observation_times = np.array([20, 40, 80]) # Observe prevalence at these time points
sir_data = pd.DataFrame({
'x': results.loc[observation_times, 'I'] / results.loc[observation_times].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)
ax[0].legend(['S', 'I', 'R'])
ax[1].plot(results['I']/results.sum(axis=1), ls='--', color='black', label='Prevalence')
ax[1].scatter(sir_data.index, sir_data['x'], marker='o', color='black', label='Observed data')
ax[1].legend()
ax[0].set_ylabel('Count')
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 | |
|---|---|
| t | |
| 20 | 0.057142 |
| 40 | 0.155425 |
| 80 | 0.045222 |
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.
# We will calibrate two parameters, beta and gamma, each allowed to take on values between 0 and 1.
calib_pars = dict(
beta = {'low': 0, 'high': 1},
gamma = {'low': 0, 'high': 1},
)
def evaluate(results, calib_data):
"""
Evaluate the results of a simulation against calibration data.
Parameters
----------
results: DataFrame
The results of a simulation.
calib_data: DataFrame
The calibration data to compare against.
Returns
-------
float
The squared error between the simulation results and the calibration data.
"""
squared_error = 0
for step, expected in calib_data.iterrows():
observed = results.loc[step, 'I'] / results.loc[step].sum()
squared_error += (observed - expected['x'])**2
return squared_error
def trial_to_pars(trial):
"""
Map an Optuna trial to a parameter spec.
Parameters
----------
trial: optuna.Trial
The trial to map.
Returns
-------
dict
The parameters dictionary with Optuna a values stored in "value."
"""
pars = sc.dcp(calib_pars)
for name, spec in pars.items():
spec['value'] = trial.suggest_float(name=name, **spec)
return pars
def objective(trial):
"""
Evaluate a trial by running the SIR model and comparing the results to the calibration data.
Parameters
----------
trial: optuna.Trial
The trial to map.
Returns
-------
float
The objective value of the trial.
"""
pars = trial_to_pars(trial) # Map the trial to parameters
result = run_sir(pars) # Run the SIR model with the parameters
objective = evaluate(result, sir_data) # Evaluate the results
return objective
# The following two lines create an Optuna study and optimize the objective function - easy!
study = optuna.create_study()
study.optimize(objective, n_trials=1000)
# Print a summary
sc.colorize(color='blue', string=f'The best parameters identified by the optimization are:\n\
* {study.best_params}\n\
These parameters should be close to the true parameters:\n\
* {true_pars}\n\
The best parameters resulted in a loss of {study.best_value}.')
# Plot the results
fig = optuna.visualization.matplotlib.plot_optimization_history(study);
fig.axes.set_yscale('log')
ax = optuna.visualization.matplotlib.plot_contour(study);
ax.axvline(true_pars['beta']['value'], ls='--', color='black')
ax.axhline(true_pars['gamma']['value'], ls='--', color='black')
ax.scatter(study.best_params['beta'], study.best_params['gamma'], 100, marker='x', color='red', zorder=10);
The best parameters identified by the optimization are:
* {'beta': 0.2002951869649306, 'gamma': 0.10020314369121844}
These parameters should be close to the true parameters:
* {'beta': {'value': 0.2}, 'gamma': {'value': 0.1}}
The best parameters resulted in a loss of 1.0534300748016926e-07.
Plot simulation with best parameter values¶
Great, it looks like the optimization was able to get pretty close to the true parameters. Let's run a simulation with the best parameter values and plot the results.
best_pars = dict(
beta = dict(value=study.best_params['beta']),
gamma = dict(value=study.best_params['gamma']),
)
results = run_sir(best_pars) # Run the model with the best parameters found by the optimizer
# And now lets plot the results
fig, ax = plt.subplots(2,1)
ax[0].plot(results)
ax[0].legend(['S', 'I', 'R'])
ax[1].plot(results['I']/results.sum(axis=1), ls='--', color='black', label='Prevalence')
ax[1].scatter(sir_data.index, sir_data['x'], marker='o', color='black', label='Observed data')
ax[1].legend()
ax[0].set_ylabel('Count')
ax[1].set_ylabel('Prevalence')
ax[1].yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x:.0%}'))
ax[1].set_xlabel('Time (days)');
You can see that the Optuna optimizer does a good job of identifying parameters that yield outputs that match the observed data. In this case, the calibration was able to re-identify the true parameters, giving us confidence to use the same workflow with real-world data in the future.