| --- |
| 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, human biting rate) from a single |
| transmission-intensity predictor plus intervention-coverage covariates. |
|
|
| Each model is a conditional normalizing flow: an MLP maps context features to |
| spline parameters, which transform a standard normal base distribution into |
| the target's distribution. Sampling a quantile `q` gives calibrated predictive |
| intervals in addition to a point (median) estimate. |
|
|
| ## Available models |
|
|
| | Model | Predictor | Target | |
| |---|---|---| |
| | `hbr_y9-eir` | `hbr_y9` | `eir` | |
| | `prev_y9-eir` | `prev_y9` | `eir` | |
| | `eir-hbr_y9` | `eir` | `hbr_y9` | |
|
|
| ## Covariates |
|
|
| Every model takes the same seven raw covariates: the predictor column first, |
| then six intervention-coverage covariates. |
|
|
| | Position | Name | Notes | |
| |---|---|---| |
| | 0 | `<predictor>` | `eir`, `prev_y9`, or `hbr_y9` — whichever the model is keyed on | |
| | 1 | `dn0_use` | | |
| | 2 | `Q0` | | |
| | 3 | `phi_bednets` | | |
| | 4 | `seasonal` | | |
| | 5 | `itn_use` | | |
| | 6 | `irs_use` | | |
|
|
| So `hbr_y9-eir` expects |
| `hbr_y9, dn0_use, Q0, phi_bednets, seasonal, itn_use, irs_use`. |
|
|
| Pass raw, unstandardized values — standardization (and log10 transforms where |
| applicable) is applied internally. Predictions are returned on the original |
| (non-log) scale, clipped at 0. |
|
|
| ## Usage |
|
|
| ```python |
| from estimint.v2.models.rqs import ConditionalRQS |
| |
| # from a local export directory |
| artifact = ConditionalRQS.from_pretrained("artifacts/hbr_y9-eir", predictor="hbr_y9", target="eir") |
| |
| # from this Hub repo |
| artifact = ConditionalRQS.from_pretrained("<org>/<repo>", predictor="hbr_y9", target="eir") |
| |
| X_raw = [ |
| {"hbr_y9": 12.4, "dn0_use": 0.5, "Q0": 0.92, "phi_bednets": 0.80, |
| "seasonal": 0, "itn_use": 0.40, "irs_use": 0.00}, |
| {"hbr_y9": 3.1, "dn0_use": 0.3, "Q0": 0.90, "phi_bednets": 0.70, |
| "seasonal": 1, "itn_use": 0.20, "irs_use": 0.10}, |
| ] |
| |
| artifact.predict(X_raw) # median prediction, shape (2,) |
| artifact.quantile(X_raw, 0.9) # single quantile |
| artifact.interval(X_raw, alpha=0.10) # (lower, upper) 90% predictive interval |
| ``` |
|
|
| `X_raw` may be: |
|
|
| - a **list of dicts**, one dict per row, keyed by the covariate names above. |
| Key order does not matter — the artifact reorders each row into the |
| training order. Each dict must contain exactly the seven covariates; |
| missing or unexpected keys raise an error. |
| - a **single dict**, treated as one row (results still come back as |
| length-1 arrays). |
| - a **`(batch, 7)` array** of raw values, already in the order above. |
|
|
| Note: `interval()` currently returns the raw `[alpha/2, 1-alpha/2]` quantile |
| band; the conformal calibration offset computed during training is not yet |
| carried over to exported artifacts, so intervals from `from_pretrained` models |
| are not conformally corrected. |
|
|
| ## License |
|
|
| MIT, matching the parent [estiMINT](https://github.com/CosmoNaught/estiMINT) repo. |
|
|
