Finding the parameters of a SIR model

Basic elements of a SIR model

In our model, each agent transitions among 3 states: Susceptible, Infectious and Recovered. At each epoch, an Infectious agent has a probability β of infecting its Susceptible neighbours, and a probability γ to transition to the Recovered state. From that moment on, it will no longer participate in the spreading of the disease.

The connectivity between agents is modeled as a Watts-Strogatz small world random graph, a regular ring lattice of mean degree K where each node has a probability r of being randomly rewired. In our model, these parameters will be input-calibrated (i.e., fixed).

# preparatory imports
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

from black_it.calibrator import Calibrator
from black_it.loss_functions.minkowski import MinkowskiLoss
from black_it.plot.plot_results import (
    plot_convergence,
    plot_losses,
    plot_losses_interact,
    plot_sampling,
    plot_sampling_interact,
)
from black_it.samplers.best_batch import BestBatchSampler
from black_it.samplers.halton import HaltonSampler
from black_it.samplers.random_forest import RandomForestSampler

cmap = matplotlib.cm.get_cmap("tab10").colors

def plotSeries(title, SIR1, SIR2=None):
    fig, ax1 = plt.subplots()
    ax1.plot(SIR1[:, 0], "-", label="susceptible", color=cmap[0])
    ax1.plot(SIR1[:, 1], "-", label="infected", color=cmap[1])
    ax1.plot(SIR1[:, 2], "-", label="recovered", color=cmap[2])
    ax1.tick_params(axis="y", labelcolor=cmap[0])
    ax1.set_ylabel("susceptible per 1000 inhabitants", color=cmap[0])
    ax1.set_xlabel("weeks")
    ax1.legend(loc=5)
    if SIR2 is not None:
        ax1.plot(SIR2[:, 0], "--", label="susceptible", color=cmap[0])
        ax1.plot(SIR2[:, 1], "--", label="infected", color=cmap[1])
        ax1.plot(SIR2[:, 2], "--", label="recovered", color=cmap[2])

    ax1.title.set_text(title)
    fig.tight_layout()
    plt.show()


def printBestParams(params):
    with np.printoptions(suppress=True, formatter={"float_kind": "{:.3f}   ".format}):
        print(f"                  brktime  beta1    beta2    gamma")
        print(f"Best parameters: {params[0]}")

Simple calibration over synthetic data

Let's start by using black-it to recover the $\beta$ and $\gamma$ parameter of a SIR model from synthetic data.

Initialize a calibrator object

To set up a calibration one needs to define the following components first:

1. a model to be calibrated
2. a loss function to measure the distance between the real time series and the simulated time series
3. a set of samplers that iteratively suggest a set of parameter values to explore
4. the parameter space that should be explored

1. Model simulator

In order to use the SIR simulator, we need to download its docker image:

1. Install Docker following the instructions given here
2. Open Docker on your computer
3. Run the following command in your terminal: docker pull bancaditalia/abmsimulator

We can then proceed to generate the data we are going to try to reproduce via calibration.

from models.sir.sir_docker import SIR

true_params = [0.1, 0.1]  # in general not known!
synth_data = SIR(true_params, 50, 0)

plotSeries("Synthetic data", synth_data)

2. Loss function

# import a quadratic loss, a simple squared difference bewteen the two series
from black_it.loss_functions.minkowski import MinkowskiLoss

loss = MinkowskiLoss()

3. Samplers

from black_it.samplers.best_batch import BestBatchSampler
from black_it.samplers.halton import HaltonSampler
from black_it.samplers.random_forest import RandomForestSampler

batch_size = 4

halton_sampler = HaltonSampler(batch_size=batch_size)
random_forest_sampler = RandomForestSampler(batch_size=batch_size)
best_batch_sampler = BestBatchSampler(batch_size=batch_size)

samplers = [halton_sampler, random_forest_sampler, best_batch_sampler]

4. Parameter space (bounds and precision)

bounds = [[0.001, 0.001], [1.00, 1.00]]

precisions = [0.001, 0.001]

Finally, initialize the Calibrator

from black_it.calibrator import Calibrator

saving_folder='sir-test'
# initialize a Calibrator object
cal = Calibrator(
    samplers=samplers,
    real_data=synth_data,
    model=SIR,
    parameters_bounds=bounds,
    parameters_precision=precisions,
    ensemble_size=1,
    loss_function=loss,
    saving_folder=saving_folder,
)

