Base Model Reference

Base class for all compartmental disease models. Subclasses must implement define_parameters(schema). They should also override prepare_initial_state() and equation(), the per-step function the solver calls.

Model(config)

Bases: ABC

Base class for all compartmental disease models.

Subclasses must implement: - define_parameters(schema) -- fill the builder with parameter declarations

For models that have migrated to declarative parameters, the following are derived automatically from the schema:

• disease_type property
• COMPARTMENT_LIST class attribute
• get_params() — returns transmission params in schema edge order
• Transmission param attributes (e.g. self.beta, self.gamma)

Subclasses should override: - prepare_initial_state() - equation() - build_travel_matrix(admin_zones) — only if the model has inter-zone mobility. The default is the identity matrix.

Non-migrated models can still override disease_type, define COMPARTMENT_LIST manually, and implement get_params().

Initialize common model state from config.

For migrated models (those with a cached parameter schema), calling super().__init__(config) extracts the common fields every model needs:

• population_matrix, compartment_list
• start_date, start_date_ordinal, n_timesteps
• admin_units, payload
• Transmission params (self.beta, self.gamma, etc.)
• disease_params — :class:DiseaseParamValues with attribute access (e.g. self.disease_params.immunity_period)
• Top-level aliases for each disease parameter (e.g. self.immunity_period) for parameter uncertainty compatibility
• intervention_dict (kept for post-processing compat)
• interventions — list of :class:Intervention runtime objects
• intervention_statuses — auto-generated from schema

Non-migrated models that don't call super().__init__() are unaffected — they continue to extract fields manually.

Subclasses override __init__ and call super().__init__(config) first, then add model-specific fields (travel matrix, demographics, etc.).

Source code in compartment/model.py

def __init__(self, config):
    """
    Initialize common model state from config.

    For **migrated models** (those with a cached parameter schema),
    calling ``super().__init__(config)`` extracts the common fields
    every model needs:

    - ``population_matrix``, ``compartment_list``
    - ``start_date``, ``start_date_ordinal``, ``n_timesteps``
    - ``admin_units``, ``payload``
    - Transmission params (``self.beta``, ``self.gamma``, etc.)
    - ``disease_params`` — :class:`DiseaseParamValues` with
      attribute access (e.g. ``self.disease_params.immunity_period``)
    - Top-level aliases for each disease parameter
      (e.g. ``self.immunity_period``) for parameter uncertainty compatibility
    - ``intervention_dict`` (kept for post-processing compat)
    - ``interventions`` — list of :class:`Intervention` runtime objects
    - ``intervention_statuses`` — auto-generated from schema

    **Non-migrated models** that don't call ``super().__init__()``
    are unaffected — they continue to extract fields manually.

    Subclasses override ``__init__`` and call ``super().__init__(config)``
    first, then add model-specific fields (travel matrix, demographics,
    etc.).
    """
    self.config = config
    self.contact_matrix = None
    self._rate_vectors = None

    schema = type(self)._get_cached_schema()
    if not schema:
        # Non-migrated model — subclass handles everything.
        return

    # -- Common fields (identical across all migrated models) --
    self.population_matrix = jnp.array(config["initial_population"]).T
    self.compartment_list = list(self.COMPARTMENTS)
    self.start_date = config["start_date"]
    self.start_date_ordinal = self.start_date.toordinal()
    self.n_timesteps = config["time_steps"]
    self.admin_units = config["admin_units"]
    self.payload = config

    # Transmission params: self.beta, self.gamma, etc.
    self._load_transmission_params(config.get("transmission_dict", {}))

    # -- Disease parameters --
    # Auto-unpack from the Disease dict using schema declarations.
    # Values are accessible via self.disease_params.<name> and also
    # as top-level aliases (self.<name>) for parameter uncertainty compatibility
    # (BatchSimulationManager._single_run uses setattr(model, key, val)).
    disease_dict = config.get("Disease", {}) or {}
    self.disease_params = DiseaseParamValues(
        schema.disease_parameters, disease_dict
    )
    for p in schema.disease_parameters:
        setattr(self, p.name, getattr(self.disease_params, p.name))

    # -- Interventions --
    self.intervention_dict = config.get("intervention_dict", {})

    # Build runtime Intervention objects from schema defs + config
    self.interventions: list[Intervention] = []
    for intv_def in schema.interventions:
        if intv_def.id in self.intervention_dict:
            self.interventions.append(
                Intervention.from_config(
                    intv_def, self.intervention_dict[intv_def.id]
                )
            )

    # Auto-generate statuses from schema intervention declarations
    self.intervention_statuses = {
        intv_def.id: False for intv_def in schema.interventions
    }

    # Demographics: prefer config-supplied weights; otherwise fall back to
    # the schema's declared default_weight values for each demographic group.
    # Downstream consumers (post-processor, output formatter) read
    # ``self.demographics`` to derive age-group labels and weights, so it
    # must always be populated when the schema declares groups.
    config_demo = {}
    case_file = config.get("case_file")
    if isinstance(case_file, dict):
        config_demo = case_file.get("demographics") or {}
    if config_demo:
        self.demographics = dict(config_demo)
    elif schema.demographic_groups:
        self.demographics = {
            g.id: g.default_weight for g in schema.demographic_groups
        }
    else:
        self.demographics = {}

    self.contact_matrix = self._build_contact_matrix(config)
    self._rate_vectors = self._build_rate_vectors(config)

