SIRD (Stochastic variant)

Here we present a stochastic variant of the SIRD model presented in SIRD (Discrete variant).

Instead of solving an ODE system, for the stochastic version of the model we perform multiple Gillespie simulation and average their evolution.

Model implementation

The input is the same as the discrete variant, but with the addition of the parameters seed (as it is a stochastic simulation) and the number of simulations with different seeds we want to average n_simulations.

@ac
def model(
    data, seed, n_simulations,
    n: Int(70_000, 500_000) = 70_000,
    initial_i: Int(1, 1_000) = 40,
    initial_r: Int(0, 1_000) = 4,
    initial_d: Int(0, 1_000) = 0,
    beta: Real(0.1, 1.0) = 0.7,
    gamma: Real(0.01, 1.0) = 0.1,
    mu: Real(0.01, 1.0) = 0.01
) -> float:
    pass

First, we implement a function that will perform random actions after a random increment of time. It receives the current amount of people in each compartment, as well as the epidemic variables, and returns the random increment of time and the variation of each compartment.

def gillespie_step(S, I, R, D, beta, gamma, mu):
    lambda_sum = I * (beta * S + gamma + mu)
    probs = {
        "infect": I * beta * S / lambda_sum,
        "rec": I * gamma / lambda_sum,
        "death": I * mu / lambda_sum,
    }

    dt = random.expovariate(lambda_sum)

    choices, weights = zip(*probs.items())
    action = random.choices(choices, weights=weights, k=1)[0]

    dS = dI = dR = dD = 0
    if action == "infect":
        dS = -1
        dI = +1
    elif action == "rec":
        dI = -1
        dR = +1
    elif action == "death":
        dI = -1
        dD = +1

    return dt, dS, dI, dR, dD

Then, we simulate a single evolution for a number of days t_total, or until the number of infected goes to 0 (as no more changes would happen).

def gillespie_simulation(
    seed, t_total,
    n, initial_i, initial_r, initial_d,
    beta, gamma, mu
):
    random.seed(seed)
    evolution = numpy.zeros(shape=(t_total, 4))

    S = n_initial_i - initial_r - initial_d
    I = initial_i
    R = initial_r
    D = initial_d
    evolution[0] = (S, I, R, D)

    t_ellapsed = 0
    while I > 0:
        dt, dS, dI, dR, dD = gillespie_step(
            S, I, R, D, beta, gamma, mu)
        if t_ellapsed + dt > t_total:
            break

        S += dS
        I += dI
        R += dR
        D += dD

        if int(t_ellapsed) != int(t_ellapsed + dt):
            # Next day, register data
            evolution[int(t_ellapsed+dt)] = (S, I, R, D)
        t_ellapsed += dt

    return evolution

Finally, the model function will execute n_simulations and average their result. Note that some simulations might end with 0 infected before simulating the entire time-span due to a bad random start. As we know we are simulating an epidemic, we discard those simulations by checking if the last day we have 0 infected. At the end of all the simulations, we average the results and return the MSE of the real infected vs the predicted infected.

def model(...):
    n_days = data.shape[0]
    # Normalize beta over effective population
    beta = beta / n

    predictions = numpy.zeros(shape=(n_simulations, n_days, 4))

    simulation = 0
    current_seed = seed

    while simulation < n_simulations:
        pred = gillespie_simulation(
            current_seed, n_days,
            n, initial_i, initial_r, initial_d,
            beta, gamma, mu)
        infected = pred[:, 1]

        current_seed += 1

        if infected[-1] != 0:
            predictions[simulation] = pred
            simulation += 1
evolution = predictions.mean(axis=0)

real_infected = data["#infected"]
pred_infected = evolution[:, 1]

return ((real_infected - pred_infected)**2).sum() / n_days

We also set the entrypoint to load the dataset using Pandas, and (for simplicity) hard-coding a n_simulations of 20:

def entrypoint(data, seed):
    data = pandas.read_csv(data, index_col='date')
    cost = model(data, seed, 20)
    print(f"Result: {cost}")

Preparing the model to be run using docker

In order to execute the model in a container using Docker, we do the following:

Dockerfile:

This project does not require additional software to be installed, so we keep the Dockerfile as it is.

requirements.txt:

We add the required packages for model.py to the requirements file:

optilog==0.3.5
scipy==1.9.3
pandas==1.5.1
numpy==1.23.4
Makefile (OPTIONAL):

We change the name of the image:

# ...
# Change the following as desired
IMAGE_NAME=sird_stochastic
TAG=latest
image_full_name=${IMAGE_NAME}:${TAG}
# ...

Example of execution using docker

[...]
Tuning process finished.
========================
Best parameters:
- model(n): 460220
- model(initial_i): 782
- model(initial_r): 54
- model(initial_d): 771
- model(beta): 0.43255751784177965
- model(gamma): 0.11660842844152228
- model(mu): 0.09385044718899432
========================
Best Config JSON:
{'_0_0__n': 460220, '_1_0__initial_i': 782, [...]}