Metadata-Version: 2.4
Name: mintstate
Version: 0.3.0
Summary: StateMINT is a state space based neural network emulator for malariasimulation.
Project-URL: Homepage, https://github.com/mrc-ide/stateMINT
Project-URL: Repository, https://github.com/mrc-ide/stateMINT
Project-URL: Weights, https://huggingface.co/dide-ic/stateMINT
Author-email: Anmol Thapar <mr.anmolthapar@gmail.com>, Cosmo Santoni <cosmo.santoni@imperial.ac.uk>
Requires-Python: >=3.12
Requires-Dist: etils>=1.0
Requires-Dist: flax>=0.12.7
Requires-Dist: huggingface-hub>=1.19.0
Requires-Dist: jax>=0.10.1
Requires-Dist: jaxtyping>=0.3.10
Requires-Dist: mamba2-jax>=1.1.1
Requires-Dist: numpy>=2.4.6
Requires-Dist: omegaconf>=2.3
Requires-Dist: orbax-checkpoint>=0.12.0
Requires-Dist: pandas>=3.0.3
Provides-Extra: all
Requires-Dist: duckdb>=1.5.3; extra == 'all'
Requires-Dist: grain>=0.2.16; extra == 'all'
Requires-Dist: hydra-core>=1.3.2; extra == 'all'
Requires-Dist: jax[cuda12]>=0.10.1; extra == 'all'
Requires-Dist: matplotlib>=3.10.9; extra == 'all'
Requires-Dist: optax>=0.2.8; extra == 'all'
Requires-Dist: tqdm>=4.67.3; extra == 'all'
Requires-Dist: wandb>=0.27.0; extra == 'all'
Provides-Extra: gpu
Requires-Dist: jax[cuda12]>=0.10.1; extra == 'gpu'
Provides-Extra: plot
Requires-Dist: matplotlib>=3.10.9; extra == 'plot'
Requires-Dist: tqdm>=4.67.3; extra == 'plot'
Provides-Extra: train
Requires-Dist: duckdb>=1.5.3; extra == 'train'
Requires-Dist: grain>=0.2.16; extra == 'train'
Requires-Dist: hydra-core>=1.3.2; extra == 'train'
Requires-Dist: optax>=0.2.8; extra == 'train'
Requires-Dist: tqdm>=4.67.3; extra == 'train'
Requires-Dist: wandb>=0.27.0; extra == 'train'
Description-Content-Type: text/markdown

# StateMINT

[![PyPI](https://img.shields.io/pypi/v/mintstate)](https://pypi.org/project/mintstate/)


StateMINT is a JAX/Flax neural emulator for
[`malariasimulation`](https://github.com/mrc-ide/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`](https://github.com/CosmoNaught/MINTelligence) RNN emulator.

## Contents

- [Features](#features)
- [Installation](#installation)
  - [From PyPI](#from-pypi-recommended)
  - [From GitHub](#from-github-latest-unreleased)
  - [For development](#for-development-from-source)
- [Quick Start: Inference](#quick-start-inference)
- [Static Covariates](#static-covariates)
- [Training Workflow](#training-workflow)
- [Configuration](#configuration)
- [Development](#development)
- [Repository Layout](#repository-layout)
- [Contributing](#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:

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

```bash
# 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`](https://github.com/astral-sh/uv):

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

Or install a subset of extras:

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

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

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

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

For local artifacts, pass the target artifact directory:

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

## Static Covariates

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

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

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

```bash
uv run python -m stateMINT.train
```

Train the cases model:

```bash
uv run python -m stateMINT.train target=cases
```

Common overrides:

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

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

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

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

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

```bash
hf auth login
```

Upload an artifact:

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

Create a release tag:

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

```bash
uv run pytest tests/
```

Skip slow or local-only tests:

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

```bash
uv run ruff check
uv run ruff format
```

## Repository Layout

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

