---
license: mit
library_name: flax
tags:
  - jax
  - flax
  - nnx
  - normalizing-flow
  - rational-quadratic-spline
  - quantile-regression
  - malaria
  - epidemiology
pipeline_tag: tabular-regression
---

# estiMINT RQS models

Conditional rational-quadratic spline (RQS) normalizing flows used by
[estiMINT](https://github.com/CosmoNaught/estiMINT) to predict quantiles of
malaria transmission intensity targets (EIR, prevalence, human biting rate)
from intervention-coverage covariates.

Each model is a `ConditionalRQS` flow (`src/estimint/v2/models/rqs.py`): an
MLP maps context features to spline parameters, which transform a standard
normal base distribution into the target's (standardized, log10) distribution.
Sampling a quantile `q` amounts to inverting the flow at `Phi^-1(q)`, which
gives calibrated predictive intervals in addition to a point (median) estimate.

## Repo layout

Each subfolder is one trained predictor -> target artifact, named
`<predictor>-<target>`:

```
<predictor>-<target>/
  config.json       # architecture hyperparams + fitted feature/target scalers
  checkpoint/        # Orbax checkpoint for the model params (nnx.State)
```

| Artifact | Predictor | Target |
|---|---|---|
| `eir-prev_y9` | `eir` | `prev_y9` |

Predictor/target values come from `PredictorType` in
`src/estimint/v2/common/types.py`: `"eir"`, `"prev_y9"`, `"hbr_y9"`.

Context features (`src/estimint/v2/data/features.py`) are the predictor
value followed by `FEATURES_BASE`:
`dn0_use, Q0, phi_bednets, seasonal, itn_use, irs_use`.

## Usage

```python
from estimint.v2.models.rqs import ConditionalRQS

# from a local export directory
artifact = ConditionalRQS.from_pretrained("artifacts", name="eir-prev_y9")

# from this Hub repo
artifact = ConditionalRQS.from_pretrained("<org>/<repo>", name="eir-prev_y9")

artifact.predict(X_raw)                 # median prediction
artifact.quantile(X_raw, 0.9)           # single quantile
artifact.interval(X_raw, alpha=0.10)    # (lower, upper) conformal band
```

`X_raw` is a `(batch, n_context)` array of raw (unstandardized) feature
values in the order described above.

## Training

Models are trained with `estimint.v2.train_base` (Hydra config in
`src/estimint/v2/conf/train_config.yaml`) and exported for sharing with
`estimint.v2.model_export` (`src/estimint/v2/conf/export_config.yaml`),
which writes the `config.json` + `checkpoint/` pair uploaded here.

## License

MIT, matching the parent [estiMINT](https://github.com/CosmoNaught/estiMINT) repo.