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