Novel Respiratory (SEIHDR)

An SEIHDR compartmental model for novel respiratory diseases with age-stratified transmission

CovidJaxModel(config)

Bases: Model

SEIHDR compartmental model with age-stratified transmission and spatial mobility.

Initialize the model and adopt the config's compartment list.

The config's compartment_list may be a variant subset, so it replaces the full schema-derived list.

Parameters:

Name	Type	Description	Default
config	dict	Validated simulation configuration, including compartment_list.	required

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

def __init__(self, config):
    """Initialize the model and adopt the config's compartment list.

    The config's ``compartment_list`` may be a variant subset, so it
    replaces the full schema-derived list.

    Args:
        config (dict): Validated simulation configuration, including
            ``compartment_list``.
    """
    super().__init__(config)

    # Override with the compartment list from config (may be a variant subset).
    self.compartment_list = config["compartment_list"]

define_parameters(schema) classmethod

Declare the SEIHDR 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. Declares the six S-E-I-H-D-R compartments, seven transmission edges, the mask-wearing / social-isolation / lockdown / vaccination interventions, the travel_sigma mobility field, and three age bands (0-17, 18-55, 56+) whose age_range tags enable the country's Prem 2021 contact matrix.

Parameters:

Name	Type	Description	Default
schema	ParameterSchemaBuilder	Schema builder to populate with model info, metadata, compartments, transmission edges, interventions, parameters, and demographic groups.	required

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

