COVID-19 Stochastic SIR

A simple stochastic SIR model for COVID-19

CovidSirStochasticModel(input)

Bases: Model

A simple stochastic SIR compartmental model for COVID-19.

Uses tau-leaping: at each timestep the number of infection and recovery events is drawn from a Poisson distribution whose mean equals the deterministic rate. This gives demographic stochasticity while keeping the same expected trajectory as the ODE model.

Sets STOCHASTIC = True so the SimulationManager uses fixed-step Euler integration instead of the adaptive ODE solver.

Initialize the model and seed its PRNG for stochastic draws.

Falls back to beta = 0.4 and a 7-day recovery period when the config omits them. A seed in the config gives reproducible trajectories; otherwise the PRNG is seeded from system entropy.

Parameters:

Name	Type	Description	Default
input	dict	Validated simulation configuration providing initial_population, transmission_dict, start_date, time_steps, admin_units, and an optional seed.	required

Source code in compartment/models/test_covid_sir_stochastic/model.py

def __init__(self, input):
    """Initialize the model and seed its PRNG for stochastic draws.

    Falls back to ``beta = 0.4`` and a 7-day recovery period when the config
    omits them. A ``seed`` in the config gives reproducible trajectories;
    otherwise the PRNG is seeded from system entropy.

    Args:
        input (dict): Validated simulation configuration providing
            ``initial_population``, ``transmission_dict``, ``start_date``,
            ``time_steps``, ``admin_units``, and an optional ``seed``.
    """
    self.population_matrix = np.array(input["initial_population"]).T
    self.compartment_list = list(self.COMPARTMENTS)

    self._load_transmission_params(input.get("transmission_dict", {}))

    if self.beta is None:
        self.beta = 0.4
    if self.gamma is None:
        self.gamma = 1.0 / 7.0

    self.start_date = input["start_date"]
    self.start_date_ordinal = self.start_date.toordinal()
    self.n_timesteps = input["time_steps"]
    self.admin_units = input["admin_units"]

    self.intervention_dict = {}
    self.intervention_statuses = {}
    self.payload = input

    # PRNG key for stochastic draws — use system entropy by default
    # so every run produces a different trajectory.  Pass "seed" in
    # the config to get reproducible results.
    import time
    seed = input.get("seed", None)
    if seed is None:
        seed = int(time.time() * 1000) % (2**31)
    self._key = jax.random.PRNGKey(seed)

define_parameters(schema) classmethod

Declare the stochastic SIR compartments, edges, and parameters.

Parameters:

Name	Type	Description	Default
schema	ParameterSchemaBuilder	Schema builder to populate with model info, the S/I/R compartments, the beta and gamma transmission edges, and the stochastic run count.	required

Source code in compartment/models/test_covid_sir_stochastic/model.py

@classmethod
def define_parameters(cls, schema):
    """Declare the stochastic SIR compartments, edges, and parameters.

    Args:
        schema (ParameterSchemaBuilder): Schema builder to populate with
            model info, the S/I/R compartments, the ``beta`` and ``gamma``
            transmission edges, and the stochastic run count.
    """
    schema.set_model_info(
        disease_type="COVID_SIR_STOCHASTIC",
        label="COVID-19 Stochastic SIR",
        description="A simple stochastic SIR model for COVID-19",
    )

    schema.add_compartment("S", "Susceptible", "Population susceptible to infection")
    schema.add_compartment("I", "Infected", "Currently infected population", infective=True)
    schema.add_compartment("R", "Recovered", "Recovered and immune population")

    schema.add_transmission_parameter(
        source="susceptible",
        target="infected",
        variable_name="beta",
        frequency_dependent=True,
        label="Transmission Rate (S->I)",
        description="Rate at which susceptible individuals become infected",
        default=0.3,
        default_min=0.2,
        default_max=0.6,
        min_value=0.01,
        max_value=2.0,
        unit="per day",
    )

    schema.add_transmission_parameter(
        source="infected",
        target="recovered",
        variable_name="gamma",
        label="Recovery Period (I->R)",
        description="Average number of days for an infected individual to recover",
        default=7.0,
        default_min=5.0,
        default_max=14.0,
        min_value=1.0,
        max_value=30.0,
        unit="days",
        value_type=ValueType.DAYS,
    )

    schema.add_parameter(
        name="num_runs",
        label="Number of Runs",
        description="Number of stochastic trajectories to simulate.",
        value_type=ValueType.INTEGER,
        default=30,
        min_value=5,
        max_value=50,
        enable_variance=False,
    )

