Age-Stratified SIR (per-age transmission)

An SIR model in which each age group has its own transmission rate, with age mixing through a contact matrix and spatial mixing through a travel matrix.

AgeStratifiedSIRModel(config)

Bases: Model

A schema-based SIR model whose transmission rate varies by age group.

Initialize the model from a validated simulation config.

Parameters:

Name	Type	Description	Default
config	dict	Validated simulation configuration produced by the framework's config loader.	required

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

def __init__(self, config):
    """Initialize the model from a validated simulation config.

    Args:
        config (dict): Validated simulation configuration produced by the
            framework's config loader.
    """
    super().__init__(config)

    # Effective age-group order for this run: config demographics if the run
    # supplied them, otherwise the schema's declared groups. This order
    # defines axis 1 of the state tensor and the rows/cols of the contact
    # matrix, so the per-age beta vector must follow it exactly.
    if getattr(self, "demographics", None):
        self._age_group_ids = list(self.demographics.keys())
    else:
        schema = type(self)._get_cached_schema()
        self._age_group_ids = [g.id for g in schema.demographic_groups]

    # Rate-dict keys the interventions target, one per age band, in state order.
    self._beta_names = [f"beta_{gid}" for gid in self._age_group_ids]

define_parameters(schema) classmethod

Declare compartments, transmission edges, per-age betas and demographics.

Called once by the framework to build the model schema, from which the config validator and parameter set are generated.

Parameters:

Name	Type	Description	Default
schema	ParameterSchemaBuilder	Schema builder to populate.	required

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

@classmethod
def define_parameters(cls, schema):
    """Declare compartments, transmission edges, per-age betas and demographics.

    Called once by the framework to build the model schema, from which the
    config validator and parameter set are generated.

    Args:
        schema (ParameterSchemaBuilder): Schema builder to populate.
    """
    schema.set_model_info(
        disease_type=cls.DISEASE_TYPE,
        label=cls.DISEASE_LABEL,
        description=cls.DISEASE_DESCRIPTION,
    )
    schema.set_model_metadata(
        model_type="Compartmental",
        diseases=["Generic directly-transmitted infection"],
        transmission_routes=["Direct contact / airborne"],
        questions_answered=[
            "How does a per-age-group transmission rate shape the outbreak "
            "in each age band?",
            "How does age mixing (contact matrix) spread infection from "
            "high-transmission bands to the rest of the population?",
            "How do transmission-reducing interventions change the "
            "age-specific attack rate?",
        ],
        key_assumptions=[
            "Closed population — no births or deaths.",
            "Four age bands (0-17, 18-49, 50-64, 65+), each with its own "
            "transmission rate.",
            "Age mixing follows the country's Prem 2021 synthetic contact "
            "matrix, aggregated to the four bands via each group's age_range.",
            "Homogeneous mixing within an (age group, region) cell; regions "
            "are coupled by a gravity travel matrix.",
            "Transmission is frequency-dependent (force of infection scales "
            "with the proportion infectious, not the raw count).",
            "Recovered individuals are fully immune (no R->S waning).",
        ],
    )

    # ---- Compartments (S-I-R) ----
    # Mark infective=True on compartments that drive the force of infection.
    schema.add_compartment("S", "Susceptible", "Population susceptible to infection")
    schema.add_compartment("I", "Infected", "Currently infectious population", infective=True)
    schema.add_compartment("R", "Recovered", "Recovered and immune")
    # Cumulative infections tracker. Normally the framework auto-generates a
    # ``<target>_total`` compartment for each transmission edge's target;
    # since the S->I force of infection is applied manually (no scalar beta
    # edge to declare), we declare I_total explicitly so cumulative incidence
    # is still tracked. ``_apply_flow`` accumulates into it automatically.
    schema.add_compartment("I_total", "Infected Total", "Cumulative infections")

    # ---- Transmission edges ----
    # Only the I->R recovery edge is declared. Its target (R) auto-generates
    # R_total. The S->I transmission is age-specific and handled manually in
    # equation() from the per-age beta parameters below, so there is no
    # single scalar/average beta edge.
    schema.add_transmission_parameter(
        source="infected",
        target="recovered",
        variable_name="gamma",
        label="Recovery Period (I->R)",
        description="Average number of days 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",
    )

    # ---- Per-age-group transmission rates (first-class parameters) ----
    # One editable transmission rate per age band -- the only transmission
    # rates in the model. Each is an absolute per-day rate with its own
    # default / bounds / uncertainty range and is surfaced as a custom field
    # in the platform UI.
    for group_id, label, _weight, _age_range, default_beta in AGE_GROUPS:
        schema.add_parameter(
            name=f"beta_{group_id}",
            label=f"Transmission Rate — {label}",
            description=(
                f"Per-day transmission rate applied to susceptibles in the "
                f"{label} band. This is the rate at which a susceptible in "
                f"this age group acquires infection per unit of contact-"
                f"weighted infectious pressure."
            ),
            value_type=ValueType.RATE,
            default=default_beta,
            min_value=0.01,
            max_value=0.9,
            default_min=0,
            default_max=1,
            unit="per day",
        )

    # ---- Spatial travel support ----
    schema.add_parameter(
        name="travel_sigma",
        label="Travel Rate (σ)",
        description="Percentage of each region's population away from home per day.",
        value_type=ValueType.PERCENTAGE,
        default=20.0,
        min_value=0.0,
        max_value=100.0,
        unit="%",
    )

    # ---- Intervention ----
    # A single transmission-reducing intervention that targets every per-age
    # transmission rate directly, so it scales each age group's beta while it
    # is active (see equation()).
    age_beta_rates = [f"beta_{g[0]}" for g in AGE_GROUPS]
    schema.add_intervention(
        id="social_isolation",
        label="Social Isolation",
        description="Reduces each age group's transmission rate while active.",
        target_rates=age_beta_rates,
        adherence=50.0,
        transmission_reduction=50.0,
    )

    # ---- Demographics + contact matrix ----
    # Declaring age_range on every group lets the framework auto-load the
    # country's Prem 2021 contact matrix, aggregated to these four bands.
    for group_id, label, weight, age_range, _default_beta in AGE_GROUPS:
        schema.add_demographic_group(
            group_id, label, default_weight=weight, age_range=age_range
        )