disease_type property

Disease type identifier (e.g. "MPOX").

Migrated models get this from the schema automatically. Non-migrated models should override this property.

__init_subclass__(**kwargs)

Auto-populate schema-derived class attributes when a subclass is defined. For migrated models this sets:

• COMPARTMENTS — :class:CompartmentRegistry with attribute access (e.g. cls.COMPARTMENTS.S).
• COMPARTMENT_LIST — plain list[str] for backward compatibility with the post-processor and other legacy code.

After building the schema the framework auto-generates _total cumulative tracking compartments for every compartment that is the target of at least one transmission edge. These are appended to the schema's compartment list before the :class:CompartmentRegistry is created, so they are accessible via cls.COMPARTMENTS and included in cls.COMPARTMENT_LIST.

Non-migrated models that raise NotImplementedError from define_parameters() are silently skipped.

Source code in compartment/model.py

def __init_subclass__(cls, **kwargs):
    """
    Auto-populate schema-derived class attributes when a subclass is
    defined.  For migrated models this sets:

    - ``COMPARTMENTS`` — :class:`CompartmentRegistry` with attribute
      access (e.g. ``cls.COMPARTMENTS.S``).
    - ``COMPARTMENT_LIST`` — plain ``list[str]`` for backward
      compatibility with the post-processor and other legacy code.

    After building the schema the framework auto-generates
    ``_total`` cumulative tracking compartments for every
    compartment that is the target of at least one transmission
    edge.  These are appended to the schema's compartment list
    before the :class:`CompartmentRegistry` is created, so they
    are accessible via ``cls.COMPARTMENTS`` and included in
    ``cls.COMPARTMENT_LIST``.

    Non-migrated models that raise ``NotImplementedError`` from
    ``define_parameters()`` are silently skipped.
    """
    super().__init_subclass__(**kwargs)

    try:
        schema = cls._build_parameter_schema()
        cls._cached_schema = schema

        # Derive the display order for the results sidebar from
        # COMPARTMENT_DELTA_GROUPING (if the model defines one) or
        # fall back to the raw compartment IDs.  Must run before
        # _add_total_compartments so that _total entries are excluded.
        grouping = getattr(cls, "COMPARTMENT_DELTA_GROUPING", None)
        if grouping:
            schema.compartment_display_order = list(grouping.keys())
        else:
            schema.compartment_display_order = [c.id for c in schema.compartments]

        # Auto-generate _total compartments for each edge target.
        # These are a framework concern (derivative/equation accumulation,
        # post-processing) — model authors never declare them.
        cls._add_total_compartments(schema)

        cls.COMPARTMENTS = CompartmentRegistry(schema.compartments)
        # Keep COMPARTMENT_LIST as a plain list for backward compat
        # with the post-processor and other downstream code.
        cls.COMPARTMENT_LIST = list(cls.COMPARTMENTS)
        # Auto-set DISEASE_TYPE from define_parameters() so models
        # only declare disease_type once (in set_model_info).
        if "DISEASE_TYPE" not in cls.__dict__ and schema.disease_type:
            cls.DISEASE_TYPE = schema.disease_type
        cls.MODEL_KEY = schema.model_key

        # Custom field registries: attribute-style access to
        # disease parameter and admin zone field names.
        cls.DISEASE_PARAMS = DiseaseParamRegistry(schema.disease_parameters)
        cls.ADMIN_ZONE_FIELDS = AdminZoneFieldRegistry(schema.admin_zone_fields)
        # Derive NUM_RUNS from the schema so run_simulation can read it
        # without importing the parameters module.
        cls.NUM_RUNS = schema.num_runs
    except (NotImplementedError, Exception):
        # Non-migrated model — leave everything as-is.
        pass

