mintstate 0.1.1

StateMINT is a state space based neural network emulator for malariasimulation.

StateMINT

StateMINT is a JAX/Flax neural emulator for malariasimulation outputs. It uses a Mamba2 state-space sequence model to predict malaria trajectories from static scenario covariates and intervention timing features. It supersedes the earlier MINTelligence RNN emulator.

Contents

• Features
• Installation
  • From PyPI
  • From GitHub
  • For development
• Quick Start: Inference
• Static Covariates
• Training Workflow
• Configuration
• Development
• Repository Layout
• Contributing

Features

• Mamba2-based sequence regressors for malaria prevalence and case-count trajectories.
• A lightweight inference path: load an exported artifact and call predict with only the runtime dependencies installed.
• Data extraction utilities for aggregating raw malariasimulation DuckDB outputs into model-ready parquet files.
• Preprocessing with target transforms, covariate scaling, and intervention-aware feature construction.
• Training, evaluation, visualization, checkpointing, and export workflows.
• Hugging Face Hub loading utilities for exported inference artifacts.

Installation

StateMINT requires Python 3.12 or newer.

Naming: the package installs as mintstate (the name statemint was already taken on PyPI) but is imported as stateMINT — e.g. pip install mintstate then from stateMINT.model import Mamba2Regressor.

From PyPI (recommended)

The base install contains only what is needed for inference — loading an exported artifact and predicting:

pip install mintstate

Optional extras add heavier, task-specific dependencies. Install only what you need:

Extra	Command	Adds
(base)	pip install mintstate	Inference: from_pretrained, predict
gpu	pip install "mintstate[gpu]"	CUDA 12 JAX wheels (Linux x86_64)
train	pip install "mintstate[train]"	Data fetch, training, export, and sweeps
plot	pip install "mintstate[plot]"	Prediction/target visualization
all	pip install "mintstate[all]"	Everything above

Extras can be combined, e.g. pip install "mintstate[train,plot]".

From GitHub (latest, unreleased)

To install the latest commit (or a specific branch/tag) straight from source without cloning:

# latest on the default branch
pip install "git+https://github.com/mrc-ide/stateMINT.git"

# a specific branch or tag, with an extra
pip install "mintstate[gpu] @ git+https://github.com/mrc-ide/stateMINT.git@v1.0.0"

Building from GitHub requires Python 3.12+. If pip resolves to an older interpreter you will see No matching distribution found; invoke it explicitly as python3.12 -m pip ....

For development (from source)

The development workflow uses uv:

git clone https://github.com/mrc-ide/stateMINT.git
cd stateMINT
uv sync --all-extras --dev

Or install a subset of extras:

uv sync --extra plot
uv sync --extra gpu

Quick Start: Inference

Load an exported artifact from the Hugging Face Hub or a local directory with Mamba2Regressor.from_pretrained.

from stateMINT.model import Mamba2Regressor

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

static_covars = [{
    "eir": 50.0,
    "dn0_use": 0.3,
    "dn0_future": 0.4,
    "Q0": 0.8,
    "phi_bednets": 0.7,
    "seasonal": 1.0,
    "routine": 0.5,
    "itn_use": 0.2,
    "irs_use": 0.1,
    "itn_future": 0.3,
    "irs_future": 0.2,
    "lsm": 0.0,
}]

predicted_prevalence = artifact.predict(static_covars)

print(predicted_prevalence.shape)  # (batch, timesteps)
print(predicted_prevalence[0])     # first trajectory

For cases, load the cases artifact and use the same input format:

artifact = Mamba2Regressor.from_pretrained(
    "dide-ic/stateMINT",
    predictor="cases",
    revision="v1.0.0",
)

predicted_cases = artifact.predict(static_covars)

By default, predictions are returned on the original target scale: prevalence as probabilities and cases on the scale used by the training data. Pass transformed=True to return model-space outputs.

raw_model_space = artifact.predict(static_covars, transformed=True)

For local artifacts, pass the target artifact directory:

artifact = Mamba2Regressor.from_pretrained(
    "artifacts/prevalence",
    predictor="prevalence",
)

Static Covariates

Inference inputs need one dictionary per scenario with these static covariates:

eir
dn0_use
dn0_future
Q0
phi_bednets
seasonal
routine
itn_use
irs_use
itn_future
irs_future
lsm

Artifacts include the fitted static scaler, timestep grid, intervention day, target transform, and other preprocessing metadata needed for inference.

Training Workflow

The commands below require the train (and optionally plot) extras, or a development install: pip install "mintstate[train,plot]" or uv sync --all-extras --dev.

Typical workflow:

1. Fetch and aggregate simulation data from DuckDB.
2. Train a target-specific model.
3. Evaluate or visualize test-set predictions.
4. Export the checkpoint into a portable inference artifact.
5. Upload the artifact to the Hugging Face Hub, if needed.

