stateMINT API

Everything you call in stateMINT is one classmethod and one dataclass, and you install mintstate but import stateMINT, capital MINT. The mismatch between the two names is the most common source of confusion in the package.

Mamba2Regressor.from_pretrained

Mamba2Regressor.from_pretrained(
    path_or_repo_id: str,
    predictor: Literal["prevalence", "cases"],
    *,
    revision: str | None = None,
    cache_dir: str | None = None,
    local_dir: str | None = None,
) -> ModelArtifact

Load a trained checkpoint from a Hugging Face repository or a local folder.

Parameter Meaning
path_or_repo_id Hugging Face repository id, "dide-ic/stateMINT", or a path to a local folder holding the same layout.
predictor Which checkpoint to load, either "prevalence" or "cases". The two are separately trained and have separate scalers. Positional or keyword.
revision Git revision of the repository, a tag, a branch or a commit. Tags run from v1.0.0 to v1.2.2. Defaults to main, which moves, so pin it.
cache_dir Override the Hugging Face cache location for this call. Prefer the HF_HOME environment variable, which applies to everything.
local_dir Load from a folder on disk instead of the Hub. Nothing is downloaded and revision is ignored.

Returns a ModelArtifact, not a Mamba2Regressor. The artifact carries the fitted scaler and the preprocessing recipe alongside the weights.

from stateMINT.model import Mamba2Regressor

artifact = Mamba2Regressor.from_pretrained(
    "dide-ic/stateMINT", predictor="prevalence",
)

print(type(artifact).__name__)
ModelArtifact

ModelArtifact

You call .predict on the artifact, never on .model directly.

Field Type What it holds
.model nnx.Module The trained Mamba2 network.
.model_config dict Architecture and the predictor it was trained for, including d_model, n_layers, d_state, input_size and dropout.
.preprocessing_config dict The input recipe, holding static_covars, after_intervention, n_steps, window_size, model_start_day, intervention_day and the scaler statistics.
.scaler StandardScaler Fitted mean and scale for the twelve covariates, as .mean_ and .scale_. Specific to this checkpoint. 
print(artifact.model_config["d_model"], "hidden units,",
      artifact.model_config["n_layers"], "layers")
print(artifact.preprocessing_config["n_steps"], "steps of",
      artifact.preprocessing_config["window_size"], "days")
print(artifact.scaler.mean_.shape, "scaler means")
256 hidden units, 2 layers
157 steps of 14 days
(12,) scaler means

The two configs are the reference for the covariate names and the time grid set out in Covariates and the time grid. Read the grid out of them, never hard-code it.

.predict

artifact.predict(
    static_covars: list[dict[str, float]],
    *,
    transformed: bool = False,
) -> np.ndarray            # (B, 157), float32

Predict the checkpoint’s six-year trajectory for each scenario in the list. A single scenario is a list of one.

Parameter Meaning
static_covars One dict per scenario, keyed by the twelve covariate names. All twelve are mandatory, and a missing key raises ValueError: Missing static covariates: [...].
transformed False (default) returns the quantity itself, prevalence in [0, 1] and cases as counts. True returns the space the network was trained in, logit for prevalence and log1p for cases.

Returns a float32 array of shape (B, 157), with rows in the order you supplied them.

Case counts may come back a fraction below zero, because expm1 of a slightly negative prediction is slightly negative. Floor them with np.maximum(y, 0.0). Prevalence needs no such treatment.

district = {
    "eir": 31.96, "dn0_use": 0.311, "dn0_future": 0.485,
    "Q0": 0.85, "phi_bednets": 0.80, "seasonal": 0,
    "routine": 0.0, "itn_use": 0.70, "irs_use": 0.0,
    "itn_future": 0.70, "irs_future": 0.0, "lsm": 0.0,
}

y = artifact.predict([district, dict(district, itn_future=0.0, dn0_future=0.0)])
print(y.shape, y.dtype)
print(f"prevalence after three years: {y[0, -1]:.3f} sustained, {y[1, -1]:.3f} withdrawn")
(2, 157) float32
prevalence after three years: 0.470 sustained, 0.498 withdrawn

.prepare_inputs

artifact.prepare_inputs(
    static_covars: list[dict[str, float]],
) -> np.ndarray            # (B, 157, 16), float32

.prepare_inputs runs the first half of .predict and stops before the network. The covariates are broadcast across the 157 windows, the five campaign covariates are masked to zero before the intervention day, the twelve are scaled, and four time features are appended to give the sixteen channels the network consumes.

Returns a float32 array of shape (B, n_steps, input_size), which is (B, 157, 16) for both current checkpoints.

import numpy as np

X = artifact.prepare_inputs([district, dict(district, lsm=0.9)])
changed = ~np.isclose(X[0], X[1]).all(axis=1)

print(X.shape)
print(f"steps affected by lsm: {int(changed.sum())} of {X.shape[1]}, from step {int(np.argmax(changed))}")
(2, 157, 16)
steps affected by lsm: 78 of 157, from step 79

See also

Calling the emulator gives the reasoning behind these arguments, and what goes wrong when the wrong checkpoint is read. Running scenarios wraps this surface into a single call.