define_parameters(schema) classmethod

Declare all parameters for this disease model.

Called by the framework with a ParameterSchemaBuilder that has been pre-populated with shared simulation parameters. Subclasses call schema.add_*() / schema.set_*() to register compartments, edges, interventions, custom fields, etc.

Parameters:

Name	Type	Description	Default
schema	ParameterSchemaBuilder	Builder instance -- call its methods to declare parameters.	required

Raises:

Type	Description
NotImplementedError	For non-migrated models that haven't adopted declarative parameter definitions yet.

Source code in compartment/model.py

@classmethod
def define_parameters(cls, schema: ParameterSchemaBuilder) -> None:
    """
    Declare all parameters for this disease model.

    Called by the framework with a ``ParameterSchemaBuilder`` that has
    been pre-populated with shared simulation parameters.  Subclasses
    call ``schema.add_*()`` / ``schema.set_*()`` to register
    compartments, edges, interventions, custom fields, etc.

    Args:
        schema: Builder instance -- call its methods to declare parameters.

    Raises:
        NotImplementedError: For non-migrated models that haven't
            adopted declarative parameter definitions yet.
    """
    raise NotImplementedError(
        f"{cls.__name__} has not implemented define_parameters(). "
        "Implement it to enable declarative parameter definitions."
    )

generate_artifact() classmethod

Generate the model artifact JSON dict.

This is the schema document consumed by the UI layer, DB seeding, and downstream Zod schema generation.

Source code in compartment/model.py

@classmethod
def generate_artifact(cls) -> dict[str, Any]:
    """
    Generate the model artifact JSON dict.

    This is the schema document consumed by the UI layer, DB seeding,
    and downstream Zod schema generation.
    """
    return cls._build_parameter_schema().to_artifact_dict()

generate_example_config(uncertainty=False) classmethod

Generate an example simulation config from parameter defaults.

The output is in the "short-form" JSON format accepted by load_config_from_json(). Pass uncertainty=True to enable every variance-capable parameter that declares default uncertainty bounds.

Source code in compartment/model.py

@classmethod
def generate_example_config(
    cls, uncertainty: bool = False
) -> dict[str, Any]:
    """
    Generate an example simulation config from parameter defaults.

    The output is in the "short-form" JSON format accepted by
    ``load_config_from_json()``. Pass ``uncertainty=True`` to enable every
    variance-capable parameter that declares default uncertainty bounds.
    """
    return cls._build_parameter_schema().to_example_config(
        uncertainty=uncertainty
    )

build_overridden_config(params)

Deep-copy this model's source config and apply LHS parameter uncertainty overrides, returning the new config.

This is the basis for parameter uncertainty sampling: instead of mutating a constructed model (whose __init__ may have already baked params into derived constants like 1/latent_period), the runner rebuilds the model from an overridden config so every value declared via add_transmission_parameter / add_parameter / add_intervention is re-derived from scratch.

Override keys are routed by kind, using the schema to disambiguate:

• "intervention.<id>.<field>" → intervention_dict[id][field] (values are fractions, matching create_intervention_dict).
• a transmission-edge variable_name → transmission_dict[name] in native units, so _to_rate is applied on rebuild (correct for RATE, DAYS, and PERCENTAGE edges alike).
• any other single-token key → the Disease dict (disease parameters), again in native units.

Parameters:

Name	Type	Description	Default
params	dict[str, Any]	{key: sampled_value} from generate_LHS_samples.	required

Returns:

Type	Description
	A deep-copied config object with overrides applied. Falls back to
	returning the original config untouched if the model exposes no
	source config (non-reconstructable model).

Source code in compartment/model.py