***
Number of free params:       2.
Explorable param space size: 1000000.
***

Selecting 4 processes for the parallel evaluation of the model

Calibration

Calibrate the model for five batches

params, losses = cal.calibrate(5)

BATCH NUMBER:   1
PARAMS SAMPLED: 0

METHOD: HaltonSampler
----> sim exec elapsed time: 3.1s
---->   min loss new params: 1093.27
---->   avg loss new params: 1387.94
----> avg loss exist params: 1387.94
---->         curr min loss: 1093.2679706997856
====>    total elapsed time: 3.1s

METHOD: RandomForestSampler
----> sim exec elapsed time: 2.6s
---->   min loss new params: 994.25
---->   avg loss new params: 1401.16
----> avg loss exist params: 1394.55
---->         curr min loss: 994.2479459016207
====>    total elapsed time: 3.5s

METHOD: BestBatchSampler
----> sim exec elapsed time: 3.1s
---->   min loss new params: 1091.99
---->   avg loss new params: 1218.39
----> avg loss exist params: 1335.83
---->         curr min loss: 994.2479459016207
====>    total elapsed time: 3.1s
Checkpoint saved in 0.0s

BATCH NUMBER:   2
PARAMS SAMPLED: 12

METHOD: HaltonSampler
----> sim exec elapsed time: 7.1s
---->   min loss new params: 1401.9
---->   avg loss new params: 2014.58
----> avg loss exist params: 1505.52
---->         curr min loss: 994.2479459016207
====>    total elapsed time: 7.1s

METHOD: RandomForestSampler
----> sim exec elapsed time: 4.9s
---->   min loss new params: 1413.63
---->   avg loss new params: 2200.89
----> avg loss exist params: 1644.59
---->         curr min loss: 994.2479459016207
====>    total elapsed time: 6.2s

METHOD: BestBatchSampler
----> sim exec elapsed time: 4.7s
---->   min loss new params: 1084.36
---->   avg loss new params: 1109.2
----> avg loss exist params: 1555.36
---->         curr min loss: 994.2479459016207
====>    total elapsed time: 4.7s
Checkpoint saved in 0.0s

BATCH NUMBER:   3
PARAMS SAMPLED: 24

METHOD: HaltonSampler
----> sim exec elapsed time: 4.4s
---->   min loss new params: 880.29
---->   avg loss new params: 1508.88
----> avg loss exist params: 1548.72
---->         curr min loss: 880.2915915967067
====>    total elapsed time: 4.4s

METHOD: RandomForestSampler
----> sim exec elapsed time: 3.5s
---->   min loss new params: 850.56
---->   avg loss new params: 1041.11
----> avg loss exist params: 1485.27
---->         curr min loss: 850.5637796224285
====>    total elapsed time: 4.9s

METHOD: BestBatchSampler
----> sim exec elapsed time: 6.4s
---->   min loss new params: 881.79
---->   avg loss new params: 934.11
----> avg loss exist params: 1424.03
---->         curr min loss: 850.5637796224285
====>    total elapsed time: 6.5s
Checkpoint saved in 0.0s

BATCH NUMBER:   4
PARAMS SAMPLED: 36

METHOD: HaltonSampler
----> sim exec elapsed time: 4.0s
---->   min loss new params: 1670.77
---->   avg loss new params: 2426.46
----> avg loss exist params: 1524.27
---->         curr min loss: 850.5637796224285
====>    total elapsed time: 4.0s

METHOD: RandomForestSampler
----> sim exec elapsed time: 3.5s
---->   min loss new params: 331.47
---->   avg loss new params: 683.43
----> avg loss exist params: 1447.83
---->         curr min loss: 331.4650895275124
====>    total elapsed time: 4.6s

METHOD: BestBatchSampler
----> sim exec elapsed time: 3.6s
---->   min loss new params: 315.78
---->   avg loss new params: 675.06
----> avg loss exist params: 1383.44
---->         curr min loss: 315.7767534283539
====>    total elapsed time: 3.6s
Checkpoint saved in 0.0s

BATCH NUMBER:   5
PARAMS SAMPLED: 48

