File size: 2,515 Bytes
4a94de1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | ---
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.
|