@classmethod
def define_parameters(cls, schema):
    """Declare the SEIHDR 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. Declares the six
    S-E-I-H-D-R compartments, seven transmission edges, the mask-wearing /
    social-isolation / lockdown / vaccination interventions, the
    ``travel_sigma`` mobility field, and three age bands (0-17, 18-55, 56+)
    whose ``age_range`` tags enable the country's Prem 2021 contact matrix.

    Args:
        schema (ParameterSchemaBuilder): Schema builder to populate with
            model info, metadata, compartments, transmission edges,
            interventions, parameters, and demographic groups.
    """
    schema.set_model_info(
        disease_type=cls.DISEASE_TYPE,
        label=cls.DISEASE_LABEL,
        description=cls.DISEASE_DESCRIPTION,
    )
    schema.set_model_metadata(
        authors=[
            {
                "name": "Daniel Langford",
                "email": "dlangford@ruvos",
                "affiliation": "Ruvos",
            }
        ],
        license="MIT",
        model_type="Compartmental",
        diseases=["Novel respiratory diseases"],
        transmission_routes=["Airborne"],
        questions_answered=[
            "How does transmission rate affect peak incidence and total infections?",
            "How do masks, social isolation, lockdown, and vaccination reduce spread?",
            "How does age structure and contact patterns affect outbreak dynamics across age groups?",
        ],
        key_assumptions=[
            "Closed population — no births or deaths.",
            "Three age bands (0–17, 18–55, 56+) with age-stratified contact rates derived from Prem 2021 synthetic matrices.",
            "Homogeneous mixing within each age group; spatial coupling via a travel matrix.",
            "Transmission is frequency-dependent (force of infection scales with the proportion infectious, not the count).",
            "Vaccination is modeled as a leaky vaccine: adherence represents coverage and transmission reduction represents vaccine efficacy.",
            "Recovered individuals are fully immune with no waning (no R→S transition).",
            "Hospitalization is a fixed fraction of all infected individuals regardless of age group.",
        ],
    )

    # ---- Compartments (S-E-I-H-D-R) ----
    schema.add_compartment(
        "S", "Susceptible", "Population susceptible to infection"
    )
    schema.add_compartment(
        "E", "Exposed", "Population exposed but not yet infectious"
    )
    schema.add_compartment(
        "I",
        "Infected",
        "Currently infected and infectious population",
        infective=True,
    )
    schema.add_compartment(
        "H", "Hospitalized", "Infected individuals requiring hospitalization"
    )
    schema.add_compartment(
        "D", "Deceased", "Population that has died from the disease"
    )
    schema.add_compartment("R", "Recovered", "Recovered and immune population")

    # ---- Transmission edges ----
    schema.add_transmission_parameter(
        source="susceptible",
        target="exposed",
        variable_name="beta",
        frequency_dependent=True,
        label="Transmission Rate (S->E)",
        description="Rate at which susceptible individuals become exposed through contact with infected individuals",
        default=0.25,
        min_value=0.01,
        max_value=2.0,
        default_min=0.2,
        default_max=0.3,
        unit="per day",
    )

    schema.add_transmission_parameter(
        source="exposed",
        target="infected",
        variable_name="theta",
        label="Incubation Period (E->I)",
        description="Average number of days from exposure to becoming infectious",
        default=5.0,
        min_value=1.0,
        max_value=100.0,
        default_min=2.0,
        default_max=14.0,
        unit="days",
        value_type=ValueType.DAYS,
    )

    schema.add_transmission_parameter(
        source="infected",
        target="hospitalized",
        variable_name="zeta",
        label="Hospitalization Rate (I->H)",
        description="Percentage of infected individuals who require hospitalization",
        default=4,
        min_value=1,
        max_value=100,
        default_min=4,
        default_max=40,
        unit="%",
        value_type=ValueType.PERCENTAGE,
    )

    schema.add_transmission_parameter(
        source="infected",
        target="deceased",
        variable_name="delta",
        label="Infection Fatality Rate (I->D)",
        description="Rate at which infected individuals die from the disease",
        default=0.0001,
        min_value=0.00001,
        max_value=0.01,
        default_min=0.0001,
        default_max=0.001,
        unit="per day",
    )

    schema.add_transmission_parameter(
        source="hospitalized",
        target="deceased",
        variable_name="epsilon",
        label="Hospital Fatality Rate (H->D)",
        description="Rate at which hospitalized individuals die",
        default=0.001,
        min_value=0.0001,
        max_value=0.01,
        default_min=0.001,
        default_max=0.005,
        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.14,
        min_value=1.0,
        max_value=100.0,
        default_min=4.0,
        default_max=10.0,
        unit="days",
        value_type=ValueType.DAYS,
    )

    schema.add_transmission_parameter(
        source="hospitalized",
        target="recovered",
        variable_name="eta",
        label="Hospital Recovery Period (H->R)",
        description="Average number of days for a hospitalized individual to recover",
        default=7.14,
        min_value=1.0,
        max_value=100.0,
        default_min=3.0,
        default_max=14.0,
        unit="days",
        value_type=ValueType.DAYS,
    )

    # ---- Interventions ----
    schema.add_intervention(
        id="mask_wearing",
        label="Mask Wearing",
        description="Reduces transmission rate through mask usage in the population",
        target_rates=["beta"],
        adherence=20.0,
        transmission_reduction=35.0,
    )

    schema.add_intervention(
        id="social_isolation",
        label="Social Isolation",
        description="Reduces transmission rate through social distancing and isolation measures",
        target_rates=["beta"],
        adherence=40.0,
        transmission_reduction=50.0,
    )

    schema.add_intervention(
        id="lock_down",
        label="Lockdown",
        description="Severe movement restrictions: halts inter-regional travel and reduces transmission through reduced social contact",
        target_rates=["beta"],
        modifies_travel=True,
        adherence=80.0,
        transmission_reduction=70.0,
    )

    schema.add_intervention(
        id="vaccination",
        label="Vaccination",
        description=(
            "Leaky-vaccine campaign: lowers the effective transmission rate "
            "while active, with adherence as vaccination coverage and "
            "transmission reduction as vaccine efficacy. There are no after "
            "effects. The reduction is applied only while the intervention is "
            "on; transmission returns to its baseline rate the moment it is "
            "turned off, and no immunity is carried forward from the campaign."
        ),
        target_rates=["beta"],
        adherence=60.0,
        transmission_reduction=80.0,
    )

    # ---- Mobility ----
    # Declared as a custom field so it is modeler-owned and per-simulation
    # editable. Consumed by build_travel_matrix() below.
    schema.add_parameter(
        name="travel_sigma",
        label="Travel Rate (σ)",
        description=(
            "Percentage of each region's population away from home on a given "
            "day. Trips are distributed across destinations by an inverse-square "
            "gravity model weighted by population and distance. 0 disables "
            "inter-regional travel. Note that if an admin 2 zone is selected, "
            "no travel takes place regardless of this value."
        ),
        value_type=ValueType.PERCENTAGE,
        default=20.0,
        min_value=0.0,
        max_value=100.0,
        unit="%",
    )

    # ---- Demographics ----
    # age_range enables auto-loading of the country's Prem 2021 contact
    # matrix (aggregated to these bands) when no explicit overrides are
    # declared.  The explicit overrides below still take precedence; the
    # age_range tags are forward-compatible for when the overrides are
    # removed in favor of country-aware defaults.
    schema.add_demographic_group("age_0_17",    "Children (0-17)", default_weight=33.3, age_range=(0, 17))
    schema.add_demographic_group("age_18_55",   "Adults (18-55)",  default_weight=44.4, age_range=(18, 55))
    schema.add_demographic_group("age_56_plus", "Elderly (56+)",   default_weight=22.3, age_range=(56, 120))

