=========================
SIRD (Stochastic variant)
=========================

Here we present a stochastic variant of the SIRD model
presented in :ref:`sird_discrete`.

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 :code:`seed` (as it is a stochastic simulation) and the number
of simulations with different seeds we want to average :code:`n_simulations`.

.. code-block:: python

    @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.

.. code-block:: python

    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 :code:`t_total`,
or until the number of infected goes to 0 (as no more changes would happen).

.. code-block:: python

    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 :code:`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.


.. code-block:: python

    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 :code:`n_simulations` of 20:

.. code-block:: python

    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:

    .. code-block::

        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:

    .. code-block:: Makefile

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


---------------------------------
Example of execution using docker
---------------------------------

.. code-block::

    [...]
    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, [...]}