def build_overridden_config(self, params: dict[str, Any]):
    """Deep-copy this model's source config and apply LHS parameter uncertainty
    overrides, returning the new config.

    This is the basis for parameter uncertainty sampling: instead of mutating a
    constructed model (whose ``__init__`` may have already baked params
    into derived constants like ``1/latent_period``), the runner rebuilds
    the model from an overridden config so every value declared via
    ``add_transmission_parameter`` / ``add_parameter`` /
    ``add_intervention`` is re-derived from scratch.

    Override keys are routed by kind, using the schema to disambiguate:

    - ``"intervention.<id>.<field>"`` → ``intervention_dict[id][field]``
      (values are fractions, matching ``create_intervention_dict``).
    - a transmission-edge ``variable_name`` → ``transmission_dict[name]``
      in **native units**, so ``_to_rate`` is applied on rebuild
      (correct for RATE, DAYS, and PERCENTAGE edges alike).
    - any other single-token key → the ``Disease`` dict (disease
      parameters), again in native units.

    Args:
        params: ``{key: sampled_value}`` from ``generate_LHS_samples``.

    Returns:
        A deep-copied config object with overrides applied. Falls back to
        returning the original config untouched if the model exposes no
        source config (non-reconstructable model).
    """
    import copy

    source = self._source_config()
    if source is None:
        return source

    config = copy.deepcopy(source)
    schema = type(self)._get_cached_schema()
    edge_vars = (
        {e.variable_name for e in schema.transmission_edges} if schema else set()
    )

    def _disease_dict():
        # config.config is the validated config dict (holds "Disease");
        # plain-dict configs hold "Disease" directly.
        inner = getattr(config, "config", config)
        return inner.setdefault("Disease", {})

    for key, val in params.items():
        parts = key.split(".")
        if len(parts) == 3 and parts[0] == "intervention":
            _, intervention_id, field_name = parts
            config.intervention_dict.setdefault(intervention_id, {})[field_name] = (
                val
            )
        elif len(parts) == 1:
            name = parts[0]
            if name in edge_vars:
                config.transmission_dict[name] = val
            else:
                _disease_dict()[name] = val
        else:
            raise ValueError(f"Invalid uncertainty parameter key: {key}")

    return config

get_params()

Return transmission parameters as a tuple in schema edge order.

For migrated models this is derived automatically from the parameter schema — no need to override. The tuple order matches the add_transmission_parameter() call order in define_parameters(), which must match the positional unpack in equation().

Non-migrated models should override this method.

Source code in compartment/model.py

def get_params(self):
    """
    Return transmission parameters as a tuple in schema edge order.

    For migrated models this is derived automatically from the
    parameter schema — no need to override.  The tuple order matches
    the ``add_transmission_parameter()`` call order in
    ``define_parameters()``, which must match the positional unpack
    in ``equation()``.

    Non-migrated models should override this method.
    """
    schema = type(self)._get_cached_schema()
    if not schema:
        return None
    return tuple(
        getattr(self, edge.variable_name) for edge in schema.transmission_edges
    )

get_initial_population(admin_zones, compartment_list, **kwargs) classmethod

Compute initial population distribution for this disease model.

Default implementation for simple S-I models. Override in subclasses for disease-specific initialization logic.

Parameters:

Name	Type	Description	Default
admin_zones		List of admin zone dicts with population data	required
compartment_list		List of compartment names	required
**kwargs		Model-specific parameters (run_mode, etc.)	{}

Returns:

Type	Description
	numpy array of shape (n_zones, n_compartments)

Source code in compartment/model.py

@classmethod
def get_initial_population(cls, admin_zones, compartment_list, **kwargs):
    """
    Compute initial population distribution for this disease model.

    Default implementation for simple S-I models.
    Override in subclasses for disease-specific initialization logic.

    Args:
        admin_zones: List of admin zone dicts with population data
        compartment_list: List of compartment names
        **kwargs: Model-specific parameters (run_mode, etc.)

    Returns:
        numpy array of shape (n_zones, n_compartments)
    """
    column_mapping = {value: index for index, value in enumerate(compartment_list)}
    initial_population = np.zeros((len(admin_zones), len(compartment_list)))

    for i, zone in enumerate(admin_zones):
        infected = round(zone["infected_population"] / 100 * zone["population"], 2)
        susceptible = zone["population"] - infected
        initial_population[i, column_mapping["S"]] = susceptible
        initial_population[i, column_mapping["I"]] = infected

    return initial_population

build_travel_matrix(admin_zones)

Return this model's (R, R) inter-zone travel matrix.

T[i, j] is the fraction of zone i's population present in zone j on a given day. Rows must sum to 1, and the diagonal is the stay-home fraction (1 - sigma for a model with an outbound travel rate sigma).

The default is the identity matrix — no inter-zone travel. Override this to define the model's mobility, reading whatever parameters define_parameters() declared. Declare them with schema.add_parameter() (by convention travel_sigma, plus any kernel-specific extras such as travel_alpha or travel_scale_km) so they surface as custom fields in the UI.

PERCENTAGE parameters arrive in native units (20.0, not 0.2) because only transmission edges are auto-converted. Convert with self._to_rate(self.travel_sigma, ValueType.PERCENTAGE).

Parameters:

