Example Disease with Stochasticity

A stochastic SIR model for an example disease

ExampleStochasticModel(config)

Bases: Model

A stochastic SIR compartmental model for Example Disease.

Setting STOCHASTIC = True tells the SimulationManager to integrate with fixed-step Euler instead of the adaptive ODE solver, and to run several trajectories and report a median with an interval band.

equation() is written by hand using tau-leaping: the number of infection and recovery events each step is drawn from a Poisson distribution whose mean equals the deterministic rate. This adds demographic randomness while keeping the same expected trajectory as the ODE version.

Initialize the model and seed its PRNG for stochastic draws.

Parameters:

Name	Type	Description	Default
config	dict	Validated simulation configuration. A seed entry gives reproducible trajectories; otherwise the PRNG is seeded from system entropy.	required

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

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

    Args:
        config (dict): Validated simulation configuration. A ``seed`` entry
            gives reproducible trajectories; otherwise the PRNG is seeded
            from system entropy.
    """
    super().__init__(config)
    # Stochastic models need a PRNG key for the random draws in equation().
    # Default to system entropy so each run differs; pass "seed" in the
    # config for reproducible trajectories.
    seed = config.get("seed") if isinstance(config, dict) else None
    if seed is None:
        seed = int(time.time() * 1000) % (2**31)
    self._key = jax.random.PRNGKey(seed)

define_parameters(schema) classmethod

Declare the model's compartments, transmission edges, and parameters.

Called once by the framework to build the model schema, from which the config validator and parameter set are generated. Also declares the aggregate I_total / R_total cumulative compartments and the stochastic run count.

Parameters:

Name	Type	Description	Default
schema	ParameterSchemaBuilder	Schema builder to populate with model info, metadata, compartments, transmission edges, disease parameters, and the intervention.	required

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

@classmethod
def define_parameters(cls, schema):
    """Declare the model's compartments, transmission edges, and parameters.

    Called once by the framework to build the model schema, from which the
    config validator and parameter set are generated. Also declares the
    aggregate ``I_total`` / ``R_total`` cumulative compartments and the
    stochastic run count.

    Args:
        schema (ParameterSchemaBuilder): Schema builder to populate with
            model info, metadata, compartments, transmission edges, disease
            parameters, and the intervention.
    """
    schema.set_model_info(
        disease_type="example_stochastic",
        label="Example Disease with Stochasticity",
        description="A stochastic SIR model for an example disease",
    )
    schema.set_model_metadata(
        authors=[
            {
                "name": "Example Author",
                "email": "example@example.com",
                "affiliation": "Example LLC",
            }
        ],
        license="MIT",
        model_type="Compartmental",
        diseases=["Example disease"],
        transmission_routes=["Airborne"],
        questions_answered=[
            "How much does demographic stochasticity change outbreak size and timing?",
            "How do asymptomatic and symptomatic infections jointly drive transmission?",
            "How does a transmission-reducing intervention change the trajectory?",
        ],
        key_assumptions=[
            "Closed population — no births or deaths.",
            "Tau-leaping: infection and recovery events are Poisson draws around the deterministic rates.",
            "Two infectious compartments (asymptomatic and symptomatic), each with its own transmission rate and its own recovery period.",
            "A fixed fraction of new infections are asymptomatic, regardless of which compartment infected them.",
            "Frequency-dependent transmission (force of infection scales with the proportion infectious).",
            "Recovered individuals are fully immune with no waning (no R→S transition).",
        ],
    )

    # --- Compartments ---
    # Two infectious compartments — asymptomatic and symptomatic — both
    # contribute to the force of infection, so both are marked infective.
    # They are combined into a single "Infected" curve for graphing via
    # COMPARTMENT_DELTA_GROUPING above.
    schema.add_compartment("S", "Susceptible", "Population susceptible to infection")
    schema.add_compartment("A", "Asymptomatic", "Infectious but showing no symptoms", infective=True)
    schema.add_compartment("Sym", "Symptomatic", "Infectious and showing symptoms", infective=True)
    schema.add_compartment("R", "Recovered", "Recovered and immune")

    # Cumulative tracker for the combined "Infected" group. Like DengueJaxModel,
    # we suppress the framework's per-edge _total compartments (see
    # _add_total_compartments) and declare our own aggregate total here. The
    # delta/"total infected" chart reads the I_total column (group name "I" ->
    # "I_total"); without it the total would collapse to the final-day
    # occupancy of A + Sym (~0 once the outbreak ends).
    schema.add_compartment(
        "I_total",
        "Infected Total",
        "Cumulative infections (asymptomatic + symptomatic combined)",
    )

    schema.add_compartment("R_total", "Recovered Total", "Cumulative recoveries")

    # --- Transmission edges ---
    # We suppress the framework's per-target _total compartments (see
    # _add_total_compartments) and hand-roll the equation, so these edges
    # declare the tunable rates. The two beta rates are the *infectiousness*
    # of the two infectious classes: each contributes to the force of
    # infection in proportion to its own prevalence. Which compartment a new
    # case lands in is decided separately by asymptomatic_fraction.
    schema.add_transmission_parameter(
        source="Susceptible",
        target="Asymptomatic",
        variable_name="beta",
        frequency_dependent=True,
        label="Transmission Rate (Asymptomatic Infectors)",
        description="Rate at which asymptomatic cases infect susceptibles they contact",
        default=0.12,
        default_min=0.04,
        default_max=0.20,
        min_value=0.01,
        max_value=2.0,
        unit="per day",
    )
    schema.add_transmission_parameter(
        source="Susceptible",
        target="Symptomatic",
        variable_name="beta_sym",
        frequency_dependent=True,
        label="Transmission Rate (Symptomatic Infectors)",
        description="Rate at which symptomatic cases infect susceptibles they contact",
        default=0.18,
        default_min=0.06,
        default_max=0.30,
        min_value=0.01,
        max_value=2.0,
        unit="per day",
    )
    # Each infectious compartment recovers on its own clock, so the two
    # gamma rates are independent. Both default to 10 days, which reproduces
    # the single shared recovery period; lengthen gamma_sym if symptomatic
    # cases should stay infectious longer.
    schema.add_transmission_parameter(
        source="Asymptomatic",
        target="Recovered",
        variable_name="gamma",
        label="Recovery Period (Asymptomatic)",
        description="Average number of days for an asymptomatic case to recover",
        default=10.0,
        default_min=5.0,
        default_max=20.0,
        min_value=1.0,
        max_value=100.0,
        value_type=ValueType.DAYS,
        unit="days",
    )
    schema.add_transmission_parameter(
        source="Symptomatic",
        target="Recovered",
        variable_name="gamma_sym",
        label="Recovery Period (Symptomatic)",
        description="Average number of days for a symptomatic case to recover",
        default=10.0,
        default_min=5.0,
        default_max=20.0,
        min_value=1.0,
        max_value=100.0,
        value_type=ValueType.DAYS,
        unit="days",
    )

    # Fraction of new infections that are asymptomatic; the rest are
    # symptomatic. Read as self.asymptomatic_fraction (a percentage). Also
    # splits the day-0 seeding in get_initial_population().
    schema.add_parameter(
        name="asymptomatic_fraction",
        label="Asymptomatic Fraction",
        description="Percentage of new infections that never develop symptoms.",
        value_type=ValueType.PERCENTAGE,
        default=40.0,
        min_value=0.0,
        max_value=100.0,
        unit="%",
        required=False,
        enable_variance=False,
    )

    # Number of stochastic trajectories to simulate. run_simulation reads
    # this to decide how many runs to average for the median + interval.
    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,
    )

    # --- Optional: spatial travel support ---
    # Declare your mobility parameters as custom fields, then define how
    # they build the matrix in build_travel_matrix() below. Without this,
    # the base class supplies an identity matrix (no inter-zone travel).
    # schema.add_parameter(
    #     name="travel_sigma",
    #     label="Travel Rate (σ)",
    #     description="Percentage of each zone's population away from home per day.",
    #     value_type=ValueType.PERCENTAGE,
    #     default=20.0,
    #     min_value=0.0,
    #     max_value=100.0,
    #     unit="%",
    # )

    # --- Optional: interventions ---
    schema.add_intervention(
        id="my_intervention",
        label="My Intervention",
        description="Reduces transmission while active",
        target_rates=["beta", "beta_sym"],
        adherence=50.0,
        transmission_reduction=50.0,
    )

get_initial_population(admin_zones, compartment_list, **kwargs) classmethod

Seed the initial infected across the two infectious compartments.

The base implementation seeds a single I compartment, which this model does not have. Infectious individuals are split between A (asymptomatic) and Sym (symptomatic) using the schema's default asymptomatic fraction.

Parameters:

Name	Type	Description	Default
admin_zones	list[dict]	Admin zone dicts providing population and the infected_population percentage.	required
compartment_list	list[str]	Ordered compartment names, used for column indexing.	required
**kwargs	Any	Additional keyword arguments (unused).	{}

Returns:

Type	Description
ndarray	Initial populations of shape (n_zones, n_compartments).

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

@classmethod
def get_initial_population(cls, admin_zones, compartment_list, **kwargs):
    """Seed the initial infected across the two infectious compartments.

    The base implementation seeds a single ``I`` compartment, which this
    model does not have. Infectious individuals are split between ``A``
    (asymptomatic) and ``Sym`` (symptomatic) using the schema's default
    asymptomatic fraction.

    Args:
        admin_zones (list[dict]): Admin zone dicts providing ``population``
            and the ``infected_population`` percentage.
        compartment_list (list[str]): Ordered compartment names, used for
            column indexing.
        **kwargs (Any): Additional keyword arguments (unused).

    Returns:
        np.ndarray: Initial populations of shape
            ``(n_zones, n_compartments)``.
    """
    # The base implementation seeds a single "I" compartment, but this model
    # has no "I" — infectious individuals live in A (asymptomatic) and Sym
    # (symptomatic). Split the initial infected between them using the
    # schema's default asymptomatic fraction.
    col = {name: i for i, name in enumerate(compartment_list)}
    initial_population = np.zeros((len(admin_zones), len(compartment_list)))

    asymp_pct = 40.0
    schema = cls._get_cached_schema()
    if schema:
        for p in schema.disease_parameters:
            if p.name == "asymptomatic_fraction":
                asymp_pct = p.default
                break
    asymp_frac = asymp_pct / 100.0

    for i, zone in enumerate(admin_zones):
        infected = round(zone["infected_population"] / 100 * zone["population"], 2)
        initial_population[i, col["S"]] = zone["population"] - infected
        initial_population[i, col["A"]] = infected * asymp_frac
        initial_population[i, col["Sym"]] = infected * (1.0 - asymp_frac)

    return initial_population

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/example_stochastic_model/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.
    """
    return self.population_matrix