METHOD: HaltonSampler
----> sim exec elapsed time: 2.9s
---->   min loss new params: 745.54
---->   avg loss new params: 1471.68
----> avg loss exist params: 1390.22
---->         curr min loss: 315.7767534283539
====>    total elapsed time: 2.9s

METHOD: RandomForestSampler
----> sim exec elapsed time: 2.5s
---->   min loss new params: 943.59
---->   avg loss new params: 1157.78
----> avg loss exist params: 1373.62
---->         curr min loss: 315.7767534283539
====>    total elapsed time: 3.5s

METHOD: BestBatchSampler
----> sim exec elapsed time: 4.7s
---->   min loss new params: 498.19
---->   avg loss new params: 574.39
----> avg loss exist params: 1320.34
---->         curr min loss: 315.7767534283539
====>    total elapsed time: 4.7s
Checkpoint saved in 0.0s

# best parameters obtained so far
params[0]

array([0.088, 0.136])

# index of the best fitting series
idxmin = np.argmin(cal.losses_samp)

Compare the synthetic and calibrated series

plotSeries("Synthetic vs calibrated series", synth_data, cal.series_samp[idxmin, 0])

Calibration of a more advanced SIR model against realistic data

In this part of the tutorial we will use black-it to find the parameters of a SIR model fitted on the italian Covid-19 epidemiological data.

We will see that a proper modelling of the first wave of the epidemic requires the introduction of a structural break in the SIR simulator i.e., a specific point in time in which an abrupt change in the parameters occurs.

This is useful to model the effect of the lockdown over the spreading of the epidemic.

Load reference data

For didactic puposes, let's load a previously prepared dataset containing a very rough estimate of the SIR data for the first 20 weeks of the italian Covid-19 epidemic.

As the official data underestimates the number of cases, Susceptible and Recovered were rescaled by a constant factor.

real_data = np.loadtxt("data/italy_20_weeks.txt")

# plot the real time series we want to reproduce
plotSeries("Real data", real_data)

Initialize a calibrator object

1. Model simulator

from models.sir.sir_docker import SIR_w_breaks

2. Loss function

# we'll use a quadratic loss, a simple squared difference bewteen the two series
loss = MinkowskiLoss()

3. Samplers

sampler_batch_size = 16
samplers = [
    HaltonSampler(batch_size=sampler_batch_size),
    RandomForestSampler(batch_size=sampler_batch_size),
    BestBatchSampler(batch_size=sampler_batch_size),
]

4. Parameter space (bounds and precision)

#    brktime, beta1, beta2, gamma
bounds_w_breaks = [
    [2, 0.1, 0,   0.1],
    [7, 0.2, 0.1, 0.3],
]
precisions_w_breaks = [1, 0.0005, 0.0005, 0.0005]

Initialize the Calibrator

saving_folder = "output"

# initialize the Calibrator
cal = Calibrator(
    samplers=samplers,
    real_data=real_data,
    model=SIR_w_breaks,
    parameters_bounds=bounds_w_breaks,
    parameters_precision=precisions_w_breaks,
    ensemble_size=1,
    loss_function=loss,
    saving_folder=saving_folder,
)

***
Number of free params:       4.
Explorable param space size: 97204806.
***

Selecting 4 processes for the parallel evaluation of the model

Calibration

Perform 5 calibration rounds.

Note that, with these parameters, you would be able to achieve a much lower loss in 30 epochs.

params, losses = cal.calibrate(5)

BATCH NUMBER:   1
PARAMS SAMPLED: 0

METHOD: HaltonSampler
----> sim exec elapsed time: 10.0s
---->   min loss new params: 78.97
---->   avg loss new params: 127.93
----> avg loss exist params: 127.93
---->         curr min loss: 78.96522708183643
====>    total elapsed time: 10.0s

METHOD: RandomForestSampler
----> sim exec elapsed time: 12.5s
---->   min loss new params: 69.4
---->   avg loss new params: 95.8
----> avg loss exist params: 111.87
---->         curr min loss: 69.39561568592492
====>    total elapsed time: 14.0s

METHOD: BestBatchSampler
----> sim exec elapsed time: 12.9s
---->   min loss new params: 67.18
---->   avg loss new params: 90.8
----> avg loss exist params: 104.85
---->         curr min loss: 67.1802120840003
====>    total elapsed time: 12.9s
Checkpoint saved in 0.0s