Name	Type	Description	Default
admin_zones		The case file's admin-zone dicts, in population-matrix column order. Each carries center_lat, center_lon and population.	required

Source code in compartment/model.py

def build_travel_matrix(self, admin_zones):
    """
    Return this model's ``(R, R)`` inter-zone travel matrix.

    ``T[i, j]`` is the fraction of zone *i*'s population present in
    zone *j* on a given day.  Rows must sum to 1, and the diagonal is
    the stay-home fraction (``1 - sigma`` for a model with an outbound
    travel rate ``sigma``).

    The default is the identity matrix — no inter-zone travel.
    Override this to define the model's mobility, reading whatever
    parameters ``define_parameters()`` declared.  Declare them with
    ``schema.add_parameter()`` (by convention ``travel_sigma``,
    plus any kernel-specific extras such as ``travel_alpha`` or
    ``travel_scale_km``) so they surface as custom fields in the UI.

    ``PERCENTAGE`` parameters arrive in native units (``20.0``, not
    ``0.2``) because only transmission edges are auto-converted.
    Convert with ``self._to_rate(self.travel_sigma, ValueType.PERCENTAGE)``.

    Args:
        admin_zones: The case file's admin-zone dicts, in
            population-matrix column order.  Each carries
            ``center_lat``, ``center_lon`` and ``population``.
    """
    return np.eye(len(admin_zones))

dataset(name) classmethod

Absolute path to a dataset this model declares in its datasets.yaml.

Read the file however suits it — the framework only resolves the path, so any format works::

df = pd.read_csv(self.dataset("kenya-contact-matrix"))
zones = json.loads(self.dataset("kenya-admin-zones").read_text())

The same call works locally and in the cloud. Locally the file arrives via python -m compartment.datasets pull --dest data/; in the cloud the release pipeline bakes it into the model image at the same relative path. Never build the path by hand — a literal Path(__file__).parent / "data" / ... resolves to nothing once the version in datasets.yaml is bumped.

Parameters:

Name	Type	Description	Default
name	str	The dataset's name: in datasets.yaml.	required

Raises:

Type	Description
ManifestError	No datasets.yaml, name not declared in it, or the file is not on disk — each with the fix.

Source code in compartment/model.py

@classmethod
def dataset(cls, name: str) -> Path:
    """
    Absolute path to a dataset this model declares in its
    ``datasets.yaml``.

    Read the file however suits it — the framework only resolves the
    path, so any format works::

        df = pd.read_csv(self.dataset("kenya-contact-matrix"))
        zones = json.loads(self.dataset("kenya-admin-zones").read_text())

    The same call works locally and in the cloud. Locally the file
    arrives via ``python -m compartment.datasets pull --dest data/``;
    in the cloud the release pipeline bakes it into the model image at
    the same relative path.  Never build the path by hand — a literal
    ``Path(__file__).parent / "data" / ...`` resolves to nothing once
    the version in ``datasets.yaml`` is bumped.

    Args:
        name: The dataset's ``name:`` in ``datasets.yaml``.

    Raises:
        ManifestError: No ``datasets.yaml``, ``name`` not declared in
            it, or the file is not on disk — each with the fix.
    """
    # Imported here rather than at module scope: compartment.model is
    # imported by every model, and this keeps yaml parsing off that path
    # for the models that use no datasets.
    from compartment.datasets import dataset_path, model_dir_for_class

    model_dir = model_dir_for_class(cls)
    if model_dir is None:
        from compartment.datasets import ManifestError

        raise ManifestError(
            f"Cannot locate the source directory for {cls.__name__}, so its "
            f"datasets.yaml cannot be read. This happens when the class is "
            f"defined dynamically rather than in a model.py on disk."
        )
    return dataset_path(model_dir, name)

equation(y, t, p)

Compute the per-step state change at time t.

This is the function the solver calls each step. For continuous ODE models it returns dy/dt; for discrete or stochastic Euler-stepped models it returns the per-step increment.

Return an array shaped like y, typically jnp.stack([derivs[c] for c in self.compartment_list]).

Source code in compartment/model.py

def equation(self, y, t, p):
    """
    Compute the per-step state change at time ``t``.

    This is the function the solver calls each step. For continuous
    ODE models it returns ``dy/dt``; for discrete or stochastic
    Euler-stepped models it returns the per-step increment.

    Return an array shaped like ``y``, typically
    ``jnp.stack([derivs[c] for c in self.compartment_list])``.
    """
    raise NotImplementedError(f"{type(self).__name__} must implement equation().")