README.md

---
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

```text
.
├── 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:

```bash
pip install stateMINT huggingface-hub
```

If installing from source:

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

pip install -e .
```

## Loading A Model

Recommended high-level API:

```python
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:

```python
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:

```python
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:

```text
(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:

```text
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:

```text
dn0_future
itn_future
irs_future
lsm
routine
```

The intervention day is:

```text
3285
```

When `use_cyclical_time` is true, each timestep uses:

```text
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`:

```python
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:

```python
prevalence = sigmoid(raw_prediction)
```

For cases:

```python
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.