1. Fetch Filtered Data

stateMINT.filter_raw_data reads raw DuckDB simulation rows, filters burn-in, aggregates fixed windows, and writes filtered_data_<predictor>.parquet.

uv run python -m stateMINT.filter_raw_data \
  --db-path /path/to/simulations.duckdb \
  --table-name simulation_results \
  --predictor prevalence \
  --window-size 14 \
  --output-folder data

Useful options:

• --predictor prevalence or --predictor cases
• --param-limit N to keep only the first N parameter indices
• --sim-limit N to sample up to N simulations per parameter

The raw table should include identifiers (parameter_index, simulation_index, global_index), daily timesteps, the static covariates above, and output columns for prevalence or cases.

2. Train a Model

Training uses Hydra; the default target is prevalence.

uv run python -m stateMINT.train

Train the cases model:

uv run python -m stateMINT.train target=cases

Common overrides:

uv run python -m stateMINT.train \
  target=prevalence \
  data_file=data/filtered_data_prevalence.parquet \
  output_dir=train_outputs/prevalence \
  use_wandb=false

Training writes checkpoints under checkpoint_dir, saves static_scaler.pkl in output_dir, and reuses a split assignment file for consistent train/validation/test splits.

3. W&B Sweeps

Sweep definitions live in stateMINT/conf/sweeps. Create a sweep, then run one or more agents with the sweep ID returned by W&B:

uv run wandb sweep stateMINT/conf/sweeps/prevalence.yaml
uv run wandb agent <entity>/stateMINT-sweep/<sweep-id>

Use stateMINT/conf/sweeps/cases.yaml for the cases target. Sweep commands set use_wandb=true and pass Hydra overrides through ${args_no_hyphens}.

4. Visualize Predictions

Compare predictions with test-set targets. checkpoint_dir is required.

uv run python -m stateMINT.visualise_predictions \
  target=prevalence \
  checkpoint_dir=train_outputs/prevalence/ckpts-YYYY-MM-DDTHH:MM:SS \
  data_file=data/filtered_data_prevalence.parquet

The default output path is viz_outputs/<predictor>/preds-vs-targets.pdf.

5. Export an Artifact

Export converts a trained Orbax checkpoint and preprocessing metadata into a self-contained artifact.

uv run python -m stateMINT.model_export \
  predictor=prevalence \
  checkpoint_dir=train_outputs/prevalence/ckpts-YYYY-MM-DDTHH:MM:SS \
  scaler_file=train_outputs/prevalence/static_scaler.pkl \
  artifact_dir=artifacts/prevalence

Export config architecture values must match the checkpoint, including d_model, d_state, n_layers, dropout, and related Mamba2 settings.

An exported artifact contains:

artifact_dir/
|-- checkpoint/
|-- model_config.json
`-- preprocessing_config.json

model_config.json stores architecture settings; preprocessing_config.json stores feature order, target transform, intervention timing, timestep construction, and static scaler parameters.

6. Upload To Hugging Face

Authenticate first:

hf auth login

Upload an artifact:

hf upload dide-ic/stateMINT artifacts/prevalence prevalence/ \
  --commit-message "Add prevalence model artifact"

Create a release tag:

hf repos tag create dide-ic/stateMINT v1.0.0 \
  --revision main \
  --message "Release v1.0.0"

Configuration

Main Hydra configs in stateMINT/conf:

• train_config.yaml for training.
• viz_config.yaml for prediction visualizations.
• export_config.yaml for artifact export.
• target/prevalence.yaml and target/cases.yaml for target-specific defaults.
• sweeps/*.yaml for Weights & Biases sweep definitions.

Select a target with target=prevalence or target=cases; override config values from the command line with Hydra syntax.

Development

Run the test suite:

uv run pytest tests/

Skip slow or local-only tests:

uv run pytest tests/ -m "not slow"
uv run pytest tests/ -m "not local"
uv run pytest tests/ -m "not slow and not local" # skip both

Run linting and formatting:

uv run ruff check
uv run ruff format

Repository Layout

stateMINT/
|-- common/              # shared dataclasses, transforms, and model helpers
|-- conf/                # Hydra configs for training, export, viz, and sweeps
|-- data/                # DuckDB fetch, preprocessing, features, and loaders
|-- eval/                # metrics and prediction/target visualization helpers
|-- model/               # Mamba2 regressor and artifact loading utilities
|-- training/            # optimizer, train/eval steps, loss, and checkpointing
|-- filter_raw_data.py   # CLI for building filtered parquet datasets
|-- train.py             # Hydra training entry point
|-- visualise_predictions.py
`-- model_export.py      # Hydra artifact export entry point

tests/                   # unit tests
artifacts/               # exported model artifact examples/metadata
viz_outputs/             # generated prediction visualization outputs

Contributing

See CONTRIBUTING.md for contribution guidelines and the development workflow.