build_travel_matrix(admin_zones)

Build an inverse-square gravity travel matrix from travel_sigma.

Parameters:

Name	Type	Description	Default
admin_zones	list[dict]	Admin-zone dicts with center_lat, center_lon and population.	required

Returns:

Type	Description
ndarray	The (n_zones, n_zones) travel matrix.

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

def build_travel_matrix(self, admin_zones):
    """Build an inverse-square gravity travel matrix from ``travel_sigma``.

    Args:
        admin_zones (list[dict]): Admin-zone dicts with ``center_lat``,
            ``center_lon`` and ``population``.

    Returns:
        np.ndarray: The ``(n_zones, n_zones)`` travel matrix.
    """
    # PERCENTAGE params arrive as 20.0, not 0.2 — convert first.
    sigma = self._to_rate(self.travel_sigma, ValueType.PERCENTAGE)
    try:
        return get_gravity_model_travel_matrix(admin_zones, sigma)
    except ValueError:
        # On some pandas/numpy versions the framework helper pivots through
        # pandas and calls np.fill_diagonal on the result; copy-on-write can
        # make DataFrame.to_numpy() return a read-only array, raising
        # "underlying array is read-only". Fall back to an equivalent
        # construction built from writable numpy arrays.
        logger.warning(
            "Gravity travel helper returned a read-only array; using the "
            "writable fallback (results are identical)."
        )
        return self._writable_gravity_travel_matrix(admin_zones, sigma)

prepare_initial_state()

Return the age-stratified initial populations for the solver.

Returns:

Type	Description
ndarray	Populations shaped (compartments, age groups, zones), including the appended _total rows.

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

def prepare_initial_state(self):
    """Return the age-stratified initial populations for the solver.

    Returns:
        jnp.ndarray: Populations shaped (compartments, age groups, zones),
            including the appended ``_total`` rows.
    """
    # Expand (K, R) -> (K, A, R) using the demographic weights and append the
    # _total rows for active compartments.
    self._prepare_demographic_state()
    return self.population_matrix

equation(y, t, p)

Compute the compartment derivatives with an age-stratified force of infection.

The standard I->R edge is handled by the framework; the S->I force of infection is built manually so it can mix age groups through the contact matrix and regions through the travel matrix, using the per-age beta vector, and is applied via :meth:_apply_flow (which also accumulates I_total). Interventions scale each per-age beta directly.

Parameters:

Name	Type	Description	Default
y	ndarray	Current compartment values, ordered by compartment_list; each row is shaped (age groups, zones).	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 derivatives (dy/dt).

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

def equation(self, y, t, p):
    """Compute the compartment derivatives with an age-stratified force of infection.

    The standard I->R edge is handled by the framework; the S->I force of
    infection is built manually so it can mix age groups through the contact
    matrix and regions through the travel matrix, using the per-age beta
    vector, and is applied via :meth:`_apply_flow` (which also accumulates
    I_total). Interventions scale each per-age beta directly.

    Args:
        y (jnp.ndarray): Current compartment values, ordered by
            ``compartment_list``; each row is shaped (age groups, zones).
        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 derivatives (dy/dt).
    """
    C = self.COMPARTMENTS
    params = self._unpack_params(p)

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

    # Per-region total population (exclude cumulative _total rows).
    non_total = [c for c in self.compartment_list if not c.endswith("_total")]
    N_total = sum(states[c] for c in non_total).sum(axis=0)  # (R,)
    I_frac = I / (N_total[None, :] + 1e-10)                   # (A, R)

    # Per-age transmission rates as a rate dict, so interventions can scale
    # each age band's beta by name (target_rates = the beta_age_* names).
    beta_vec = self._beta_by_age()                           # (A,)
    rates = {name: beta_vec[i] for i, name in enumerate(self._beta_names)}
    prop_infective_scalar = I.sum() / (N_total.sum() + 1e-10)
    rates, travel_matrix = self._apply_interventions(
        t, rates, prop_infective_scalar
    )
    # Reassemble the (post-intervention) per-age beta vector in state order.
    beta_vec = jnp.asarray([rates[name] for name in self._beta_names])  # (A,)
    rates["gamma"] = params["gamma"]

    # Standard edges via the framework (only the I->R gamma edge exists; the
    # extra beta_age_* keys in `rates` are ignored — they aren't edges).
    derivs = self._compute_equations(states, rates)

    # Manual force of infection: spatial mixing then age mixing.
    contact_matrix = (
        self.contact_matrix
        if self.contact_matrix is not None
        else jnp.eye(S.shape[0])
    )
    infectious_pressure = (travel_matrix @ I_frac.T).T        # (A, R) spatial mix
    mixed = contact_matrix @ infectious_pressure              # (A, R) age mix
    flow_foi = S * beta_vec[:, None] * mixed                  # (A, R)

    self._apply_flow(derivs, C.S, C.I, flow_foi)

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