Replace xgboost by absternator · Pull Request #3 · mrc-ide/estiMINT-python

absternator · 2026-07-17T14:58:07Z

The following PR replaces the old XGBoost models with an RQS Conditional Normalising flow. The logic was coppied from paper on Neural Spine Flows

The reason for change was to improve the accuracy of the models. The normalising flow also provides us with uncertainty that we didn't have prior.

RQS Conditonal was used so we could retain desired monotonic behaviour as well as getting uncertainty.

Most of the new code is in estimint/v2 .

This reverts commit 8a01454.

…ccordingly

…with updated model configurations and calibration methods

- Updated `calib_frac` in `train_config.yaml` from 0.1 to 0.06 for improved calibration. - Reduced `num_epochs` from 300 to 200 in `train_config.yaml` to optimize training duration. - Fixed checkpoint directory path in `train_config.yaml` by removing the extra slash. - Added a blank line in `train_step.py` for better code readability. - Deleted obsolete test file `test_estimint_nn.py` which was not integrated into the package.

…ss structures

- Moved checkpoint saving logic to a separate function `save_checkpoint` in `checkpoint.py`. - Updated `train_model` function in `train_base.py` to utilize the new checkpointing function. - Removed unused checkpoint session management code for clarity. - Introduced a new YAML configuration file for sweep parameters to facilitate hyperparameter tuning.

…ning - Added `types.py` to define `ModelFactory` and `ModelArtifact` protocols. - Created `export_config.yaml` for model export configurations. - Updated `train_config.yaml` to include dynamic naming for output files based on predictor and target. - Refactored feature handling in `features.py` to dynamically include the predictor in the feature list. - Modified `preprocess.py` to fit feature scalers based on the updated feature list. - Enhanced metrics computation in `metrics.py` to utilize the `ModelArtifact` protocol. - Implemented a factory method in `rqs.py` for model instantiation from configuration. - Updated `train_base.py` to use the new model factory method and removed redundant feature list handling. - Created `calibrate.py` for conformal calibration offset calculations.

- Introduced a new sweep configuration file `hbr_eir.yaml` for hyperparameter tuning using Bayesian optimization. - Implemented `model_export.py` to export trained models along with their scalers and configuration for sharing. - Created `hub.py` to facilitate loading model artifacts from local or Hugging Face repositories. - Added logging for training processes in `train_base.log` to capture model training details and performance metrics.

- Updated version in pyproject.toml from 1.5.4 to 1.6.4. - Added 'gpu' extra dependency to the 'all' group in pyproject.toml. - Introduced new module `eir_models.py` for loading and inference of EIR models. - Refactored `hbr.py` to utilize the new EIR models for predictions. - Updated `scenarios.py` to integrate with the new EIR models and streamline predictions. - Modified type hints and internal functions to support the new model structure. - Enhanced tests to cover the new EIR model functionality and ensure correctness.

…ding OOF and test metrics for both EIR and HBR. These files contained performance metrics that are no longer relevant to the current model evaluations.

…update import paths

Copilot

Pull request overview

Replaces legacy XGBoost inference and training with conditional RQS models distributed through Hugging Face.

Changes:

Adds RQS training, calibration, checkpointing, export, and inference.
Rewires scenario and mosquito-density pipelines to use RQS artifacts.
Removes legacy XGBoost code and expands tests and documentation.

Reviewed changes