prepare_initial_state()

Return the initial compartment populations for the solver.

Returns:

Type	Description
ndarray	Population matrix (compartments x admin zones) used as the solver's initial state.

Source code in compartment/models/test_covid_sir_stochastic/model.py

def prepare_initial_state(self):
    """Return the initial compartment populations for the solver.

    Returns:
        jnp.ndarray: Population matrix (compartments x admin zones)
            used as the solver's initial state.
    """
    # No inter-zone travel — the base class supplies the identity matrix.
    return self.population_matrix

equation(y, t, p)

Tau-leaping stochastic SIR step.

Returns the per-step change (delta); the Euler integrator applies y_{t+1} = y_t + dt * equation(...). Infection and recovery counts are Poisson draws around the deterministic rates, clamped so no more individuals move than exist.

Parameters:

Name	Type	Description	Default
y	ndarray	Current compartment values, ordered by compartment_list.	required
t	float	Current time in days since the simulation start date.	required
p	tuple	Packed parameter tuple, unpacked via _unpack_params.	required

Returns:

Type	Description
ndarray	Stacked per-compartment deltas for this step.

Source code in compartment/models/test_covid_sir_stochastic/model.py

def equation(self, y, t, p):
    """Tau-leaping stochastic SIR step.

    Returns the per-step change (delta); the Euler integrator applies
    ``y_{t+1} = y_t + dt * equation(...)``. Infection and recovery counts are
    Poisson draws around the deterministic rates, clamped so no more
    individuals move than exist.

    Args:
        y (jnp.ndarray): Current compartment values, ordered by
            ``compartment_list``.
        t (float): Current time in days since the simulation start date.
        p (tuple): Packed parameter tuple, unpacked via ``_unpack_params``.

    Returns:
        jnp.ndarray: Stacked per-compartment deltas for this step.
    """
    C = self.COMPARTMENTS
    params = self._unpack_params(p)

    states = {c: y[i] for i, c in enumerate(C)}
    S = states[C.S]
    I = states[C.I]  # noqa: E741

    non_total = [c for c in C if not c.endswith("_total")]
    N = sum(states[c] for c in non_total)

    beta = params["beta"]
    gamma = params["gamma"]

    # Expected events per day
    infection_rate = beta * S * I / (N + 1e-10)
    recovery_rate = gamma * I

    # Poisson draws
    self._key, k1, k2 = jax.random.split(self._key, 3)
    new_infections = jax.random.poisson(k1, infection_rate).astype(S.dtype)
    new_recoveries = jax.random.poisson(k2, recovery_rate).astype(I.dtype)

    # Clamp so we don't move more people than exist
    new_infections = np.minimum(new_infections, S)
    new_recoveries = np.minimum(new_recoveries, I)

    # Build derivatives for all compartments (including auto-generated _totals)
    zero = np.zeros_like(S)
    derivs = {c: zero for c in C}

    derivs[C.S] = -new_infections
    derivs[C.I] = new_infections - new_recoveries
    derivs[C.R] = new_recoveries

    # Auto-generated cumulative tracking
    if C.I_total in derivs:
        derivs[C.I_total] = new_infections
    if C.R_total in derivs:
        derivs[C.R_total] = new_recoveries

    return np.stack([derivs[c] for c in C])