equation(y, t, p)

Tau-leaping stochastic step with two infectious compartments.

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. Asymptomatic (A) and symptomatic (Sym) cases transmit at their own rates (beta and beta_sym) and both feed a single force of infection; new cases are then split by asymptomatic_fraction. Each compartment also recovers on its own clock (gamma and gamma_sym). The two compartments are combined into one "Infected" curve for graphing via COMPARTMENT_DELTA_GROUPING.

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/example_stochastic_model/model.py

def equation(self, y, t, p):
    """Tau-leaping stochastic step with two infectious compartments.

    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. Asymptomatic (``A``) and
    symptomatic (``Sym``) cases transmit at their own rates (``beta`` and
    ``beta_sym``) and both feed a single force of infection; new cases are
    then split by ``asymptomatic_fraction``. Each compartment also recovers
    on its own clock (``gamma`` and ``gamma_sym``). The two compartments are
    combined into one "Infected" curve for graphing via
    ``COMPARTMENT_DELTA_GROUPING``.

    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(self.compartment_list)}
    S = states[C.S]
    A = states[C.A]
    Sym = states[C.Sym]

    # Both infectious compartments drive the force of infection, but at their
    # own rates, so track their proportions separately as well as combined.
    # prop_infective is what threshold-triggered interventions watch.
    non_total = [c for c in C if not c.endswith("_total")]
    N_total = sum(states[c] for c in non_total)
    denom = N_total.sum() + 1e-10
    prop_A = A.sum() / denom
    prop_Sym = Sym.sum() / denom
    prop_infective = prop_A + prop_Sym

    # _apply_interventions scales target_rates and returns the updated travel
    # matrix. With no interventions configured it returns both unchanged.
    rates, self.travel_matrix = self._apply_interventions(
        t, {"beta": params["beta"], "beta_sym": params["beta_sym"]}, prop_infective
    )
    beta = rates["beta"]
    beta_sym = rates["beta_sym"]
    gamma = params["gamma"]
    gamma_sym = params["gamma_sym"]
    # PERCENTAGE params arrive as e.g. 40.0 — convert to a 0-1 fraction.
    asymp_frac = self._to_rate(self.asymptomatic_fraction, ValueType.PERCENTAGE)

    # Expected events per day. Asymptomatic and symptomatic cases transmit at
    # different rates, so each contributes to the force of infection in
    # proportion to its own prevalence.
    foi = beta * prop_A + beta_sym * prop_Sym
    expected_infections = S * foi
    expected_recoveries_A = gamma * A
    expected_recoveries_Sym = gamma_sym * Sym

    # Draw the actual event counts from Poisson distributions. Split the key
    # each step so successive draws are independent; Euler runs this in a
    # plain Python loop, so mutating self._key here is safe.
    self._key, k_inf, k_rec_a, k_rec_s = jax.random.split(self._key, 4)
    new_infections = jax.random.poisson(k_inf, expected_infections).astype(S.dtype)
    new_infections = jnp.minimum(new_infections, S)

    # Whether a new case is asymptomatic is independent of who infected it,
    # so split the total by asymptomatic_fraction rather than by source.
    new_asymp = new_infections * asymp_frac
    new_sym = new_infections - new_asymp

    # Recoveries drawn separately for each infectious compartment.
    new_rec_a = jax.random.poisson(k_rec_a, expected_recoveries_A).astype(A.dtype)
    new_rec_s = jax.random.poisson(k_rec_s, expected_recoveries_Sym).astype(Sym.dtype)
    new_rec_a = jnp.minimum(new_rec_a, A)
    new_rec_s = jnp.minimum(new_rec_s, Sym)

    # Build derivatives for every compartment (start at zero, add flows).
    derivs = {c: jnp.zeros_like(S) for c in self.compartment_list}
    derivs[C.S] = -new_infections
    derivs[C.A] = new_asymp - new_rec_a
    derivs[C.Sym] = new_sym - new_rec_s
    derivs[C.R] = new_rec_a + new_rec_s

    # Cumulative combined infected tracker (inflow only) so the
    # "total infected" chart reads cumulative incidence
    derivs["I_total"] = new_infections
    derivs["R_total"] = new_rec_a + new_rec_s

    return jnp.stack([derivs[c] for c in self.compartment_list])