build_travel_matrix(admin_zones)

Build an inverse-square gravity travel matrix from travel_sigma.

Rows sum to 1 with the diagonal at 1 - sigma, so the force of infection in equation() conserves each region's population.

Parameters:

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

Returns:

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

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

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

    Rows sum to 1 with the diagonal at ``1 - sigma``, so the force of
    infection in ``equation()`` conserves each region's population.

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

    Returns:
        np.ndarray: The ``(n_zones, n_zones)`` travel matrix.
    """
    sigma = self._to_rate(self.travel_sigma, ValueType.PERCENTAGE)
    return get_gravity_model_travel_matrix(admin_zones, sigma)

prepare_initial_state()

Return the age-stratified initial populations for the solver.

Returns:

Type	Description
ndarray	Initial populations shaped (compartments, age groups, zones), including the _total rows appended for the active compartments.

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

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

    Returns:
        np.ndarray: Initial populations shaped
            ``(compartments, age groups, zones)``, including the ``_total``
            rows appended for the active compartments.
    """
    # Expand (K, R) → (K, A, R) and append _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.

Uses the conditional pattern: _compute_equations() handles the six standard edges (E->I, I->H, I->D, H->D, I->R, H->R), while the S->E force of infection mixes zones through the travel matrix and age groups through the contact matrix and is applied manually via _apply_flow(). Deceased and _total compartments are excluded from the population that scales the frequency-dependent force of infection.

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 derivatives (dy/dt).

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

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

    Uses the conditional pattern: ``_compute_equations()`` handles the six
    standard edges (E->I, I->H, I->D, H->D, I->R, H->R), while the S->E
    force of infection mixes zones through the travel matrix and age groups
    through the contact matrix and is applied manually via
    ``_apply_flow()``. Deceased and ``_total`` compartments are excluded
    from the population that scales the frequency-dependent force of infection.

    Args:
        y (np.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:
        np.ndarray: Stacked per-compartment derivatives (dy/dt).
    """
    C = self.COMPARTMENTS
    params = self._unpack_params(p)

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

    # Population for FOI (exclude deaths and _total compartments)
    non_total = [c for c in self.compartment_list if not c.endswith("_total")]
    pop_terms = [states[c] for c in non_total if c != "D"]
    N_total = sum(pop_terms).sum(axis=0)
    I_frac = I / N_total[None, :]

    # --- Interventions (schema-driven) ---
    rates = {"beta": params["beta"]}
    prop_infective_scalar = I.sum() / N_total.sum()
    rates, travel_matrix = self._apply_interventions(
        t, rates, prop_infective_scalar
    )

    # Add all active (non-None) rates for standard edges
    for name, value in params.items():
        if name != "beta" and value is not None:
            rates[name] = value

    # --- Standard edges via framework (skip manual FOI edge) ---
    derivs = self._compute_equations(states, rates, skip_edges={"beta"})

    # --- Manual FOI: spatial travel mixing + demographic contact matrix ---
    BETA = ((rates["beta"] * travel_matrix) @ I_frac.T).T
    omega = self.contact_matrix @ BETA
    flow_foi = S * omega

    # Route FOI to E if present, otherwise directly to I
    target = "E" if "E" in states else "I"
    self._apply_flow(derivs, "S", target, flow_foi)

    return np.stack([derivs[comp] for comp in self.compartment_list])