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