Copilot reviewed 52 out of 62 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`tests/v2/training/test_train_step.py`	Tests optimization and training loops.
`tests/v2/training/test_checkpoint.py`	Tests checkpoint persistence.
`tests/v2/training/test_calibrate.py`	Tests conformal calibration.
`tests/v2/training/__init__.py`	Adds training test package.
`tests/v2/models/test_rqs.py`	Tests RQS behavior and artifact inference.
`tests/v2/models/__init__.py`	Adds model test package.
`tests/v2/eval/test_metrics.py`	Tests metric aggregation.
`tests/v2/eval/__init__.py`	Adds evaluation test package.
`tests/v2/data/test_preprocess.py`	Tests preprocessing and splitting.
`tests/v2/data/test_features.py`	Tests scalers and feature selection.
`tests/v2/data/test_dataset.py`	Tests Grain data loading.
`tests/v2/data/__init__.py`	Adds data test package.
`tests/v2/__init__.py`	Adds v2 test package.
`tests/test_scenarios.py`	Migrates scenario tests to RQS models.
`tests/test_flows.py`	Migrates end-to-end model tests.
`tests/test_estimint.py`	Removes legacy global-model tests.
`src/estimint/v2/training/train_step.py`	Adds model training loop.
`src/estimint/v2/training/checkpoint.py`	Adds Orbax checkpoint handling.
`src/estimint/v2/training/calibrate.py`	Adds conformal interval correction.
`src/estimint/v2/train_base.py`	Adds RQS training entry point.
`src/estimint/v2/models/rqs.py`	Implements RQS model and artifact API.
`src/estimint/v2/models/mlp.py`	Adds conditioning MLP.
`src/estimint/v2/models/hub.py`	Loads local or Hugging Face artifacts.
`src/estimint/v2/model_export.py`	Exports portable model artifacts.
`src/estimint/v2/eval/metrics.py`	Adds evaluation metrics.
`src/estimint/v2/data/preprocess.py`	Adds splitting, scaling, and record preparation.
`src/estimint/v2/data/features.py`	Adds feature definitions and scalers.
`src/estimint/v2/data/dataset.py`	Adds Grain loader construction.
`src/estimint/v2/conf/train_config.yaml`	Configures training runs.
`src/estimint/v2/conf/sweeps/sweep.yaml`	Configures W&B sweeps.
`src/estimint/v2/conf/export_config.yaml`	Configures artifact export.
`src/estimint/v2/common/types.py`	Adds v2 protocols and type aliases.
`src/estimint/utils.py`	Adds bias and median percentage metrics.
`src/estimint/types.py`	Adds `PreparedScenario`.
`src/estimint/train.py`	Removes legacy XGBoost training.
`src/estimint/storage.py`	Removes legacy model storage.
`src/estimint/scenarios.py`	Uses hosted RQS models in scenarios.
`src/estimint/run.py`	Removes XGBoost inference.
`src/estimint/models.py`	Removes XGBoost model helpers.
`src/estimint/hbr.py`	Migrates mosquito-density estimation to RQS.
`src/estimint/globals.py`	Removes legacy constants.
`src/estimint/eir_models.py`	Adds RQS model loading and prediction.
`src/estimint/data/models-tag.txt`	Removes legacy model release tag.
`src/estimint/data/models-checksums.csv`	Removes legacy model checksums.
`src/estimint/data/model_checksum.txt`	Removes bundled-model checksum.
`src/estimint/data_processing.py`	Retains only value weighting.
`src/estimint/__init__.py`	Removes legacy public APIs.
`README.md`	Documents RQS inference and training.
`pyproject.toml`	Updates v2 dependencies and extras.
`models/prevalence/train.py`	Removes legacy prevalence training.
`models/prevalence/README.md`	Removes legacy prevalence documentation.
`models/prevalence/prepare.py`	Removes legacy prevalence preparation.
`models/hbr/train_hbr_to_eir.py`	Removes legacy HBR-to-EIR training.
`models/hbr/train_eir_to_hbr.py`	Removes legacy EIR-to-HBR training.
`models/hbr/README.md`	Removes legacy HBR documentation.
`models/hbr/prepare.py`	Removes legacy HBR preparation.
`.gitignore`	Ignores new training artifacts.
`.github/workflows/tests.yml`	Installs training and scenario extras in CI.

Suppressed comments (5)

src/estimint/v2/data/preprocess.py:377