BATCH NUMBER:   2
PARAMS SAMPLED: 48

METHOD: HaltonSampler
----> sim exec elapsed time: 15.2s
---->   min loss new params: 58.32
---->   avg loss new params: 122.98
----> avg loss exist params: 109.38
---->         curr min loss: 58.32156435411211
====>    total elapsed time: 15.2s

METHOD: RandomForestSampler
----> sim exec elapsed time: 10.4s
---->   min loss new params: 57.7
---->   avg loss new params: 85.93
----> avg loss exist params: 104.69
---->         curr min loss: 57.70356830875495
====>    total elapsed time: 12.0s

METHOD: BestBatchSampler
----> sim exec elapsed time: 11.2s
---->   min loss new params: 58.89
---->   avg loss new params: 86.32
----> avg loss exist params: 101.63
---->         curr min loss: 57.70356830875495
====>    total elapsed time: 11.2s
Checkpoint saved in 0.0s

BATCH NUMBER:   3
PARAMS SAMPLED: 96

METHOD: HaltonSampler
----> sim exec elapsed time: 13.0s
---->   min loss new params: 64.77
---->   avg loss new params: 113.03
----> avg loss exist params: 103.26
---->         curr min loss: 57.70356830875495
====>    total elapsed time: 13.0s

METHOD: RandomForestSampler
----> sim exec elapsed time: 11.7s
---->   min loss new params: 42.16
---->   avg loss new params: 69.24
----> avg loss exist params: 99.01
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 14.5s

METHOD: BestBatchSampler
----> sim exec elapsed time: 9.3s
---->   min loss new params: 47.7
---->   avg loss new params: 80.67
----> avg loss exist params: 96.97
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 9.3s
Checkpoint saved in 0.0s

BATCH NUMBER:   4
PARAMS SAMPLED: 144

METHOD: HaltonSampler
----> sim exec elapsed time: 8.6s
---->   min loss new params: 78.91
---->   avg loss new params: 131.24
----> avg loss exist params: 100.4
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 8.6s

METHOD: RandomForestSampler
----> sim exec elapsed time: 8.2s
---->   min loss new params: 50.61
---->   avg loss new params: 69.1
----> avg loss exist params: 97.55
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 9.5s

METHOD: BestBatchSampler
----> sim exec elapsed time: 15.5s
---->   min loss new params: 42.16
---->   avg loss new params: 82.51
----> avg loss exist params: 96.3
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 15.5s
Checkpoint saved in 0.1s

BATCH NUMBER:   5
PARAMS SAMPLED: 192

METHOD: HaltonSampler
----> sim exec elapsed time: 10.6s
---->   min loss new params: 77.07
---->   avg loss new params: 124.54
----> avg loss exist params: 98.47
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 10.7s

METHOD: RandomForestSampler
----> sim exec elapsed time: 10.5s
---->   min loss new params: 53.1
---->   avg loss new params: 65.98
----> avg loss exist params: 96.15
---->         curr min loss: 42.162915251802474
====>    total elapsed time: 12.4s

METHOD: BestBatchSampler
----> sim exec elapsed time: 7.7s
---->   min loss new params: 41.11
---->   avg loss new params: 72.03
----> avg loss exist params: 94.54
---->         curr min loss: 41.11389045820134
====>    total elapsed time: 7.7s
Checkpoint saved in 0.1s

# best parameters obtained so far
printBestParams(params)

                  brktime  beta1    beta2    gamma
Best parameters: [7.000    0.175    0.036    0.234   ]

Compare the original and the calibrated time series

idxmin = np.argmin(cal.losses_samp)
plotSeries("real vs calibrated", real_data, cal.series_samp[idxmin, 0])

Plots

plot_sampling(saving_folder)

plot_sampling_interact(saving_folder)

interactive(children=(Dropdown(description='batch_nums', options={'from 0 to 1': [0, 1], 'from 2 to 3': [2, 3]…

plot_losses(saving_folder)

plot_losses_interact(saving_folder)

interactive(children=(Dropdown(description='method_num', options={'HaltonSampler': 0, 'RandomForestSampler': 1…

plot_convergence(saving_folder)