README.md

metadata

library_name: flax
tags:
  - jax
  - flax
  - orbax
  - mamba
  - malaria
  - simulation
  - time-series
  - regression
pipeline_tag: time-series-forecasting

StateMINT

StateMINT is a neural emulator for malariasimulation outputs. This model repository contains two exported inference artifacts:

• prevalence/: predicts malaria prevalence over time.
• cases/: predicts malaria cases over time.

Both artifacts use the same Mamba2Regressor architecture but have separate weights and preprocessing metadata. Users should load the folder that matches the target they want to predict.

Repository Layout

.
├── prevalence/
│   ├── checkpoint/
│   ├── model_config.json
│   └── preprocessing_config.json
├── cases/
│   ├── checkpoint/
│   ├── model_config.json
│   └── preprocessing_config.json
└── README.md

Each target folder is self-contained:

• checkpoint/ contains model-only Orbax checkpoint data.
• model_config.json contains the model architecture settings needed to instantiate Mamba2Regressor.
• preprocessing_config.json contains feature ordering, intervention timing, target transform settings, and the fitted static covariate scaler.

Intended Use

These models are intended for emulating trajectories generated by malariasimulation-style simulation inputs. They are designed for research and analysis workflows where fast approximate prediction of simulated prevalence or cases is useful.

They are not intended for direct clinical decision-making or for use on real-world surveillance data without additional validation.

Installation

Install StateMINT and the Hugging Face Hub client:

pip install stateMINT huggingface-hub

If installing from source:

git clone https://github.com/mrc-ide/stateMINT.git
cd stateMINT

pip install -e .

Loading A Model

Recommended high-level API:

from stateMINT.model import Mamba2Regressor

model = Mamba2Regressor.from_pretrained(
    "mrc-ide/stateMINT",
    predictor="prevalence",
    revision="v1.0.0",
)

model.eval()

To load the cases model:

from stateMINT.model import Mamba2Regressor

model = Mamba2Regressor.from_pretrained(
    "absternator/stateMINT",
    predictor="cases",
    revision="v1.0.0",
)

model.eval()

If you need the preprocessing metadata as well:

artifact = Mamba2Regressor.from_pretrained(
    "absternator/stateMINT",
    predictor="prevalence",
    revision="v1.0.0",
    return_artifact=True,
)

model = artifact.model
preprocessing = artifact.preprocessing
scaler = artifact.scaler

The model expects already-prepared arrays with shape:

(batch, time, input_size)

For the exported artifacts in this repository, input_size is 16.

Preprocessing Contract

The model was trained on transformed inputs, not raw covariates. To reproduce training-time behavior, users must apply the same preprocessing described in preprocessing_config.json.

The static covariates are expected in this order:

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

The following covariates are zeroed before the intervention day:

dn0_future
itn_future
irs_future
lsm
routine

The intervention day is:

3285

When use_cyclical_time is true, each timestep uses:

sin(day_of_year), cos(day_of_year), scaled_static_covariates, post_intervention_flag, years_since_intervention

The static covariates must be standardized using the fitted scaler stored in preprocessing_config.json:

scaled_static = (raw_static - scaler_mean) / scaler_scale

Do not refit the scaler for inference. The exported scaler is part of the trained model.

Prediction Scale

The model predicts in the transformed target space used during training.

For prevalence:

prevalence = sigmoid(raw_prediction)

For cases:

cases = expm1(raw_prediction)

StateMINT utilities may perform this inverse transform for you depending on the prediction helper being used.

General Notes

• The prevalence and cases folders have separate checkpoints and separate fitted scalers. Always load the folder corresponding to the target being predicted.

Citation

If you use StateMINT in research, please cite the StateMINT repository and the underlying malariasimulation work used to generate the training data.