The inverse-frequency weights are computed before the split from every target value, so training weights encode the validation, calibration, and test target distributions. This leaks held-out information into fitting and early stopping. Create/load the split first, then derive training weights using training pairs only (and define validation weights without consulting the other held-out splits).
src/estimint/v2/data/preprocess.py:295
All targets are passed through log10, but preprocessing only filters on prevalence and does not ensure the selected target is positive. In the supported eir -> hbr_y9 workflow, a zero/negative HBR produces -inf here and contaminates the scaler and training loss. Filter or explicitly reject non-positive target pairs before stratification and scaling.
src/estimint/v2/models/rqs.py:150
Array inputs are not checked against the artifact's feature count. A transposed (C, 1) array, for example, broadcasts against scaler statistics to (C, C) and returns multiple plausible-looking predictions instead of rejecting the malformed input. Require a 2-D array whose final dimension equals len(feature_names).
src/estimint/v2/models/rqs.py:154
Quantile probabilities are accepted without validation. Values at 0 or 1 produce infinite predictions through norm.ppf, while out-of-range values produce NaNs, so invalid user input looks like a model result. Reject probabilities outside the open interval (0, 1).
src/estimint/v2/models/rqs.py:167
alpha is not validated independently of the derived quantiles. Values between 1 and 2 pass quantile validation but reverse the lower and upper bounds, while other invalid values yield infinities/NaNs. Require 0 < alpha < 1 before constructing the interval.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

…; update configuration and tests accordingly

absternator added 16 commits July 6, 2026 08:45

Revert "chore: drop committed coverage and metrics"

f3439e0

This reverts commit 8a01454.

fix: update minimum prevalence filter to 0.01 and adjust assertions a…

90378e8

…ccordingly

update to correct data preprocessing

7d037e8

update data loading

d3f9d70

feat: consolidate "replace XGBoost" findings into a single test file …

7ba02b4

…with updated model configurations and calibration methods

get working umnn

19d5e2a

get rqs v1

65173d4

cleanup rqs logic

5aaeec5

get working version

900a047

Refactor model artifacts and update training functions to use new cla…

9806997

…ss structures

remove bad code and fix rqs

d071554

fix checkpointing dir

75588fd

refactor + remove umnn

3ac2fdd

absternator marked this pull request as draft July 17, 2026 14:59

absternator added 13 commits July 17, 2026 15:14

changes to config

d0dd8e3

fix: remove .log from .gitignore to clean up ignored files

fb49413

tests + scaler updates

9ccb0cd

save conformal so can be used by hf loaded weights

2f71f87

random

1313eb3

remove: delete unused random.py file

b42665a

get working with scenarios

bf99fa0

Remove outdated metric CSV files for HBR and prevalence models, inclu…

f6a0f54

…ding OOF and test metrics for both EIR and HBR. These files contained performance metrics that are no longer relevant to the current model evaluations.

Update project installation step to include extra scenarios for testing

43e1a19

remove all related to old xgboost

ea79df2

Implement code changes to enhance functionality and improve performance

4a74097

absternator marked this pull request as ready for review August 13, 2026 15:24

Refactor test cases in test_estimint.py for improved readability and …

d035c90

…update import paths

absternator requested a balanced review from Copilot August 13, 2026 15:35

Copilot started reviewing on behalf of absternator August 13, 2026 15:35 View session

Copilot AI reviewed Aug 13, 2026

View reviewed changes

absternator requested a review from SteamedGit August 13, 2026 15:41

absternator added 2 commits August 14, 2026 13:49

easy pr changes

1ad2124

Refactor loss functions and training steps to support weighted losses…

8cf0906

…; update configuration and tests accordingly

absternator deployed to pypi August 16, 2026 09:42 — with GitHub Actions Active

absternator mentioned this pull request Aug 16, 2026

Upgrade estimint v2.1.0 mrc-ide/mint-v2#51

Open

Navigation Menu

Uh oh!

Replace xgboost - #3

Replace xgboost#3
absternator wants to merge 32 commits into
masterfrom
replace-xgboost

absternator commented Jul 17, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

absternator commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

absternator commented Jul 17, 2026 •

edited

Loading