diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d32b077..67c8d0b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -32,7 +32,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install project - run: uv sync --locked + run: uv sync --locked --extra train --extra scenarios - name: Run tests run: uv run pytest diff --git a/.gitignore b/.gitignore index 3ca51a6..8a1d118 100644 --- a/.gitignore +++ b/.gitignore @@ -40,12 +40,15 @@ uv.lock # Training outputs output/ scripts/output/ +*logs/ +wandb/ # Model training artifacts (regenerable; shipped copies live in src/estimint/data/) models/**/*.parquet models/**/*.pkl models/**/*.model models/**/plots/ + models/**/metrics/ # Test / coverage @@ -53,3 +56,12 @@ models/**/metrics/ .coverage.* htmlcov/ .pytest_cache/ +train_outputs/ +outputs/ +*.out +test.ipynb +artifacts/ +datasets/ +# Miscellaneous +*.log +slurm.sh diff --git a/README.md b/README.md index d6425d0..8f46897 100644 --- a/README.md +++ b/README.md @@ -1,37 +1,35 @@ # estiMINT -Package for EIR (Entomological Inoculation Rate) estimation using machine learning. +estiMINT estimates malaria transmission intensity from malaria prevalence or +human biting rate (HBR), accounting for intervention coverage. Its pretrained +models are conditional rational-quadratic spline (RQS) flows: they provide a +median prediction, arbitrary quantiles, and prediction intervals with optional +conformal calibration. -It estimates EIR from prevalence, converts between EIR and human biting rate (including the effect of changes in mosquito density), and turns a bednet specification (net type and resistance level) into the `dn0` killing parameter. +The package also provides: -## Installation +- EIR estimation from year-9 prevalence or HBR +- EIR-to-HBR conversion and projected EIR changes from mosquito-density changes +- Bednet resistance and coverage conversion to the `dn0` transmission covariate +- A `run_scenarios` pipeline that combines estiMINT estimates with the stateMINT + prevalence and case-burden emulator -```bash -pip install estimint # core: inference only (numpy, pandas, xgboost, scipy) -``` +## Installation -Optional extras, by use case: +estiMINT requires Python 3.12 or newer. ```bash -pip install "estimint[train]" # data prep + model training (duckdb, scikit-learn, pyarrow) -pip install "estimint[viz]" # plotting (matplotlib) -pip install "estimint[scenarios]" # run_scenarios pipeline (stateMINT emulator) -pip install "estimint[all]" -pip install "estimint[dev]" # test/lint/type-check toolchain +pip install estimint ``` -The `run_scenarios` pipeline also needs the stateMINT emulator (Python 3.12+). For now it -comes from the `mamba2-train` branch. With uv this is handled for you: +Optional extras are available for specific workflows: ```bash -uv sync --extra scenarios -``` - -With plain pip, install stateMINT from the branch yourself, then estiMINT: - -```bash -pip install "git+https://github.com/mrc-ide/stateMINT.git@mamba2-train" -pip install estimint +pip install "estimint[train]" # training and data preparation +pip install "estimint[gpu]" # CUDA-enabled JAX +pip install "estimint[viz]" # plotting +pip install "estimint[scenarios]" # stateMINT scenario emulator +pip install "estimint[all]" # train, viz, download, scenarios. gpu ``` For local development with [uv](https://docs.astral.sh/uv/): @@ -40,207 +38,334 @@ For local development with [uv](https://docs.astral.sh/uv/): uv sync --all-extras --dev ``` -## Data & retraining pipeline - -All training data lives in `datasets/estimint_simulations_y9.parquet`. Two model folders -derive their views from it and train: - -``` -datasets/ # training data (see datasets/README.md) -models/ - prevalence/ # prev_y9 -> EIR (estiMINT_model.pkl) - hbr/ # HBR<->EIR sub-models (estiMINT_HBR_model.pkl, estiMINT_EIR_to_HBR_model.pkl) -``` +The RQS weights are downloaded from Hugging Face when first requested and are +cached by `huggingface_hub`. The high-level scenario pipeline additionally +loads stateMINT only when `run_scenarios` is called. -Retrain a model end-to-end, e.g. the prevalence model: +## Pretrained RQS models -```bash -python models/prevalence/prepare.py # derive the training view from the parquet -python models/prevalence/train.py # train -> estiMINT_model.pkl + metrics/ + plots/ -``` +The default artifacts are hosted in [`dide-ic/estiMINT`](https://huggingface.co/dide-ic/estiMINT). -The deployed models shipped with the package live in `src/estimint/data/` and are loaded by -name (`prevalence`, `hbr`, `eir_to_hbr`). This is independent of the training pipeline above. +| Artifact | Input | Output | +| --- | --- | --- | +| `prev_y9-eir` | year-9 prevalence (`prev_y9`) | EIR (`eir`) | +| `hbr_y9-eir` | year-9 HBR (`hbr_y9`) | EIR (`eir`) | +| `eir-hbr_y9` | EIR (`eir`) | year-9 HBR (`hbr_y9`) | -## API Reference +All artifacts use the supplied transmission measure plus these intervention +covariates: `dn0_use`, `Q0`, `phi_bednets`, `seasonal`, `itn_use`, and +`irs_use`. Pass raw values; feature scaling and any required log transforms +are applied by the artifact. -### Inference +### Direct inference ```python -from estimint import load_xgb_model, run_xgb_model -import pandas as pd - -# Load a bundled model by name: "prevalence", "hbr", or "eir_to_hbr" -model = load_xgb_model("prevalence") +from estimint.v2.models.rqs import ConditionalRQS -# Prepare input data -new_data = pd.DataFrame({ - "dn0_use": [0.5], - "Q0": [0.3], - "phi_bednets": [0.6], - "seasonal": [1], - "itn_use": [0.7], - "irs_use": [0.2], - "prev_y9": [0.15] # or "prevalence" -}) +artifact = ConditionalRQS.from_pretrained( + "dide-ic/estiMINT", + predictor="prev_y9", + target="eir", +) -# Run prediction -eir_predictions = run_xgb_model(new_data, model) -print(f"Predicted EIR: {eir_predictions[0]:.2f}") +inputs = { + "prev_y9": 0.30, + "dn0_use": 0.33, + "Q0": 0.87, + "phi_bednets": 0.82, + "seasonal": 0.0, + "itn_use": 0.60, + "irs_use": 0.0, +} + +median_eir = artifact.predict(inputs)[0] +upper_quantile_eir = artifact.quantile(inputs, 0.90)[0] +lower_eir, upper_eir = artifact.interval(inputs, alpha=0.10) # arrays, one entry per row ``` -### Using Global Model +`predict()` returns the median. `quantile()` evaluates a requested probability +level, and `interval(alpha=0.10)` returns the lower and upper bounds of a 90% +prediction interval. All three return one value per input row. +Inputs may be a single feature dictionary, a list of dictionaries, or a NumPy +array whose columns are already in `artifact.feature_names` order. Dictionary +rows must contain exactly the artifact's expected features; this prevents +accidental feature-order or feature-name mismatches. -```python -from estimint import load_xgb_model, run_xgb_model, set_global_model +`interval()` widens the raw quantile band by the conformal offset stored in +`artifact.conformal`, which training calibrates for `alpha=0.10` only. If an +artifact carries no offset for the requested `alpha` the offset is zero and the +returned band is the uncorrected model quantile band. Inspect +`artifact.conformal` to see which levels a given artifact has calibrated; the +artifacts currently published on the Hub carry none. -# Set global model once -model = load_xgb_model("prevalence") -set_global_model(model) +To load an exported local artifact instead, pass its directory to +`from_pretrained`: -# Run predictions without passing model -predictions = run_xgb_model(new_data) # Uses global model +```python +artifact = ConditionalRQS.from_pretrained( + "artifacts/prev_y9-eir", + predictor="prev_y9", + target="eir", +) ``` -### Bednet to dn0 +## Bednet covariates -Turn a bednet specification (a mix of net types and an insecticide resistance level) into -the `dn0` covariate, the probability a mosquito dies on contact, along with total ITN usage. +`calculate_dn0` converts an insecticide-resistance level and a bednet-usage +mix into the `dn0` killing parameter and total ITN use. ```python from estimint import calculate_dn0, net_types -net_types() # ['pyrethroid_only', 'pyrethroid_pbo', 'pyrethroid_ppf', 'pyrethroid_pyrrole'] -res = calculate_dn0(0.5, py_only=0.4, py_pbo=0.3, py_pyrrole=0.2, py_ppf=0.1) -res.dn0, res.itn_use # weighted dn0, total net usage +net_types() +result = calculate_dn0(0.5, py_only=0.4, py_pbo=0.3, py_pyrrole=0.2, py_ppf=0.1) + +print(result.dn0, result.itn_use) ``` -### Run scenarios +The short names `py_only`, `py_pbo`, `py_pyrrole`, and `py_ppf` are accepted, +as are their canonical `pyrethroid_*` names. -`run_scenarios` runs the whole pipeline in one call. You give it a list of scenarios and -get back a DataFrame. For each scenario it works out the bednet killing effect, estimates -the EIR (from prevalence, from biting rate, or taken directly), optionally adjusts for a -change in mosquito density, then runs the stateMINT emulator forward to the prevalence and -cases trajectories. +## Run scenarios -This needs the [stateMINT](https://github.com/mrc-ide/stateMINT) package installed as well -as estiMINT. estiMINT only loads it when you call `run_scenarios`, and the model weights -download from HuggingFace. +`run_scenarios` estimates EIR and then runs the stateMINT emulator. A scenario +can start from prevalence, HBR, or a supplied EIR. A mosquito-density change is +applied only for prevalence inputs: estiMINT estimates baseline HBR, scales it +by `1 + mosquito_delta`, converts it back to EIR, and preserves the baseline +EIR estimate through that relative change. ```python -from estimint import run_scenarios -from estimint.scenarios import Scenario, EirTarget +from estimint import EirTarget, Scenario, run_scenarios scenarios = [ - Scenario(name="PBO nets, prevalence input, 60% more mosquitoes", - eir_target=EirTarget(0.30, "prevalence"), - res_use=0.55, py_pbo=0.85, - Q0=0.90, phi=0.85, seasonal=1, irs=0.40, lsm=0.0, - mosquito_delta=0.60), - Scenario(name="Biting rate input, mixed nets", - eir_target=EirTarget(250000.0, "hbr"), - res_use=0.45, py_only=0.30, py_ppf=0.20, - Q0=0.80, phi=0.82, seasonal=0, irs=0.0), - Scenario(name="EIR supplied directly, no nets", - eir_target=EirTarget(20.0, "eir"), - res_use=0.0, - Q0=0.88, phi=0.78, seasonal=1, irs=0.60), + Scenario( + name="PBO campaign with higher mosquito density", + eir_target=EirTarget(0.30, "prevalence"), + res_use=0.55, + py_pbo=0.85, + Q0=0.90, + phi=0.85, + seasonal=1.0, + irs=0.40, + mosquito_delta=0.60, + net_type_future="pyrethroid_pbo", + itn_future=0.85, + irs_future=0.40, + ), + Scenario( + name="HBR input", + eir_target=EirTarget(250000.0, "hbr"), + res_use=0.45, + py_only=0.30, + py_ppf=0.20, + Q0=0.80, + phi=0.82, + seasonal=0.0, + irs=0.0, + ), + Scenario( + name="Supplied EIR", + eir_target=EirTarget(20.0, "eir"), + res_use=0.0, + Q0=0.88, + phi=0.78, + seasonal=1.0, + irs=0.60, + ), ] -df = run_scenarios(scenarios) -print(df[["name", "eir_baseline", "eir_final", "prev_y9", "cases_endline"]]) +results = run_scenarios(scenarios) +print(results[["name", "eir_baseline", "eir_final", "prev_y9", "cases_endline"]]) ``` -Every scenario is a `Scenario` and needs `name`, `res_use`, `eir_target`, `Q0`, -`phi`, `seasonal` and `irs`. `lsm`, `routine` and `irs_future` default to 0 (note -`irs_future` does **not** default to `irs` — set it explicitly if you want IRS to -continue). **Current nets:** give a net-type usage mix (`py_only`, `py_pbo`, -`py_pyrrole`, `py_ppf` shares), or leave the net keys out for none; current and -future legs share the same `res_use`. **Future nets:** give `net_type_future` + -`itn_future` to switch net type; omit `net_type_future` and the future leg is zeroed -(it does **not** carry the current mix forward), or set `itn_future=0` to remove -nets explicitly. `mosquito_delta` only applies when `eir_target.input_mode` is `"prevalence"`. +Each `Scenario` requires `name`, `res_use`, `Q0`, `phi`, `seasonal`, `irs`, and +an `EirTarget`. Current bednet coverage is represented by a mix of `py_only`, +`py_pbo`, `py_pyrrole`, and `py_ppf`. The optional future leg is separate: +set both `net_type_future` and `itn_future` to specify future nets. It does not +inherit the current net mix. `irs_future`, `routine`, and `lsm` each default to +zero. PPF coverage additionally contributes to the emulator's LSM covariate. -The returned DataFrame has one row per scenario. Alongside the inputs it gives the -estimated EIR (`eir_baseline`, and `eir_final` after any mosquito-density change) and the -stateMINT output. That output is year-9 prevalence (`prev_y9`), endline prevalence and -cases, and the full 157-step `prevalence` and `cases` series. What you do with it is up to -you. +The returned DataFrame has one row per scenario and includes: -The `estimint.scenarios` module is also where the simulation-based inference and experiment -code will go. +- scenario and intervention covariates, including `dn0_use` and `dn0_future` +- `eir_baseline` and `eir_final` +- `hbr_baseline` and `hbr_new` for prevalence scenarios with a mosquito change +- `prev_y9`, `prev_endline`, and `cases_endline` +- 157-step `prevalence` and non-negative `cases` NumPy arrays -## Utility Functions +Call `preload_models()` before repeated scenario runs to download and cache the +estiMINT and stateMINT models explicitly. -```python -from estimint import ( - r2, rmse, mse, mae, median_ae, mae_rel, rmsle, smape, - fit_qmap_w, predict_qmap_w, scale_pos -) +## Training workflow + +> These commands require the training extra or a development installation: +> `pip install "estimint[train]"` or `uv sync --all-extras --dev`. -# Calculate metrics -y_true = [1, 2, 3, 4, 5] -y_pred = [1.1, 2.2, 2.9, 4.1, 4.8] +RQS training and artifact export code lives in `estimint.v2`. The standard +workflow is: -print(f"R²: {r2(y_true, y_pred):.4f}") -print(f"RMSE: {rmse(y_true, y_pred):.4f}") -print(f"MAE: {mae(y_true, y_pred):.4f}") +1. Prepare a simulation parquet dataset and train one of the three mappings. +2. Evaluate the generated test metrics and prediction-interval coverage. +3. Export the checkpoint and fitted preprocessing metadata as an inference + artifact. +4. Upload the artifact to Hugging Face. -# Quantile mapping calibration -cal = fit_qmap_w(y_pred, y_true) -y_calibrated = predict_qmap_w(y_pred, cal) +### 1. Train a model + +The default configuration in `estimint.v2.conf.train_config` trains the +`prev_y9 -> eir` model from `datasets/estimint_simulations_y9.parquet`. +It expects simulation identifiers (`parameter_index`, `simulation_index`), the +model's predictor and target columns, and the six intervention covariates +listed in [Pretrained RQS models](#pretrained-rqs-models). + +```bash +uv run python -m estimint.v2.train_base \ + predictor=prev_y9 \ + target=eir \ + output_dir=train_outputs/prev_y9-eir \ + use_wandb=false ``` -## Data Processing +Train the remaining deployed mappings by changing `predictor` and `target`: + +```bash +uv run python -m estimint.v2.train_base predictor=hbr_y9 target=eir +uv run python -m estimint.v2.train_base predictor=eir target=hbr_y9 +``` -These functions need the training extras. Install them with `pip install "estimint[train]"`, -which adds duckdb and scikit-learn. +Useful overrides include `data_file`, `split_file`, `use_existing_split`, +`stratify`, `num_epochs`, `batch_size`, `lr`, and the RQS architecture settings +`width`, `depth`, `n_bins`, `rqs_bounds`, `mlp_residual`, and `dropout_rate`. +Hydra prints the resolved configuration before training begins. + +Training creates or reuses the configured split CSV, writes fitted scalers to +the output directory, calculates the conformal offset, and saves the Orbax +checkpoint. With the default time-based run ID, the important outputs are: + +```text +train_outputs/prev_y9-eir/ +|-- features_scaler.pkl +|-- target_scaler.pkl +|-- conformal-.json +`-- ckpts-/ + `-- RQS/ +``` -```python -from estimint import load_and_filter, make_value_weights, strata_and_split +Use the same `` when exporting. The architecture arguments given to +export must exactly match the checkpoint's training configuration. -# Load and filter parquet data -result = load_and_filter("data.parquet", thr_lo=0.02, thr_hi=0.95) -df = result["DT"] -df_excluded = result["DT_excluded"] +### 2. Export an inference artifact -# Create inverse-frequency weights -weights = make_value_weights(df["eir"].values, digits=3) +Set the predictor, target, and training run ID. The default architecture values +already match the default training configuration; pass explicit architecture +overrides when the model was trained with non-default values. -# Stratified split -df["eir_log10"] = np.log10(df["eir"]) -df = strata_and_split(df, k_strata=16, seed=42) +```bash +RUN_ID=2026-08-13T12:00:00 + +uv run python -m estimint.v2.model_export \ + predictor=prev_y9 \ + target=eir \ + timestamp="$RUN_ID" \ + output_dir=train_outputs/prev_y9-eir \ + artifact_dir=artifacts/prev_y9-eir ``` -## Testing +The exporter restores the selected checkpoint, embeds the feature and target +scalers plus conformal offsets in `config.json`, and writes a portable artifact: + +```text +artifacts/prev_y9-eir/ +|-- config.json +`-- checkpoint/ + `-- RQS/ +``` + +Load this directory locally with `ConditionalRQS.from_pretrained`, as shown in +[Direct inference](#direct-inference). The artifact contains all inference +preprocessing metadata, so the training dataset and pickle scaler files are not +needed at prediction time. + +### 3. Upload to Hugging Face + +Authenticate with an account that can write to the target model repository: + +```bash +hf auth login +``` + +Upload the artifact under a subdirectory named exactly +`-`. This layout is required by +`ConditionalRQS.from_pretrained` when it downloads an artifact from the Hub. + +```bash +hf upload dide-ic/estiMINT \ + artifacts/prev_y9-eir \ + prev_y9-eir/ \ + --commit-message "Add prev_y9 to EIR RQS artifact" +``` + +Repeat this for `hbr_y9-eir` and `eir-hbr_y9` when publishing the complete +model set. To pin a set of Hub artifacts for reproducible inference, create a +repository tag and pass it as the `revision` argument to `from_pretrained`: + +```bash +hf repos tag create dide-ic/estiMINT v1.0.0 \ + --revision main \ + --message "Release RQS model set v1.0.0" +``` + +```python +artifact = ConditionalRQS.from_pretrained( + "dide-ic/estiMINT", + predictor="prev_y9", + target="eir", + revision="v1.0.0", +) +``` + +### 4. W&B sweeps + +`estimint.v2.conf.sweeps.sweep.yaml` defines a Bayesian sweep over learning +rate, batch size, dropout, and the RQS architecture (`width`, `depth`, +`n_bins`, `rqs_bounds`, `mlp_residual`), minimising `val/loss`. Create a sweep, +then run agents using the ID returned by Weights & Biases: ```bash -uv sync --extra dev # or: pip install -e ".[dev]" -uv run pytest # or: pytest +uv run wandb sweep src/estimint/v2/conf/sweeps/sweep.yaml +uv run wandb agent /estimint-sweep/ ``` -This covers the metric and utility helpers, the EIR estimators (prevalence, HBR and direct -EIR), the mosquito-density HBR pipeline, and the bednet calculation. +The sweep configuration currently targets `eir -> hbr_y9`; edit its final +Hydra overrides to sweep a different mapping. -## CI and releases +### Configuration reference -The test suite runs on every push and pull request across Python 3.10 to 3.14, defined in -[`.github/workflows/tests.yml`](.github/workflows/tests.yml). +- `src/estimint/v2/conf/train_config.yaml` defines data paths, split behavior, + model architecture, optimization, checkpoint location, and W&B settings. +- `src/estimint/v2/conf/export_config.yaml` maps a training run's timestamp, + checkpoint, scalers, and conformal JSON to an artifact directory. +- `src/estimint/v2/conf/sweeps/sweep.yaml` defines the W&B hyperparameter + search space. -Releases publish to PyPI from [`.github/workflows/publish.yml`](.github/workflows/publish.yml). -It builds with `uv build` and uploads with `uv publish` using -[PyPI trusted publishing](https://docs.astral.sh/uv/guides/integration/github/#publishing-to-pypi), -so no token is stored. To cut a release, bump `version` in `pyproject.toml` and publish a -GitHub Release. The first time, register this repository as a trusted publisher in the PyPI -project settings. +Run either entry point with `--help` to see the available Hydra options: -## Key Differences from R Version +```bash +uv run python -m estimint.v2.train_base --help +uv run python -m estimint.v2.model_export --help +``` + +## Testing + +```bash +uv sync --all-extras --dev +uv run pytest +``` -1. **File format**: Models saved as `.pkl` (pickle) instead of `.rds` -2. **Data handling**: Uses pandas instead of data.table -3. **Plotting**: Uses matplotlib instead of ggplot2 -4. **Global model**: Use `set_global_model()` / `get_global_model()` instead of `.GlobalEnv` +The EIR-estimation and mosquito-delta tests download the published estiMINT +artifacts on first run. The full `run_scenarios` test is skipped unless the +`scenarios` extra is installed. ## License -MIT License +MIT License \ No newline at end of file diff --git a/models/hbr/README.md b/models/hbr/README.md deleted file mode 100644 index db612df..0000000 --- a/models/hbr/README.md +++ /dev/null @@ -1,17 +0,0 @@ -# models/hbr - -The HBR feature's two sub-models, both used by `estimate_eir_with_mosquito_delta` -(`src/estimint/hbr.py`) to answer "what happens to EIR if mosquito density changes by X%?". - -| Sub-model | Direction | Bundle name | File | -|---|---|---|---| -| `train_hbr_to_eir.py` | HBR + interventions → EIR | `hbr` | `estiMINT_HBR_model.pkl` | -| `train_eir_to_hbr.py` | EIR + interventions → HBR | `eir_to_hbr` | `estiMINT_EIR_to_HBR_model.pkl` | - -```bash -python models/hbr/prepare.py # source -> hbr_training.parquet + eir_to_hbr_training.parquet -python models/hbr/train_hbr_to_eir.py # -> estiMINT_HBR_model.pkl -python models/hbr/train_eir_to_hbr.py # -> estiMINT_EIR_to_HBR_model.pkl -``` - -Deployed copies live in `src/estimint/data/`. diff --git a/models/hbr/prepare.py b/models/hbr/prepare.py deleted file mode 100644 index 5535ff1..0000000 --- a/models/hbr/prepare.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Derive the two HBR training views from datasets/estimint_simulations_y9.parquet. - -- hbr_training: HBR->EIR model (hbr_y9 > 0) -- eir_to_hbr_training: EIR->HBR model (prev_y9 >= 0.01 AND hbr_y9 > 0) - -Both sorted by key for a deterministic, reproducible view. -""" - -from pathlib import Path - -import pandas as pd - -HERE = Path(__file__).parent -SOURCE = HERE.parents[1] / "datasets" / "estimint_simulations_y9.parquet" - -KEYS = ["parameter_index", "simulation_index"] -COLS = KEYS + ["eir", "dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use", "hbr_y9"] -MIN_PREVALENCE = 0.01 - - -def prepare(): - src = pd.read_parquet(SOURCE) - - hbr = src[src.hbr_y9 > 0][COLS].sort_values(KEYS).reset_index(drop=True) - assert len(hbr) == 16384, f"expected 16,384 hbr rows, got {len(hbr):,}" - assert not hbr.isna().any().any(), "unexpected NaN in hbr view" - hbr.to_parquet(HERE / "hbr_training.parquet", index=False) - print(f"hbr view: {len(hbr):,} rows -> models/hbr/hbr_training.parquet") - - e2h = src[(src.prev_y9 >= MIN_PREVALENCE) & (src.hbr_y9 > 0)][COLS].sort_values(KEYS).reset_index(drop=True) - assert len(e2h) == 12874, f"expected 12,874 eir_to_hbr rows, got {len(e2h):,}" - assert not e2h.isna().any().any(), "unexpected NaN in eir_to_hbr view" - e2h.to_parquet(HERE / "eir_to_hbr_training.parquet", index=False) - print(f"eir_to_hbr view: {len(e2h):,} rows -> models/hbr/eir_to_hbr_training.parquet") - - -if __name__ == "__main__": - prepare() diff --git a/models/hbr/train_eir_to_hbr.py b/models/hbr/train_eir_to_hbr.py deleted file mode 100644 index 4f087ff..0000000 --- a/models/hbr/train_eir_to_hbr.py +++ /dev/null @@ -1,247 +0,0 @@ -"""Train the EIR->HBR model (eir + interventions -> hbr_y9). - -The reverse model: given baseline EIR and interventions, predict the human biting -rate so percentage mosquito-density changes can be applied. XGBoost with k-means -strata on log10(HBR), 10-fold CV, QMAP+scale calibration (no monotone constraint). -Reads models/hbr/eir_to_hbr_training.parquet; writes artifacts into this folder. -""" - -import sys -import pickle -import numpy as np -import pandas as pd -import xgboost as xgb -from pathlib import Path -from sklearn.cluster import KMeans - -sys.path.insert(0, str(Path(__file__).parents[2] / "src")) - -from estimint.utils import ( - ts, r2, rmse, mae, fit_qmap_w, predict_qmap_w, scale_pos -) -from estimint.data_processing import make_value_weights -from estimint.plotting import plot_obs_pred - -HERE = Path(__file__).parent -DATA_PATH = HERE / "eir_to_hbr_training.parquet" -OUTPUT_DIR = HERE -K_FOLDS = 10 -K_STRATA = 16 -SEED = 42 - - -def main(): - dir_plots = OUTPUT_DIR / "plots" - dir_metric = OUTPUT_DIR / "metrics" - for d in [dir_plots, dir_metric]: - d.mkdir(parents=True, exist_ok=True) - - ts("Loading training data...") - df = pd.read_parquet(DATA_PATH) - print(f"Loaded {len(df):,} rows") - - features = ["eir", "dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use"] - - df["hbr_log10"] = np.log10(df["hbr_y9"]) - - xgb_params = { - "objective": "reg:squarederror", - "eval_metric": "rmse", - "tree_method": "hist", - "max_depth": 6, - "eta": 0.05, - "subsample": 0.8, - "colsample_bytree": 0.8, - "min_child_weight": 1.0, - "lambda": 1.0, - "seed": SEED, - } - - ts("Creating %d strata on log10(HBR) and 70/15/15 split...", K_STRATA) - np.random.seed(SEED) - - hbr_log10 = df["hbr_log10"].values.reshape(-1, 1) - km = KMeans(n_clusters=K_STRATA, n_init=50, max_iter=5000, random_state=SEED) - km.fit(hbr_log10) - - centers = km.cluster_centers_.flatten() - ord_idx = np.argsort(centers) - id_map = {old_id: new_id + 1 for new_id, old_id in enumerate(ord_idx)} - df["strat_bin"] = np.array([id_map[c] for c in km.labels_]) - - df["split"] = None - for b in sorted(df["strat_bin"].unique()): - idx = df[df["strat_bin"] == b].index.tolist() - n_b = len(idx) - n_tr = int(np.floor(0.70 * n_b)) - n_val = int(np.floor(0.15 * n_b)) - - np.random.shuffle(idx) - - tr_idx = idx[:n_tr] if n_tr > 0 else [] - val_idx = idx[n_tr:n_tr + n_val] if n_val > 0 else [] - te_idx = idx[n_tr + n_val:] - - df.loc[tr_idx, "split"] = "train" - df.loc[val_idx, "split"] = "val" - df.loc[te_idx, "split"] = "test" - - df["split"] = df["split"].fillna("train") - - df_test = df[df["split"] == "test"] - X_test = df_test[features].values.astype(np.float64) - y_test = df_test["hbr_log10"].values - obs_hbr_test = np.power(10, y_test) - - ts("Test set: %d rows", len(df_test)) - - ts("Assigning %d-fold CV within TRAIN+VAL strata...", K_FOLDS) - dfcv = df[df["split"] != "test"].copy() - - np.random.seed(SEED + 1) - - dfcv["fold"] = 0 - for b in dfcv["strat_bin"].unique(): - mask = dfcv["strat_bin"] == b - n_b = mask.sum() - idx = dfcv.index[mask].tolist() - np.random.shuffle(idx) - folds = np.tile(np.arange(1, K_FOLDS + 1), int(np.ceil(n_b / K_FOLDS)))[:n_b] - np.random.shuffle(folds) - dfcv.loc[idx, "fold"] = folds - - ts("Running %d-fold CV with early stopping...", K_FOLDS) - oof_pred_raw = np.full(len(dfcv), np.nan) - best_iters = np.zeros(K_FOLDS, dtype=int) - - for k in range(1, K_FOLDS + 1): - ts(" Fold %d / %d", k, K_FOLDS) - - idx_val = dfcv["fold"] == k - idx_tr = dfcv["fold"] != k - - X_tr = dfcv.loc[idx_tr, features].values.astype(np.float64) - y_tr = dfcv.loc[idx_tr, "hbr_log10"].values - X_va = dfcv.loc[idx_val, features].values.astype(np.float64) - y_va = dfcv.loc[idx_val, "hbr_log10"].values - - w_tr = make_value_weights(np.power(10, y_tr), digits=3) - w_va = make_value_weights(np.power(10, y_va), digits=3) - - dtr = xgb.DMatrix(X_tr, label=y_tr, weight=w_tr) - dva = xgb.DMatrix(X_va, label=y_va, weight=w_va) - - mdl = xgb.train( - params=xgb_params, - dtrain=dtr, - num_boost_round=5000, - evals=[(dtr, "train"), (dva, "val")], - early_stopping_rounds=100, - verbose_eval=False, - ) - - best_iters[k - 1] = mdl.best_iteration - pred_log10_va = mdl.predict(dva) - oof_pred_raw[idx_val.values] = np.power(10, pred_log10_va) - - obs_cv_raw = np.power(10, dfcv["hbr_log10"].values) - - ts("Fitting final calibrator (QMAP + positive scale) on OOF...") - cal_oof = fit_qmap_w(oof_pred_raw, obs_cv_raw, ngrid=1024, round_digits=8) - oof_pred_cal = predict_qmap_w(oof_pred_raw, cal_oof) - a_oof = scale_pos(obs_cv_raw, oof_pred_cal) - oof_pred_final = np.maximum(0, a_oof * oof_pred_cal) - - oof_metrics = pd.DataFrame({ - "set": ["OOF_uncalibrated", "OOF_calibrated"], - "R2": [r2(obs_cv_raw, oof_pred_raw), r2(obs_cv_raw, oof_pred_final)], - "bias": [np.mean(oof_pred_raw - obs_cv_raw), np.mean(oof_pred_final - obs_cv_raw)], - "RMSE": [rmse(obs_cv_raw, oof_pred_raw), rmse(obs_cv_raw, oof_pred_final)], - "MAE": [mae(obs_cv_raw, oof_pred_raw), mae(obs_cv_raw, oof_pred_final)], - }) - oof_metrics.to_csv(dir_metric / f"hbr_OOF_metrics_K{K_FOLDS}CV.csv", index=False) - print("\n" + str(oof_metrics)) - - ts("Training final model on TRAIN+VAL with nrounds = median(best_iteration)...") - best_nrounds = int(np.round(np.median(best_iters))) - print(f"Best nrounds: {best_nrounds}") - - df_trcv = df[df["split"] != "test"] - X_trcv = df_trcv[features].values.astype(np.float64) - y_trcv = df_trcv["hbr_log10"].values - w_trcv = make_value_weights(np.power(10, y_trcv), digits=3) - - dtrcv = xgb.DMatrix(X_trcv, label=y_trcv, weight=w_trcv) - - xgb_final = xgb.train( - params=xgb_params, - dtrain=dtrcv, - num_boost_round=best_nrounds, - verbose_eval=False, - ) - xgb_final.save_model(str(OUTPUT_DIR / "hbr_xgb_FINAL.model")) - - dtest = xgb.DMatrix(X_test, label=y_test) - pred_log10_test_raw = xgb_final.predict(dtest) - pred_raw_test = np.power(10, pred_log10_test_raw) - pred_hbr_test = predict_qmap_w(pred_raw_test, cal_oof) - pred_hbr_test = np.maximum(0, a_oof * pred_hbr_test) - - test_metrics = pd.DataFrame({ - "set": ["Test"], - "R2": [r2(obs_hbr_test, pred_hbr_test)], - "bias": [np.mean(pred_hbr_test - obs_hbr_test)], - "RMSE": [rmse(obs_hbr_test, pred_hbr_test)], - "MAE": [mae(obs_hbr_test, pred_hbr_test)], - }) - test_metrics.to_csv(dir_metric / "hbr_test_metrics.csv", index=False) - print("\n" + str(test_metrics)) - - plot_obs_pred( - obs_hbr_test, pred_hbr_test, - f"HBR — Observed vs Predicted (XGBoost, K={K_FOLDS} CV, QMAP+Scale, test)", - str(dir_plots / "hbr_obs_vs_pred_xgb_QMAP_SCALE_test.png"), - xlab="Observed HBR", ylab="Predicted HBR" - ) - - cal_bundle = { - "kind": "qmap+scale", - "qmap": {"xq": cal_oof["xq"], "yq": cal_oof["yq"]}, - "scale": a_oof - } - - preprocess = { - "features": features, - "target": "hbr_y9", - "transform": "log10", - "inverse": "pow10", - "training_data": { - "source": "datasets/estimint_simulations_y9.parquet (prev_y9 >= 0.01 AND hbr_y9 > 0)", - "n_rows": len(df), - "n_params": df["parameter_index"].nunique() - }, - "cv": { - "K": K_FOLDS, - "stratify_by": f"strat_bin (k-means on log10(HBR), centers={K_STRATA})", - "best_iteration_median": best_nrounds - }, - } - - model_bundle = { - "class": "estiMINT_EIR_to_HBR_model", - "booster": xgb_final, - "calibrator": cal_bundle, - "features": features, - "best_nrounds": best_nrounds, - "preprocess": preprocess, - } - - with open(OUTPUT_DIR / "estiMINT_EIR_to_HBR_model.pkl", "wb") as f: - pickle.dump(model_bundle, f, protocol=pickle.HIGHEST_PROTOCOL) - - print(f"\nModel saved to: {OUTPUT_DIR}/estiMINT_EIR_to_HBR_model.pkl") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/models/hbr/train_hbr_to_eir.py b/models/hbr/train_hbr_to_eir.py deleted file mode 100644 index b42ea80..0000000 --- a/models/hbr/train_hbr_to_eir.py +++ /dev/null @@ -1,252 +0,0 @@ -"""Train the HBR->EIR model (hbr_y9 + interventions -> EIR). - -Same pipeline as the prevalence model but with hbr_y9 as the monotone feature. -Reads models/hbr/hbr_training.parquet; writes the model, booster, metrics and plot -into this folder. -""" - -import sys -import pickle -import numpy as np -import pandas as pd -import xgboost as xgb -from pathlib import Path -from sklearn.cluster import KMeans - -sys.path.insert(0, str(Path(__file__).parents[2] / "src")) - -from estimint.utils import ( - ts, r2, rmse, mae, fit_qmap_w, predict_qmap_w, scale_pos -) -from estimint.data_processing import make_value_weights -from estimint.plotting import plot_obs_pred - -HERE = Path(__file__).parent -DATA_PATH = HERE / "hbr_training.parquet" -OUTPUT_DIR = HERE -K_FOLDS = 10 -K_STRATA = 16 -SEED = 42 - - -def main(): - dir_plots = OUTPUT_DIR / "plots" - dir_metric = OUTPUT_DIR / "metrics" - for d in [dir_plots, dir_metric]: - d.mkdir(parents=True, exist_ok=True) - - ts("Loading training data...") - df = pd.read_parquet(DATA_PATH) - print(f"Loaded {len(df):,} rows") - - features = ["dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use", "hbr_y9"] - - df["eir_log10"] = np.log10(df["eir"]) - - # monotone constraint: hbr_y9 (index 6) positively correlated with EIR - xgb_params = { - "objective": "reg:squarederror", - "eval_metric": "rmse", - "tree_method": "hist", - "max_bin": 4096, - "max_depth": 8, - "eta": 0.02, - "subsample": 0.8, - "colsample_bytree": 0.8, - "min_child_weight": 1.0, - "lambda": 1.0, - "seed": SEED, - "monotone_constraints": "(0,0,0,0,0,0,1)", - } - - ts("Creating %d strata on log10(EIR) and 70/15/15 split...", K_STRATA) - np.random.seed(SEED) - - eir_log10 = df["eir_log10"].values.reshape(-1, 1) - km = KMeans(n_clusters=K_STRATA, n_init=50, max_iter=5000, random_state=SEED) - km.fit(eir_log10) - - centers = km.cluster_centers_.flatten() - ord_idx = np.argsort(centers) - id_map = {old_id: new_id + 1 for new_id, old_id in enumerate(ord_idx)} - df["strat_bin"] = np.array([id_map[c] for c in km.labels_]) - - df["split"] = None - for b in sorted(df["strat_bin"].unique()): - idx = df[df["strat_bin"] == b].index.tolist() - n_b = len(idx) - n_tr = int(np.floor(0.70 * n_b)) - n_val = int(np.floor(0.15 * n_b)) - - np.random.shuffle(idx) - - tr_idx = idx[:n_tr] if n_tr > 0 else [] - val_idx = idx[n_tr:n_tr + n_val] if n_val > 0 else [] - te_idx = idx[n_tr + n_val:] - - df.loc[tr_idx, "split"] = "train" - df.loc[val_idx, "split"] = "val" - df.loc[te_idx, "split"] = "test" - - df["split"] = df["split"].fillna("train") - - df_test = df[df["split"] == "test"] - X_test = df_test[features].values.astype(np.float64) - y_test = df_test["eir_log10"].values - obs_eir_test = np.power(10, y_test) - - ts("Test set: %d rows", len(df_test)) - - ts("Assigning %d-fold CV within TRAIN+VAL strata...", K_FOLDS) - dfcv = df[df["split"] != "test"].copy() - - np.random.seed(SEED + 1) - - dfcv["fold"] = 0 - for b in dfcv["strat_bin"].unique(): - mask = dfcv["strat_bin"] == b - n_b = mask.sum() - idx = dfcv.index[mask].tolist() - np.random.shuffle(idx) - folds = np.tile(np.arange(1, K_FOLDS + 1), int(np.ceil(n_b / K_FOLDS)))[:n_b] - np.random.shuffle(folds) - dfcv.loc[idx, "fold"] = folds - - ts("Running %d-fold CV with early stopping...", K_FOLDS) - oof_pred_raw = np.full(len(dfcv), np.nan) - best_iters = np.zeros(K_FOLDS, dtype=int) - - for k in range(1, K_FOLDS + 1): - ts(" Fold %d / %d", k, K_FOLDS) - - idx_val = dfcv["fold"] == k - idx_tr = dfcv["fold"] != k - - X_tr = dfcv.loc[idx_tr, features].values.astype(np.float64) - y_tr = dfcv.loc[idx_tr, "eir_log10"].values - X_va = dfcv.loc[idx_val, features].values.astype(np.float64) - y_va = dfcv.loc[idx_val, "eir_log10"].values - - w_tr = make_value_weights(np.power(10, y_tr), digits=3) - w_va = make_value_weights(np.power(10, y_va), digits=3) - - dtr = xgb.DMatrix(X_tr, label=y_tr, weight=w_tr) - dva = xgb.DMatrix(X_va, label=y_va, weight=w_va) - - mdl = xgb.train( - params=xgb_params, - dtrain=dtr, - num_boost_round=15000, - evals=[(dtr, "train"), (dva, "val")], - early_stopping_rounds=200, - verbose_eval=False, - ) - - best_iters[k - 1] = mdl.best_iteration - pred_log10_va = mdl.predict(dva) - oof_pred_raw[idx_val.values] = np.power(10, pred_log10_va) - - obs_cv_raw = np.power(10, dfcv["eir_log10"].values) - - ts("Fitting final calibrator (QMAP + positive scale) on OOF...") - cal_oof = fit_qmap_w(oof_pred_raw, obs_cv_raw, ngrid=1024, round_digits=8) - oof_pred_cal = predict_qmap_w(oof_pred_raw, cal_oof) - a_oof = scale_pos(obs_cv_raw, oof_pred_cal) - oof_pred_final = np.maximum(0, a_oof * oof_pred_cal) - - oof_metrics = pd.DataFrame({ - "set": ["OOF_uncalibrated", "OOF_calibrated"], - "R2": [r2(obs_cv_raw, oof_pred_raw), r2(obs_cv_raw, oof_pred_final)], - "bias": [np.mean(oof_pred_raw - obs_cv_raw), np.mean(oof_pred_final - obs_cv_raw)], - "RMSE": [rmse(obs_cv_raw, oof_pred_raw), rmse(obs_cv_raw, oof_pred_final)], - "MAE": [mae(obs_cv_raw, oof_pred_raw), mae(obs_cv_raw, oof_pred_final)], - }) - oof_metrics.to_csv(dir_metric / f"eir_OOF_metrics_K{K_FOLDS}CV.csv", index=False) - print("\n" + str(oof_metrics)) - - ts("Training final model on TRAIN+VAL with nrounds = median(best_iteration)...") - best_nrounds = int(np.round(np.median(best_iters))) - print(f"Best nrounds: {best_nrounds}") - - df_trcv = df[df["split"] != "test"] - X_trcv = df_trcv[features].values.astype(np.float64) - y_trcv = df_trcv["eir_log10"].values - w_trcv = make_value_weights(np.power(10, y_trcv), digits=3) - - dtrcv = xgb.DMatrix(X_trcv, label=y_trcv, weight=w_trcv) - - xgb_final = xgb.train( - params=xgb_params, - dtrain=dtrcv, - num_boost_round=best_nrounds, - verbose_eval=False, - ) - xgb_final.save_model(str(OUTPUT_DIR / "eir_xgb_HBR_FINAL.model")) - - dtest = xgb.DMatrix(X_test, label=y_test) - pred_log10_test_raw = xgb_final.predict(dtest) - pred_raw_test = np.power(10, pred_log10_test_raw) - pred_eir_test = predict_qmap_w(pred_raw_test, cal_oof) - pred_eir_test = np.maximum(0, a_oof * pred_eir_test) - - test_metrics = pd.DataFrame({ - "set": ["Test"], - "R2": [r2(obs_eir_test, pred_eir_test)], - "bias": [np.mean(pred_eir_test - obs_eir_test)], - "RMSE": [rmse(obs_eir_test, pred_eir_test)], - "MAE": [mae(obs_eir_test, pred_eir_test)], - }) - test_metrics.to_csv(dir_metric / "eir_test_metrics.csv", index=False) - print("\n" + str(test_metrics)) - - plot_obs_pred( - obs_eir_test, pred_eir_test, - f"EIR — Observed vs Predicted (XGBoost HBR, K={K_FOLDS} CV, QMAP+Scale, test)", - str(dir_plots / "eir_obs_vs_pred_xgb_HBR_QMAP_SCALE_test.png"), - xlab="Observed EIR", ylab="Predicted EIR" - ) - - cal_bundle = { - "kind": "qmap+scale", - "qmap": {"xq": cal_oof["xq"], "yq": cal_oof["yq"]}, - "scale": a_oof - } - - preprocess = { - "features": features, - "target": "eir", - "transform": "log10", - "inverse": "pow10", - "hbr_filter": { - "note": "Trained on MINTelligence data, HBR > 0 (Im > 0)" - }, - "training_data": { - "source": "datasets/estimint_simulations_y9.parquet (hbr_y9 > 0)", - "n_rows": len(df), - "n_params": df["parameter_index"].nunique() - }, - "cv": { - "K": K_FOLDS, - "stratify_by": f"strat_bin (k-means on log10(EIR), centers={K_STRATA})", - "best_iteration_median": best_nrounds - }, - } - - model_bundle = { - "class": "estiMINT_HBR_model", - "booster": xgb_final, - "calibrator": cal_bundle, - "features": features, - "best_nrounds": best_nrounds, - "preprocess": preprocess, - } - - with open(OUTPUT_DIR / "estiMINT_HBR_model.pkl", "wb") as f: - pickle.dump(model_bundle, f, protocol=pickle.HIGHEST_PROTOCOL) - - print(f"\nModel saved to: {OUTPUT_DIR}/estiMINT_HBR_model.pkl") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/models/prevalence/README.md b/models/prevalence/README.md deleted file mode 100644 index e43e317..0000000 --- a/models/prevalence/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# models/prevalence - -Prevalence → EIR model — the default estiMINT emulator. - -Predicts **EIR** from year-9 prevalence (`prev_y9`) + 6 interventions. Bundled as -`prevalence` → `estiMINT_model.pkl`. - -```bash -python models/prevalence/prepare.py # datasets source -> training.parquet (prev_y9 >= 0.02) -python models/prevalence/train.py # -> estiMINT_model.pkl, eir_xgb_FINAL.model, metrics/, plots/ -``` - -Deployed copy lives in `src/estimint/data/estiMINT_model.pkl`. diff --git a/models/prevalence/prepare.py b/models/prevalence/prepare.py deleted file mode 100644 index 08b6675..0000000 --- a/models/prevalence/prepare.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Derive the prevalence->EIR training view from datasets/estimint_simulations_y9.parquet. - -Filters prev_y9 >= 0.02 and sorts by key for a deterministic, reproducible view. -""" - -from pathlib import Path - -import pandas as pd - -ROOT = Path(__file__).parents[2] -SOURCE = ROOT / "datasets" / "estimint_simulations_y9.parquet" -OUT = Path(__file__).parent / "training.parquet" - -KEYS = ["parameter_index", "simulation_index"] -COLS = KEYS + ["eir", "dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use", "prev_y9"] -MIN_PREVALENCE = 0.02 - - -def prepare(): - src = pd.read_parquet(SOURCE) - view = src[src.prev_y9 >= MIN_PREVALENCE][COLS].sort_values(KEYS).reset_index(drop=True) - - assert len(view) == 12429, f"expected 12,429 rows, got {len(view):,}" - assert not view.isna().any().any(), "unexpected NaN in prevalence view" - view.to_parquet(OUT, index=False) - print(f"prevalence view: {len(view):,} rows -> {OUT}") - - -if __name__ == "__main__": - prepare() diff --git a/models/prevalence/train.py b/models/prevalence/train.py deleted file mode 100644 index 3cdb65c..0000000 --- a/models/prevalence/train.py +++ /dev/null @@ -1,253 +0,0 @@ -"""Train the prevalence->EIR model (prev_y9 + interventions -> EIR). - -XGBoost with k-means strata on log10(EIR), 10-fold CV, and QMAP+scale calibration. -Reads models/prevalence/training.parquet; writes the model, booster, metrics and plot -into this folder. -""" - -import sys -import pickle -import numpy as np -import pandas as pd -import xgboost as xgb -from pathlib import Path -from sklearn.cluster import KMeans - -sys.path.insert(0, str(Path(__file__).parents[2] / "src")) - -from estimint.utils import ( - ts, r2, rmse, mae, fit_qmap_w, predict_qmap_w, scale_pos -) -from estimint.data_processing import make_value_weights -from estimint.plotting import plot_obs_pred - -HERE = Path(__file__).parent -DATA_PATH = HERE / "training.parquet" -OUTPUT_DIR = HERE -K_FOLDS = 10 -K_STRATA = 16 -SEED = 42 - - -def main(): - dir_plots = OUTPUT_DIR / "plots" - dir_metric = OUTPUT_DIR / "metrics" - for d in [dir_plots, dir_metric]: - d.mkdir(parents=True, exist_ok=True) - - ts("Loading training data...") - df = pd.read_parquet(DATA_PATH) - print(f"Loaded {len(df):,} rows") - - features = ["dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use", "prev_y9"] - - df["eir_log10"] = np.log10(df["eir"]) - - # monotone constraint: prev_y9 (index 6) positively correlated with EIR - xgb_params = { - "objective": "reg:squarederror", - "eval_metric": "rmse", - "tree_method": "hist", - "max_bin": 4096, - "max_depth": 8, - "eta": 0.02, - "subsample": 0.8, - "colsample_bytree": 0.8, - "min_child_weight": 1.0, - "lambda": 1.0, - "seed": SEED, - "monotone_constraints": "(0,0,0,0,0,0,1)", - } - - ts("Creating %d strata on log10(EIR) and 70/15/15 split...", K_STRATA) - np.random.seed(SEED) - - eir_log10 = df["eir_log10"].values.reshape(-1, 1) - km = KMeans(n_clusters=K_STRATA, n_init=50, max_iter=5000, random_state=SEED) - km.fit(eir_log10) - - centers = km.cluster_centers_.flatten() - ord_idx = np.argsort(centers) - id_map = {old_id: new_id + 1 for new_id, old_id in enumerate(ord_idx)} - df["strat_bin"] = np.array([id_map[c] for c in km.labels_]) - - df["split"] = None - for b in sorted(df["strat_bin"].unique()): - idx = df[df["strat_bin"] == b].index.tolist() - n_b = len(idx) - n_tr = int(np.floor(0.70 * n_b)) - n_val = int(np.floor(0.15 * n_b)) - - np.random.shuffle(idx) - - tr_idx = idx[:n_tr] if n_tr > 0 else [] - val_idx = idx[n_tr:n_tr + n_val] if n_val > 0 else [] - te_idx = idx[n_tr + n_val:] - - df.loc[tr_idx, "split"] = "train" - df.loc[val_idx, "split"] = "val" - df.loc[te_idx, "split"] = "test" - - df["split"] = df["split"].fillna("train") - - df_test = df[df["split"] == "test"] - X_test = df_test[features].values.astype(np.float64) - y_test = df_test["eir_log10"].values - obs_eir_test = np.power(10, y_test) - - ts("Test set: %d rows", len(df_test)) - - ts("Assigning %d-fold CV within TRAIN+VAL strata...", K_FOLDS) - dfcv = df[df["split"] != "test"].copy() - - np.random.seed(SEED + 1) - - dfcv["fold"] = 0 - for b in dfcv["strat_bin"].unique(): - mask = dfcv["strat_bin"] == b - n_b = mask.sum() - idx = dfcv.index[mask].tolist() - np.random.shuffle(idx) - folds = np.tile(np.arange(1, K_FOLDS + 1), int(np.ceil(n_b / K_FOLDS)))[:n_b] - np.random.shuffle(folds) - dfcv.loc[idx, "fold"] = folds - - ts("Running %d-fold CV with early stopping...", K_FOLDS) - oof_pred_raw = np.full(len(dfcv), np.nan) - best_iters = np.zeros(K_FOLDS, dtype=int) - - for k in range(1, K_FOLDS + 1): - ts(" Fold %d / %d", k, K_FOLDS) - - idx_val = dfcv["fold"] == k - idx_tr = dfcv["fold"] != k - - X_tr = dfcv.loc[idx_tr, features].values.astype(np.float64) - y_tr = dfcv.loc[idx_tr, "eir_log10"].values - X_va = dfcv.loc[idx_val, features].values.astype(np.float64) - y_va = dfcv.loc[idx_val, "eir_log10"].values - - w_tr = make_value_weights(np.power(10, y_tr), digits=3) - w_va = make_value_weights(np.power(10, y_va), digits=3) - - dtr = xgb.DMatrix(X_tr, label=y_tr, weight=w_tr) - dva = xgb.DMatrix(X_va, label=y_va, weight=w_va) - - mdl = xgb.train( - params=xgb_params, - dtrain=dtr, - num_boost_round=15000, - evals=[(dtr, "train"), (dva, "val")], - early_stopping_rounds=200, - verbose_eval=False, - ) - - best_iters[k - 1] = mdl.best_iteration - pred_log10_va = mdl.predict(dva) - oof_pred_raw[idx_val.values] = np.power(10, pred_log10_va) - - obs_cv_raw = np.power(10, dfcv["eir_log10"].values) - - ts("Fitting final calibrator (QMAP + positive scale) on OOF...") - cal_oof = fit_qmap_w(oof_pred_raw, obs_cv_raw, ngrid=1024, round_digits=8) - oof_pred_cal = predict_qmap_w(oof_pred_raw, cal_oof) - a_oof = scale_pos(obs_cv_raw, oof_pred_cal) - oof_pred_final = np.maximum(0, a_oof * oof_pred_cal) - - oof_metrics = pd.DataFrame({ - "set": ["OOF_uncalibrated", "OOF_calibrated"], - "R2": [r2(obs_cv_raw, oof_pred_raw), r2(obs_cv_raw, oof_pred_final)], - "bias": [np.mean(oof_pred_raw - obs_cv_raw), np.mean(oof_pred_final - obs_cv_raw)], - "RMSE": [rmse(obs_cv_raw, oof_pred_raw), rmse(obs_cv_raw, oof_pred_final)], - "MAE": [mae(obs_cv_raw, oof_pred_raw), mae(obs_cv_raw, oof_pred_final)], - }) - oof_metrics.to_csv(dir_metric / f"eir_OOF_metrics_K{K_FOLDS}CV.csv", index=False) - print("\n" + str(oof_metrics)) - - ts("Training final model on TRAIN+VAL with nrounds = median(best_iteration)...") - best_nrounds = int(np.round(np.median(best_iters))) - print(f"Best nrounds: {best_nrounds}") - - df_trcv = df[df["split"] != "test"] - X_trcv = df_trcv[features].values.astype(np.float64) - y_trcv = df_trcv["eir_log10"].values - w_trcv = make_value_weights(np.power(10, y_trcv), digits=3) - - dtrcv = xgb.DMatrix(X_trcv, label=y_trcv, weight=w_trcv) - - xgb_final = xgb.train( - params=xgb_params, - dtrain=dtrcv, - num_boost_round=best_nrounds, - verbose_eval=False, - ) - xgb_final.save_model(str(OUTPUT_DIR / "eir_xgb_FINAL.model")) - - dtest = xgb.DMatrix(X_test, label=y_test) - pred_log10_test_raw = xgb_final.predict(dtest) - pred_raw_test = np.power(10, pred_log10_test_raw) - pred_eir_test = predict_qmap_w(pred_raw_test, cal_oof) - pred_eir_test = np.maximum(0, a_oof * pred_eir_test) - - test_metrics = pd.DataFrame({ - "set": ["Test"], - "R2": [r2(obs_eir_test, pred_eir_test)], - "bias": [np.mean(pred_eir_test - obs_eir_test)], - "RMSE": [rmse(obs_eir_test, pred_eir_test)], - "MAE": [mae(obs_eir_test, pred_eir_test)], - }) - test_metrics.to_csv(dir_metric / "eir_test_metrics.csv", index=False) - print("\n" + str(test_metrics)) - - plot_obs_pred( - obs_eir_test, pred_eir_test, - f"EIR — Observed vs Predicted (XGBoost, K={K_FOLDS} CV, QMAP+Scale, test)", - str(dir_plots / "eir_obs_vs_pred_xgb_QMAP_SCALE_test.png"), - xlab="Observed EIR", ylab="Predicted EIR" - ) - - cal_bundle = { - "kind": "qmap+scale", - "qmap": {"xq": cal_oof["xq"], "yq": cal_oof["yq"]}, - "scale": a_oof - } - - preprocess = { - "features": features, - "target": "eir", - "transform": "log10", - "inverse": "pow10", - "prevalence_filter": { - "min_prev_input": 0.02, - "note": "Trained on MINTelligence data with prev >= 0.02" - }, - "training_data": { - "source": "datasets/estimint_simulations_y9.parquet (prev_y9 >= 0.02)", - "n_rows": len(df), - "n_params": df["parameter_index"].nunique() - }, - "cv": { - "K": K_FOLDS, - "stratify_by": f"strat_bin (k-means on log10(EIR), centers={K_STRATA})", - "best_iteration_median": best_nrounds - }, - } - - model_bundle = { - "class": "estiMINT_model", - "booster": xgb_final, - "calibrator": cal_bundle, - "features": features, - "best_nrounds": best_nrounds, - "preprocess": preprocess, - } - - with open(OUTPUT_DIR / "estiMINT_model.pkl", "wb") as f: - pickle.dump(model_bundle, f, protocol=pickle.HIGHEST_PROTOCOL) - - print(f"\nModel saved to: {OUTPUT_DIR}/estiMINT_model.pkl") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml index 46c921a..5b78359 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "estimint" -version = "1.5.4" +version = "2.1.0" description = "EIR Estimation using Machine learning interventions " readme = "README.md" license = "MIT" @@ -17,23 +17,31 @@ keywords = [ "malaria", "EIR", "machine-learning", - "xgboost", + "flows", "epidemiology", ] dependencies = [ "numpy>=1.20.0", "pandas>=1.3.0", - "xgboost>=1.6.0", - "scipy>=1.7.0", + "jax>=0.10.1", + "flax>=0.12.7", + "jaxtyping>=0.3.10", + "omegaconf>=2.3", + "orbax-checkpoint>=0.12.0", + "wandb>=0.28.0", + "huggingface_hub>=0.24.0", ] [project.optional-dependencies] train = [ "duckdb>=0.8.0", - "scikit-learn>=1.0.0", - "pyarrow>=10.0.0", + "grain>=0.2.16", + "hydra-core>=1.3.2", + "optax>=0.2.8", + "tqdm>=4.67.3", ] +gpu = ["jax[cuda12]>=0.10.1"] viz = [ "matplotlib>=3.4.0", ] @@ -44,7 +52,7 @@ download = [ scenarios = [ "mintstate>=0.3.0" ] -all = ["estimint[train,viz,download,scenarios]"] +all = ["estimint[train,viz,download,scenarios,gpu]"] [dependency-groups] dev = [ diff --git a/src/estimint/__init__.py b/src/estimint/__init__.py index 95c1834..3a6718a 100644 --- a/src/estimint/__init__.py +++ b/src/estimint/__init__.py @@ -1,18 +1,3 @@ -""" -estiMINT - EIR Estimation using Machine learning INTerventions - -This package provides tools for training and running XGBoost models -to predict Entomological Inoculation Rate (EIR) from malaria intervention data. - -Dependencies ------------- -Core (inference): numpy, pandas, xgboost, scipy. -Optional extras: -- train: duckdb, scikit-learn, pyarrow (data prep + model training) -- viz: matplotlib (plotting) -- download: requests, appdirs (fetch published models) -""" - __package_name__ = "estiMINT" # Public API exports @@ -32,26 +17,9 @@ scale_pos, ) -from .data_processing import ( - load_and_filter, - make_value_weights, - strata_and_split, -) - -from .models import train_eir_xgboost - -from .train import train_xgb_model from .plotting import plot_obs_pred -from .storage import ( - load_xgb_model, - save_xgb_model, - bundle_model, -) - -from .run import run_xgb_model, set_global_model, get_global_model - from .hbr import estimate_eir_with_mosquito_delta from .bednet import calculate_dn0, net_types, DN0Result @@ -74,24 +42,8 @@ "fit_qmap_w", "predict_qmap_w", "scale_pos", - # data_processing - "load_and_filter", - "make_value_weights", - "strata_and_split", - # models - "train_eir_xgboost", - # train - "train_xgb_model", # plotting "plot_obs_pred", - # storage - "load_xgb_model", - "save_xgb_model", - "bundle_model", - # run - "run_xgb_model", - "set_global_model", - "get_global_model", # hbr "estimate_eir_with_mosquito_delta", # bednet diff --git a/src/estimint/data/estiMINT_EIR_to_HBR_model.pkl b/src/estimint/data/estiMINT_EIR_to_HBR_model.pkl deleted file mode 100644 index 604bebd..0000000 Binary files a/src/estimint/data/estiMINT_EIR_to_HBR_model.pkl and /dev/null differ diff --git a/src/estimint/data/estiMINT_HBR_model.pkl b/src/estimint/data/estiMINT_HBR_model.pkl deleted file mode 100644 index e6551c8..0000000 Binary files a/src/estimint/data/estiMINT_HBR_model.pkl and /dev/null differ diff --git a/src/estimint/data/estiMINT_model.pkl b/src/estimint/data/estiMINT_model.pkl deleted file mode 100644 index ab72a29..0000000 Binary files a/src/estimint/data/estiMINT_model.pkl and /dev/null differ diff --git a/src/estimint/data/model_checksum.txt b/src/estimint/data/model_checksum.txt deleted file mode 100644 index 19c887e..0000000 --- a/src/estimint/data/model_checksum.txt +++ /dev/null @@ -1 +0,0 @@ -b626baaaa2c260a86f2be5e8f2ee588e estiMINT_model.pkl diff --git a/src/estimint/data/models-checksums.csv b/src/estimint/data/models-checksums.csv deleted file mode 100644 index cec8fb3..0000000 --- a/src/estimint/data/models-checksums.csv +++ /dev/null @@ -1,2 +0,0 @@ -path,md5,size_B -estiMINT_model.pkl,d657c95899d82aae28773b2b4a658ab0,17066281 diff --git a/src/estimint/data/models-tag.txt b/src/estimint/data/models-tag.txt deleted file mode 100644 index f65def0..0000000 --- a/src/estimint/data/models-tag.txt +++ /dev/null @@ -1 +0,0 @@ -bundled-v1 diff --git a/src/estimint/data_processing.py b/src/estimint/data_processing.py index 2044093..4a87448 100644 --- a/src/estimint/data_processing.py +++ b/src/estimint/data_processing.py @@ -1,109 +1,18 @@ -""" -Data processing functions for estiMINT package. - -Equivalent to: data_processing.R -""" - -from typing import Dict import numpy as np -import pandas as pd - -# duckdb / scikit-learn imported lazily in the functions below (estimint[train]). - - -def load_and_filter( - in_parquet: str, - thr_lo: float = 0.02, - thr_hi: float = 0.95 -) -> Dict[str, pd.DataFrame]: - """ - Load parquet file and apply prevalence filters. - - Equivalent to R's load_and_filter() function. - - Parameters - ---------- - in_parquet : str - Path to input parquet file - thr_lo : float, optional - Lower prevalence threshold (inclusive, default: 0.02) - thr_hi : float, optional - Upper prevalence threshold (inclusive, default: 0.95) - - Returns - ------- - dict - Dictionary with keys: - - 'DT': DataFrame with rows passing filters - - 'DT_excluded': DataFrame with rows failing filters - """ - try: - import duckdb - except ImportError: - raise ImportError( - "load_and_filter() requires duckdb. " - "Install the training extras: pip install estimint[train]" - ) - - con = duckdb.connect(database=":memory:") - - try: - # Set DuckDB parameters - con.execute("PRAGMA threads=8; PRAGMA memory_limit='16GB';") - - qry = f""" - WITH base AS (SELECT * FROM read_parquet('{in_parquet}')), - avg_prev AS ( - SELECT parameter_index, - AVG(CASE WHEN year BETWEEN 1 AND 8 THEN prevalence_annual_mean END) AS prev_avg_1_8 - FROM base GROUP BY parameter_index - ), - y9 AS ( - SELECT b.parameter_index, - b.dn0_use, b.Q0, b.phi_bednets, b.seasonal, b.itn_use, b.irs_use, - b.prevalence_annual_mean AS prev_y9, - b.eir - FROM base b WHERE b.year = 9 - ) - SELECT y9.*, avg_prev.prev_avg_1_8 - FROM y9 JOIN avg_prev USING (parameter_index); - """ - - df_all = con.execute(qry).fetchdf() - - finally: - con.close() - - # Remove rows with any NaN - DT0 = df_all.dropna() - - # Apply prevalence filters - excluded_mask = ( - (DT0["prev_avg_1_8"] < thr_lo) | - (DT0["prev_avg_1_8"] > thr_hi) | - (DT0["prev_y9"] < thr_lo) | - (DT0["prev_y9"] > thr_hi) - ) - - DT_excluded = DT0[excluded_mask].copy().reset_index(drop=True) - DT = DT0[~excluded_mask].copy().reset_index(drop=True) - - return {"DT": DT, "DT_excluded": DT_excluded} - def make_value_weights(eir_raw: np.ndarray, digits: int = 3) -> np.ndarray: """ Create inverse-frequency weights based on EIR values. - + Equivalent to R's make_value_weights() function. - + Parameters ---------- eir_raw : array-like Raw EIR values digits : int, optional Number of digits for rounding (default: 3) - + Returns ------- np.ndarray @@ -111,90 +20,16 @@ def make_value_weights(eir_raw: np.ndarray, digits: int = 3) -> np.ndarray: """ eir_raw = np.asarray(eir_raw) key = np.round(eir_raw, digits) - + # Count frequency of each rounded value unique_vals, counts = np.unique(key, return_counts=True) freq_dict = dict(zip(unique_vals, counts)) - + # Inverse frequency weights w = np.array([1.0 / freq_dict[k] for k in key]) - + # Normalize to mean = 1 w = w / np.mean(w) - - return w - - -def strata_and_split( - DT: pd.DataFrame, - k_strata: int = 16, - seed: int = 42 -) -> pd.DataFrame: - """ - Create strata using k-means on log10(EIR) and perform stratified train/val/test split. - - Equivalent to R's strata_and_split() function. - - Parameters - ---------- - DT : pd.DataFrame - Input DataFrame (must have 'eir_log10' column) - k_strata : int, optional - Number of strata for k-means (default: 16) - seed : int, optional - Random seed (default: 42) - - Returns - ------- - pd.DataFrame - DataFrame with added 'strat_bin' and 'split' columns - """ - try: - from sklearn.cluster import KMeans - except ImportError: - raise ImportError( - "strata_and_split() requires scikit-learn. " - "Install the training extras: pip install estimint[train]" - ) - DT = DT.copy() - np.random.seed(seed) + return w - # K-means clustering on log10(EIR) - eir_log10 = DT["eir_log10"].values.reshape(-1, 1) - km = KMeans(n_clusters=k_strata, n_init=50, max_iter=5000, random_state=seed) - km.fit(eir_log10) - - # Reorder cluster IDs by center value (ascending) - centers = km.cluster_centers_.flatten() - ord_idx = np.argsort(centers) - id_map = {old_id: new_id + 1 for new_id, old_id in enumerate(ord_idx)} - - DT["strat_bin"] = np.array([id_map[c] for c in km.labels_]) - - # Initialize split column - DT["split"] = None - - # Stratified split within each bin - for b in sorted(DT["strat_bin"].unique()): - idx = DT[DT["strat_bin"] == b].index.tolist() - n_b = len(idx) - n_tr = int(np.floor(0.70 * n_b)) - n_val = int(np.floor(0.15 * n_b)) - - # Shuffle indices - np.random.shuffle(idx) - - # Assign splits - tr_idx = idx[:n_tr] if n_tr > 0 else [] - val_idx = idx[n_tr:n_tr + n_val] if n_val > 0 else [] - te_idx = idx[n_tr + n_val:] - - DT.loc[tr_idx, "split"] = "train" - DT.loc[val_idx, "split"] = "val" - DT.loc[te_idx, "split"] = "test" - - # Fill any remaining NaN splits with 'train' - DT["split"] = DT["split"].fillna("train") - - return DT diff --git a/src/estimint/eir_models.py b/src/estimint/eir_models.py new file mode 100644 index 0000000..4448c05 --- /dev/null +++ b/src/estimint/eir_models.py @@ -0,0 +1,85 @@ +"""Loading and inference for the estiMINT EIR models. + +Three conditional RQS flows are published on the Hugging Face hub. Each takes the +shared intervention covariates plus one measurement, and is keyed here by the name +of that measurement: + + prev_y9 -> eir baseline prevalence to EIR + hbr_y9 -> eir human biting rate to EIR + eir -> hbr_y9 EIR to human biting rate +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from estimint.v2.common.types import PredictorType, TargetType +from estimint.v2.models.rqs import ConditionalRQS, RQSArtifact + +from .types import Input_Mode, PreparedScenario + +ESTIMINT_HF_REPO = "dide-ic/estiMINT" + +# What each model predicts from the measurement it is named after. +EIR_MODEL_TARGETS: dict[PredictorType, TargetType] = { + "prev_y9": "eir", + "hbr_y9": "eir", + "eir": "hbr_y9", +} + +# The measurement a scenario supplies, per input mode. +INPUT_MODE_TO_PREDICTOR: dict[Input_Mode, PredictorType] = { + "prevalence": "prev_y9", + "hbr": "hbr_y9", + "eir": "eir", +} + +EirModels = dict[PredictorType, RQSArtifact] + +_MODEL_CACHE: dict[str, EirModels] = {} + + +def load_eir_models(hf_repo: str = ESTIMINT_HF_REPO) -> EirModels: + """Load the three estiMINT RQS artifacts, caching them per repo. + + Args: + hf_repo: HuggingFace repo ID (or local folder) holding the model artifacts. + + Returns: + The artifacts keyed by the measurement each one takes as input. + """ + if hf_repo not in _MODEL_CACHE: + _MODEL_CACHE[hf_repo] = { + predictor: ConditionalRQS.from_pretrained(hf_repo, predictor, target) + for predictor, target in EIR_MODEL_TARGETS.items() + } + return _MODEL_CACHE[hf_repo] + + +def predict_from_measurements( + eir_models: EirModels, + predictor: PredictorType, + prepared_scenarios: Sequence[PreparedScenario], + values: Sequence[float] | np.ndarray, +) -> np.ndarray: + """Predict ``EIR_MODEL_TARGETS[predictor]`` for a batch of scenarios. + + Each model row is the scenario's intervention covariates plus the supplied + measurement; the artifact reorders them into the training feature order. + + Args: + eir_models: Artifacts from :func:`load_eir_models`. + predictor: The measurement carried by *values*, e.g. ``"prev_y9"``. + prepared_scenarios: Scenarios supplying the intervention covariates. + values: One *predictor* measurement per scenario, in the same order. + + Returns: + Median predictions, one per scenario, in input order. + """ + records = [ + {**prepared_scenario.eir_model_features, predictor: float(value)} + for prepared_scenario, value in zip(prepared_scenarios, values, strict=True) + ] + return eir_models[predictor].predict(records) diff --git a/src/estimint/globals.py b/src/estimint/globals.py deleted file mode 100644 index e1a2215..0000000 --- a/src/estimint/globals.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Global variables and constants for estiMINT package. - -Equivalent to: globals.R - -In R, globalVariables() silences CMD check warnings for NSE columns. -In Python, we define these as module-level constants for documentation -and type-checking purposes. -""" - -from typing import List - -# Column names used throughout the package (equivalent to R's globalVariables) -COLUMN_NAMES: List[str] = [ - "row_id", - "true_value", - "case_range", - "model", - "prediction", - "error", - "true", - "pred", - "bin", - "year", -] - -# Feature column names -FEATURE_COLUMNS: List[str] = [ - "Feature", - "Gain_scaled", - "feature", - "importance_scaled", -] - -# Metric column names -METRIC_COLUMNS: List[str] = [ - "Model", - "Quantile", - "RMSE", -] - -# Default features for EIR prediction -DEFAULT_FEATURES: List[str] = [ - "dn0_use", - "Q0", - "phi_bednets", - "seasonal", - "itn_use", - "irs_use", - "prev_y9", -] - -# Default thresholds -DEFAULT_THR_LO: float = 0.02 -DEFAULT_THR_HI: float = 0.95 - -# Default k-means strata -DEFAULT_K_STRATA: int = 16 - -# Default CV folds -DEFAULT_K_FOLDS: int = 10 - -# Default random seed -DEFAULT_SEED: int = 42 diff --git a/src/estimint/hbr.py b/src/estimint/hbr.py index fa70135..efe20d5 100644 --- a/src/estimint/hbr.py +++ b/src/estimint/hbr.py @@ -4,117 +4,96 @@ Answers the question: "What happens to EIR if mosquito density changes by X%?" Pipeline: -1. prev_y9 + interventions -> EIR_baseline (prevalence model) -2. EIR_baseline + interventions -> HBR_baseline (EIR-to-HBR model) -3. HBR_new = HBR_baseline * (1 + mosquito_delta) (user's mosquito change, pos or neg) -4. HBR model predicts EIR at both HBR values (ratio approach) -5. EIR_new = EIR_baseline * (EIR_scaled / EIR_roundtrip) +1. prev_y9 + interventions -> EIR_baseline (prev_y9 -> eir model) +2. EIR_baseline + interventions -> HBR_baseline (eir -> hbr_y9 model) +3. HBR_new = HBR_baseline * (1 + mosquito_delta) (user's mosquito change, pos or neg) +4. EIR predicted at both HBR values (hbr_y9 -> eir model) +5. EIR_new = EIR_baseline * (EIR_new_raw / EIR_roundtrip) + +Step 5 applies the *ratio* of the two step-4 predictions rather than using EIR_new_raw +directly, so the result stays anchored to the cleaner step-1 baseline and any bias in +the EIR -> HBR -> EIR round trip cancels out. """ -from typing import Any +import numpy as np -import pandas as pd +from .eir_models import EirModels, predict_from_measurements +from .types import PreparedScenario -from .run import run_xgb_model - -def estimate_eir_with_mosquito_delta(inputs: pd.DataFrame, *, models: dict[str, Any]) -> pd.DataFrame: +def estimate_eir_with_mosquito_delta( + prepared_scenarios: list[PreparedScenario], *, eir_models: EirModels +) -> list[dict[str, float]]: """ - Estimate new EIR after a change in mosquito density for multiple scenarios. - - The HBR model predicts EIR at both the - baseline and scaled HBR, then applies the relative multiplier to the clean - baseline EIR from the prevalence model. + Estimate the new EIR after a change in mosquito density, for a batch of scenarios. Parameters ---------- - inputs : pd.DataFrame - One row per scenario. Required columns: - - - ``prevalence`` : baseline malaria prevalence (prev_y9), e.g. 0.30 for 30%. - - ``mosquito_delta`` : fractional change in mosquito density, e.g. 0.10 - for +10%, -0.50 for -50%. Must be > -1 per row. - - ``dn0_use`` : bednet contact reduction parameter. - - ``Q0`` : human blood index. - - ``phi_bednets`` : proportion of bites on humans while in bed. - - ``seasonal`` : seasonality flag (0.0 or 1.0). - - ``itn_use`` : ITN coverage (0-1). - - ``irs_use`` : IRS coverage (0-1). - - models : dict - Pre-loaded model dictionary with keys ``"prevalence"``, ``"hbr"``, - and ``"eir_to_hbr"``. + prepared_scenarios : list[PreparedScenario] + Scenarios with ``input_mode == "prevalence"``. Each supplies: + + - ``eir_target.input_value`` : baseline malaria prevalence (prev_y9), e.g. 0.30 for 30%. + - ``mosquito_density_change`` : fractional change in mosquito density, e.g. 0.10 + for +10%, -0.50 for -50%. Must be > -1. + - ``eir_model_features`` : the intervention covariates shared by all three models + (``dn0_use``, ``Q0``, ``phi_bednets``, ``seasonal``, ``itn_use``, ``irs_use``). + + eir_models : EirModels + Artifacts from :func:`estimint.eir_models.load_eir_models`. Returns ------- - pd.DataFrame - Same index as *inputs*, with columns: + list[dict[str, float]] + One entry per scenario, in input order, with keys: - ``eir_baseline`` : baseline EIR from prevalence. - - ``eir_new`` : new EIR after mosquito density change. + - ``eir_new`` : new EIR after the mosquito density change. - ``eir_multiplier`` : ratio of new EIR to baseline. - ``hbr_baseline`` : estimated baseline HBR. - - ``hbr_new`` : HBR after mosquito density change. + - ``hbr_new`` : HBR after the mosquito density change. Examples -------- - >>> import pandas as pd - >>> from estimint import estimate_eir_with_mosquito_delta - >>> inputs = pd.DataFrame([ - ... {"prevalence": 0.30, "mosquito_delta": 0.25, - ... "dn0_use": 0.33, "Q0": 0.87, "phi_bednets": 0.82, - ... "seasonal": 0.0, "itn_use": 0.6, "irs_use": 0.0}, - ... ]) - >>> result = estimate_eir_with_mosquito_delta(inputs, models=models) - >>> print(result[["eir_baseline", "eir_new"]]) + >>> from estimint.eir_models import load_eir_models + >>> from estimint.hbr import estimate_eir_with_mosquito_delta + >>> results = estimate_eir_with_mosquito_delta(prepared_scenarios, eir_models=load_eir_models()) + >>> results[0]["eir_new"] """ - features = [ - "dn0_use", - "Q0", - "phi_bednets", - "seasonal", - "itn_use", - "irs_use", - ] - intervention_data = inputs[features] - # Step 1: prevalence -> EIR baseline - prevalence_data = intervention_data.assign(prev_y9=inputs["prevalence"].to_numpy()) - eir_baseline = run_xgb_model(prevalence_data, models["prevalence"]) + prevalences = [prepared_scenario.eir_target.input_value for prepared_scenario in prepared_scenarios] + eir_baselines = predict_from_measurements(eir_models, "prev_y9", prepared_scenarios, prevalences) # Step 2: EIR -> HBR baseline - eir_data = intervention_data.assign(eir=eir_baseline) - hbr_baseline = run_xgb_model(eir_data, models["eir_to_hbr"]) + hbr_baselines = predict_from_measurements(eir_models, "eir", prepared_scenarios, eir_baselines) # Step 3: apply mosquito delta (positive or negative) - hbr_new = hbr_baseline * (1 + inputs["mosquito_delta"].to_numpy()) - - # Step 4: ratio approach — batch both HBR values in one call so they - # share the same smooth PCHIP curve - hbr_data = pd.concat( - [ - intervention_data.assign(hbr_y9=hbr_baseline), - intervention_data.assign(hbr_y9=hbr_new), - ], - ignore_index=True, + mosquito_deltas = np.array( + [prepared_scenario.mosquito_density_change for prepared_scenario in prepared_scenarios] ) - eir_from_hbr = run_xgb_model(hbr_data, models["hbr"]) - - count = len(inputs) - eir_rt = eir_from_hbr[:count] - eir_new_raw = eir_from_hbr[count:] + hbr_new = hbr_baselines * (1 + mosquito_deltas) + + # Step 4: both HBR values go back through the HBR -> EIR model in one batched call + eir_from_hbr = predict_from_measurements( + eir_models, + "hbr_y9", + [*prepared_scenarios, *prepared_scenarios], + np.concatenate([hbr_baselines, hbr_new]), + ) + eir_roundtrip, eir_new_raw = np.split(eir_from_hbr, 2) - # Step 5: multiplier applied to clean baseline - multiplier = eir_new_raw / eir_rt - eir_new = eir_baseline * multiplier + # Step 5: multiplier applied to the clean baseline + eir_multipliers = eir_new_raw / eir_roundtrip + eir_news = eir_baselines * eir_multipliers - return pd.DataFrame( + return [ { - "eir_baseline": eir_baseline, - "eir_new": eir_new, - "eir_multiplier": multiplier, - "hbr_baseline": hbr_baseline, - "hbr_new": hbr_new, - }, - index=inputs.index, - ) + "eir_baseline": float(eir_baseline), + "eir_new": float(eir_new), + "eir_multiplier": float(eir_multiplier), + "hbr_baseline": float(hbr_baseline), + "hbr_new": float(hbr_adjusted), + } + for eir_baseline, eir_new, eir_multiplier, hbr_baseline, hbr_adjusted in zip( + eir_baselines, eir_news, eir_multipliers, hbr_baselines, hbr_new, strict=True + ) + ] diff --git a/src/estimint/models.py b/src/estimint/models.py deleted file mode 100644 index c8850e1..0000000 --- a/src/estimint/models.py +++ /dev/null @@ -1,158 +0,0 @@ -""" -Model training functions for estiMINT package. - -Equivalent to: models.R -""" - -from typing import Dict, Any, Optional -import numpy as np -import xgboost as xgb -from numpy.typing import ArrayLike - - -def train_eir_xgboost( - X_train: np.ndarray, - y_train: np.ndarray, - X_val: Optional[np.ndarray] = None, - y_val: Optional[np.ndarray] = None, - tune_params: bool = True -) -> Dict[str, Any]: - """ - Train XGBoost model for EIR prediction. - - Equivalent to R's train_eir_xgboost() function. - - Parameters - ---------- - X_train : np.ndarray - Numeric matrix of training features - y_train : np.ndarray - Numeric vector of training targets - X_val : np.ndarray, optional - Numeric matrix of validation features (default: None) - y_val : np.ndarray, optional - Numeric vector of validation targets (default: None) - tune_params : bool, optional - Whether to tune hyperparameters (default: True) - - Returns - ------- - dict - Dictionary containing: - - 'model': trained XGBoost model - - 'params': best parameters - - 'nrounds': number of training rounds - - 'importance': feature importance DataFrame - - 'transform': function to transform target (log10(y + 1)) - - 'inverse_transform': inverse transform function (10^y - 1) - """ - X_train = np.asarray(X_train) - y_train = np.asarray(y_train) - - # Transform target - y_train_log = np.log10(y_train + 1) - dtrain = xgb.DMatrix(data=X_train, label=y_train_log) - - # Build watchlist - evals = [(dtrain, "train")] - if X_val is not None and y_val is not None: - X_val = np.asarray(X_val) - y_val = np.asarray(y_val) - dval = xgb.DMatrix(data=X_val, label=np.log10(y_val + 1)) - evals.append((dval, "eval")) - - base_params = { - "objective": "reg:squarederror", - "eval_metric": "rmse", - "eta": 0.05, - "max_depth": 4, - "min_child_weight": 5, - "subsample": 0.7, - "colsample_bytree": 0.7, - "gamma": 0.1, - "alpha": 0.1, - "lambda": 1.0, - "seed": 42, - } - - if tune_params: - best_rmse = float("inf") - best_params = base_params.copy() - best_nrounds = 100 - - for depth in [3, 4, 5]: - for eta in [0.01, 0.05, 0.1]: - for subsample in [0.6, 0.7, 0.8]: - params = base_params.copy() - params.update({ - "max_depth": depth, - "eta": eta, - "subsample": subsample, - }) - - cv_results = xgb.cv( - params=params, - dtrain=dtrain, - num_boost_round=500, - nfold=5, - early_stopping_rounds=20, - verbose_eval=False, - seed=42, - ) - - # Get best iteration metrics - best_iter = len(cv_results) - 1 - cv_rmse = cv_results["test-rmse-mean"].iloc[best_iter] - - if cv_rmse < best_rmse: - best_rmse = cv_rmse - best_params = params.copy() - best_nrounds = best_iter + 1 - - params = best_params - nrounds = best_nrounds - print( - f"Best XGBoost: depth={params['max_depth']}, " - f"eta={params['eta']:.3f}, subsample={params['subsample']:.2f}, " - f"nrounds={nrounds}, CV-RMSE={best_rmse:.4f}" - ) - else: - cv_results = xgb.cv( - params=base_params, - dtrain=dtrain, - num_boost_round=500, - nfold=5, - early_stopping_rounds=20, - verbose_eval=False, - seed=42, - ) - params = base_params - nrounds = len(cv_results) - - # Train final model - model = xgb.train( - params=params, - dtrain=dtrain, - num_boost_round=nrounds, - evals=evals, - verbose_eval=False, - ) - - # Get feature importance - importance = model.get_score(importance_type="gain") - - # Define transform functions - def transform(y: ArrayLike) -> np.ndarray: - return np.log10(np.asarray(y) + 1) - - def inverse_transform(y: ArrayLike) -> np.ndarray: - return np.power(10, np.asarray(y)) - 1 - - return { - "model": model, - "params": params, - "nrounds": nrounds, - "importance": importance, - "transform": transform, - "inverse_transform": inverse_transform, - } diff --git a/src/estimint/run.py b/src/estimint/run.py deleted file mode 100644 index 8944060..0000000 --- a/src/estimint/run.py +++ /dev/null @@ -1,244 +0,0 @@ -""" -Model inference functions for estiMINT package. - -Equivalent to: run.R -""" - -from typing import Optional, Dict, Any, Union -import numpy as np -import pandas as pd -import xgboost as xgb - -from .utils import predict_qmap_w - - -# Global model storage (equivalent to R's .GlobalEnv) -_global_model: Optional[Dict[str, Any]] = None - - -def set_global_model(model: Dict[str, Any]) -> None: - """ - Set the global estiMINT model. - - Parameters - ---------- - model : dict - An 'estiMINT_model' object - """ - global _global_model - _global_model = model - - -def get_global_model() -> Optional[Dict[str, Any]]: - """ - Get the global estiMINT model. - - Returns - ------- - dict or None - The global model if set, None otherwise - """ - return _global_model - - -def run_xgb_model( - new_data: Union[pd.DataFrame, Dict[str, Any]], - model: Optional[Dict[str, Any]] = None -) -> np.ndarray: - """ - Run XGBoost model with initial conditions. - - Equivalent to R's run_xgb_model() function. - - Parameters - ---------- - new_data : pd.DataFrame or dict - Data frame with columns: prevalence (or prev_y9), dn0_use, Q0, - phi_bednets, seasonal, itn_use, irs_use - model : dict, optional - An 'estiMINT_model' object; if None, tries global 'estiMINT_model' - - Returns - ------- - np.ndarray - Numeric array of calibrated EIR predictions - - Raises - ------ - ValueError - If no model is provided and no global model exists - ValueError - If required columns are missing from new_data - """ - # Get model - if model is None: - model = get_global_model() - if model is None: - raise ValueError( - "No model provided and 'estiMINT_model' not found in the global context. " - "Either pass a model or call set_global_model() first." - ) - - # Get required features - req = model["features"] - - # Convert to DataFrame if necessary - if isinstance(new_data, dict): - nd = pd.DataFrame(new_data) - else: - nd = new_data.copy() - - # Handle prevalence -> prev_y9 alias - if "prevalence" in nd.columns and "prev_y9" not in nd.columns: - nd["prev_y9"] = nd["prevalence"] - - # Check for missing columns - missing = set(req) - set(nd.columns) - if missing: - raise ValueError(f"Missing required columns: {', '.join(sorted(missing))}") - - # Extract feature matrix - X = nd[req].values.astype(np.float64) - - # Detect monotonic feature. For models with a monotonic input - # (prev_y9 or hbr_y9), we always predict via a dense internal sweep - # + PCHIP smoothing so that every query — even a single point — - # returns a value from a smooth, monotone curve rather than the raw - # XGBoost staircase. - # Fixed sweep ranges. prev_y9 is linear; hbr_y9 is log-spaced because - # HBR spans four orders of magnitude in the training data (8k–57M). - _MONO_SWEEP = { - "prev_y9": {"lo": 0.005, "hi": 0.80, "log": False}, - "hbr_y9": {"lo": 5000, "hi": 70_000_000, "log": True}, - } - mono_fidx = None - mono_cfg = None - for fname, cfg in _MONO_SWEEP.items(): - if fname in req: - mono_fidx = req.index(fname) - mono_cfg = cfg - break - - if mono_fidx is None: - # No monotonic feature (e.g. EIR-to-HBR model) — predict directly - return _predict_direct(X, model) - - # Group rows by unique intervention combo (all columns except the - # monotonic feature) so each group gets its own smooth curve. - other_cols = [i for i in range(X.shape[1]) if i != mono_fidx] - other_vals = X[:, other_cols] - - if len(X) == 1: - unique_combos = other_vals - combo_labels = np.array([0]) - else: - unique_combos, combo_labels = np.unique( - other_vals, axis=0, return_inverse=True - ) - - pred_final = np.empty(len(X), dtype=np.float64) - n_sweep = 501 - - for ci in range(len(unique_combos)): - mask = combo_labels == ci - query_fvals = X[mask, mono_fidx] - intv = unique_combos[ci] - - # Dense sweep with fixed range (same curve for every call) - if mono_cfg["log"]: - sweep_fvals = np.logspace( - np.log10(mono_cfg["lo"]), np.log10(mono_cfg["hi"]), n_sweep - ) - else: - sweep_fvals = np.linspace(mono_cfg["lo"], mono_cfg["hi"], n_sweep) - - # Build the sweep feature matrix (insert mono column at correct index) - X_sweep = np.tile(intv, (n_sweep, 1)) - X_sweep = np.insert(X_sweep, mono_fidx, sweep_fvals, axis=1) - - # Predict sweep → staircase → PCHIP smooth curve - sweep_preds = _predict_direct(X_sweep, model) - sweep_smooth = _smooth_staircase(sweep_fvals, sweep_preds) - - # Interpolate query points from the smooth curve - pred_final[mask] = np.interp(query_fvals, sweep_fvals, sweep_smooth) - - return pred_final.astype(np.float64) - - -def _predict_direct(X: np.ndarray, model: dict) -> np.ndarray: - """Raw model prediction + QMAP calibration (no smoothing).""" - dnew = xgb.DMatrix(X) - pred_log10 = model["booster"].predict(dnew) - pred_raw = np.power(10, pred_log10) - pred_cal = predict_qmap_w(pred_raw, model["calibrator"]["qmap"]) - return np.maximum(0, model["calibrator"]["scale"] * pred_cal).astype(np.float64) - - -def _smooth_staircase(fvals: np.ndarray, preds: np.ndarray) -> np.ndarray: - """Smooth a monotone staircase into a visually smooth monotone curve. - - Two-stage smoothing: - 1. PCHIP through staircase midpoint knots — removes hard edges, gives - C1-continuous monotone curve. - 2. Gaussian kernel smoothing — spreads steep PCHIP transitions over a - wider range so that 1% user steps always show gradual change, never - "nothing then big jump". - 3. Re-enforce strict monotonicity after Gaussian pass. - """ - from scipy.interpolate import PchipInterpolator - from scipy.ndimage import gaussian_filter1d - - order = np.argsort(fvals) - x = fvals[order] - y = preds[order] - - # Find constant segments (within relative tolerance) - segments = [] - i = 0 - while i < len(y): - j = i + 1 - ref = max(abs(y[i]), 1e-10) - while j < len(y) and abs(y[j] - y[i]) / ref < 1e-6: - j += 1 - segments.append((i, j, y[i])) - i = j - - if len(segments) <= 2: - return preds # too few segments to interpolate - - # Knot at the midpoint of each segment's feature range - knots_x = [(x[s[0]] + x[s[1] - 1]) / 2 for s in segments] - knots_y = [s[2] for s in segments] - - # Pin the boundary knots to the actual data range so PCHIP covers - # the full sweep without extrapolation - knots_x[0] = x[0] - knots_x[-1] = x[-1] - - knots_x = np.array(knots_x) - knots_y = np.array(knots_y) - - # Stage 1: PCHIP — smooth monotone cubic interpolation - pchip = PchipInterpolator(knots_x, knots_y) - y_smooth = pchip(x) - - # Stage 2: Gaussian smoothing in log-space to spread steep transitions. - # Work in log-space so the kernel acts on relative (multiplicative) - # changes rather than absolute — this prevents the Gaussian from - # under-smoothing at the low end and over-smoothing at the high end. - # sigma=12 on a 501-point sweep ≈ 2.4% of the range, enough to spread - # a sharp step across ~5% of the sweep (several user-facing 1% steps). - y_log = np.log(np.maximum(y_smooth, 1e-10)) - y_log_smooth = gaussian_filter1d(y_log, sigma=12, mode="nearest") - y_smooth = np.exp(y_log_smooth) - - # Stage 3: enforce strict monotonicity (Gaussian can create tiny dips) - for i in range(1, len(y_smooth)): - if y_smooth[i] <= y_smooth[i - 1]: - y_smooth[i] = y_smooth[i - 1] * (1 + 1e-10) - - # Restore original order - result = np.empty_like(preds) - result[order] = y_smooth - return result diff --git a/src/estimint/scenarios.py b/src/estimint/scenarios.py index d853d4b..1c3bbeb 100644 --- a/src/estimint/scenarios.py +++ b/src/estimint/scenarios.py @@ -1,26 +1,30 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Dict +from typing import Any import numpy as np import pandas as pd from .bednet import DN0Result, calculate_dn0 +from .eir_models import ( + ESTIMINT_HF_REPO, + INPUT_MODE_TO_PREDICTOR, + EirModels, + load_eir_models, + predict_from_measurements, +) from .hbr import estimate_eir_with_mosquito_delta -from .run import run_xgb_model -from .storage import load_xgb_model -from .types import EirTarget, Scenario +from .types import Input_Mode, PreparedScenario, Scenario from collections import defaultdict ####################### Constants and global storage ################### -HF_REPO = "dide-ic/stateMINT" +STATEMINT_HF_REPO = "dide-ic/stateMINT" # 157 windows of 14 days from day 2190; intervention at day 3285. _ABS_TIME = 2190 + 14 * np.arange(157) _IDX_Y9 = int(np.argmin(np.abs(_ABS_TIME - 3285))) -_EIR_MODEL_CACHE: Dict[str, Any] = {} -_EMULATOR_MODEL_CACHE: Dict[str, Dict[str, Any]] = {} +_EMULATOR_MODEL_CACHE: dict[str, dict[str, Any]] = {} _NET_KEYS = ( "py_only", @@ -29,30 +33,10 @@ "py_ppf", ) -_REQUIRED_EIR_MODEL_NAMES = ("prevalence", "hbr", "eir_to_hbr") - - -@dataclass(frozen=True) -class _EirInputModelConfig: - feature_column: str - model_name: str - - -_EIR_INPUT_MODEL_CONFIG = { - "prevalence": _EirInputModelConfig(feature_column="prev_y9", model_name="prevalence"), - "hbr": _EirInputModelConfig(feature_column="hbr_y9", model_name="hbr"), -} - ######################## Internal helpers ######################## -def _load_eir_hbr_models() -> Dict[str, Any]: - if not _EIR_MODEL_CACHE: - for model_name in _REQUIRED_EIR_MODEL_NAMES: - _EIR_MODEL_CACHE[model_name] = load_xgb_model(model_name) - return _EIR_MODEL_CACHE - -def _load_emulators(hf_repo: str) -> Dict[str, Any]: +def _load_emulators(hf_repo: str) -> dict[str, Any]: if hf_repo not in _EMULATOR_MODEL_CACHE: try: from stateMINT.model import Mamba2Regressor @@ -60,7 +44,7 @@ def _load_emulators(hf_repo: str) -> Dict[str, Any]: raise ImportError( "run_scenarios needs stateMINT. Install it with: " "uv sync --extra scenarios (or pip install " - '"git+https://github.com/mrc-ide/stateMINT.git@mamba2-train").' + '"mintstate>=0.3.0")' ) from error _EMULATOR_MODEL_CACHE[hf_repo] = { outcome_name: Mamba2Regressor.from_pretrained(hf_repo, predictor=outcome_name) @@ -94,17 +78,7 @@ def _calculate_bednet_effects(scenario: Scenario) -> _BedNetEffects: future=future_effect, ) - -@dataclass -class _PreparedScenario: - eir_target: EirTarget - mosquito_density_change: float - eir_model_features: Dict[str, float] - summary_values: dict[str, Any] - emulator_covariates: dict[str, float] - - -def _prepare_scenario_inputs(scenario: Scenario) -> _PreparedScenario: +def _prepare_scenario_inputs(scenario: Scenario) -> PreparedScenario: """Compute the model inputs and initial output values for one scenario.""" bednet_effects = _calculate_bednet_effects(scenario) current_dn0 = bednet_effects.current.dn0 @@ -161,7 +135,7 @@ def _prepare_scenario_inputs(scenario: Scenario) -> _PreparedScenario: "hbr_new": np.nan, } - return _PreparedScenario( + return PreparedScenario( eir_target=scenario.eir_target, mosquito_density_change=scenario.mosquito_delta, eir_model_features=eir_model_features, @@ -171,7 +145,7 @@ def _prepare_scenario_inputs(scenario: Scenario) -> _PreparedScenario: def _record_eir_estimate( - prepared_scenario: _PreparedScenario, + prepared_scenario: PreparedScenario, *, eir_baseline: float, eir_final: float, @@ -186,19 +160,15 @@ def _record_eir_estimate( def _predict_eir_from_measurements( - prepared_scenarios: list[_PreparedScenario], *, input_mode: str, eir_models: Dict[str, Any] + prepared_scenarios: list[PreparedScenario], *, input_mode: Input_Mode, eir_models: EirModels ) -> None: """Predict EIR from baseline prevalence or HBR measurements.""" - model_config = _EIR_INPUT_MODEL_CONFIG[input_mode] - model_input_records = [ - { - **prepared_scenario.eir_model_features, - model_config.feature_column: prepared_scenario.eir_target.input_value, - } - for prepared_scenario in prepared_scenarios - ] - - eir_predictions = run_xgb_model(pd.DataFrame(model_input_records), eir_models[model_config.model_name]) + eir_predictions = predict_from_measurements( + eir_models, + INPUT_MODE_TO_PREDICTOR[input_mode], + prepared_scenarios, + [prepared_scenario.eir_target.input_value for prepared_scenario in prepared_scenarios], + ) for prepared_scenario, eir_prediction in zip(prepared_scenarios, eir_predictions): _record_eir_estimate( @@ -208,7 +178,7 @@ def _predict_eir_from_measurements( ) -def _classify_prepared_scenario(prepared_scenario: _PreparedScenario) -> str: +def _classify_prepared_scenario(prepared_scenario: PreparedScenario) -> str: """Return the EIR estimation method for a prepared scenario.""" if prepared_scenario.eir_target.input_mode == "eir": return "eir" @@ -217,18 +187,10 @@ def _classify_prepared_scenario(prepared_scenario: _PreparedScenario) -> str: return prepared_scenario.eir_target.input_mode # "prevalence" or "hbr" -def _apply_mosquito_delta_batch(prepared_scenarios: list[_PreparedScenario], eir_models: Dict[str, Any]) -> None: - inputs = pd.DataFrame( - [ - { - "prevalence": prepared_scenario.eir_target.input_value, - "mosquito_delta": prepared_scenario.mosquito_density_change, - **prepared_scenario.eir_model_features, - } - for prepared_scenario in prepared_scenarios - ] - ) - estimates = estimate_eir_with_mosquito_delta(inputs, models=eir_models).to_dict(orient="records") +def _apply_mosquito_delta_batch(prepared_scenarios: list[PreparedScenario], eir_models: EirModels) -> None: + """Estimate EIR for a batch of scenarios with prevalence input and a mosquito-density change.""" + + estimates = estimate_eir_with_mosquito_delta(prepared_scenarios, eir_models=eir_models) for prepared_scenario, estimate in zip(prepared_scenarios, estimates): _record_eir_estimate( prepared_scenario, @@ -239,10 +201,10 @@ def _apply_mosquito_delta_batch(prepared_scenarios: list[_PreparedScenario], eir ) -def _estimate_eir(scenarios: list[Scenario], eir_models: Dict[str, Any]) -> list[_PreparedScenario]: +def _estimate_eir(scenarios: list[Scenario], eir_models: EirModels) -> list[PreparedScenario]: """Estimate EIR for many scenarios, dispatching each to one of three paths: - "eir": supplied directly, passed through unchanged - - "prevalence" / "hbr": predicted from baseline measurements via XGBoost + - "prevalence" / "hbr": predicted from baseline measurements by the matching RQS model - "mosquito_delta": prevalence input with a projected mosquito-density change """ if any(scenario.eir_target.input_mode not in {"prevalence", "eir", "hbr"} for scenario in scenarios): @@ -250,7 +212,7 @@ def _estimate_eir(scenarios: list[Scenario], eir_models: Dict[str, Any]) -> list prepared_scenarios = [_prepare_scenario_inputs(scenario) for scenario in scenarios] - scenario_groups: dict[str, list[_PreparedScenario]] = defaultdict(list) + scenario_groups: dict[str, list[PreparedScenario]] = defaultdict(list) for prepared_scenario in prepared_scenarios: scenario_groups[_classify_prepared_scenario(prepared_scenario)].append(prepared_scenario) @@ -269,18 +231,16 @@ def _estimate_eir(scenarios: list[Scenario], eir_models: Dict[str, Any]) -> list ######################### Public API ######################## -def preload_models(*, hf_repo: str = HF_REPO) -> tuple[Dict[str, Any], Dict[str, Any]]: +def preload_models(*, statemint_hf_repo: str = STATEMINT_HF_REPO, estimint_hf_repo: str = ESTIMINT_HF_REPO) -> tuple[EirModels, dict[str, Any]]: """Preload the models used by run_scenarios.""" - eir_models = _load_eir_hbr_models() - emulator_models = _load_emulators(hf_repo) - - return eir_models, emulator_models + return load_eir_models(estimint_hf_repo), _load_emulators(statemint_hf_repo) def run_scenarios( scenarios: list[Scenario], *, - hf_repo: str = HF_REPO, + statemint_hf_repo: str = STATEMINT_HF_REPO, + estimint_hf_repo: str = ESTIMINT_HF_REPO ) -> pd.DataFrame: """Run a list of scenarios through the estiMINT -> stateMINT pipeline. @@ -293,8 +253,10 @@ def run_scenarios( scenarios: Scenarios to evaluate. Each ``Scenario`` describes intervention coverages (ITN, IRS, LSM, etc.) and an ``EirTarget`` specifying the baseline transmission intensity. - hf_repo: HuggingFace repo ID from which emulator model weights are - downloaded. Defaults to the package-level ``HF_REPO`` constant. + statemint_hf_repo: HuggingFace repo ID from which stateMINT emulator model weights are + downloaded. Defaults to the package-level ``STATEMINT_HF_REPO`` constant. + estimint_hf_repo: HuggingFace repo ID from which estiMINT model weights are + downloaded. Defaults to the package-level ``ESTIMINT_HF_REPO`` constant. Returns: A ``pd.DataFrame`` with one row per scenario containing: @@ -349,7 +311,7 @@ def run_scenarios( if not scenarios: return pd.DataFrame() - eir_models, emulator_models = preload_models(hf_repo=hf_repo) + eir_models, emulator_models = preload_models(statemint_hf_repo=statemint_hf_repo, estimint_hf_repo=estimint_hf_repo) scenario_estimates = _estimate_eir(scenarios, eir_models) emulator_covariates = [estimate.emulator_covariates for estimate in scenario_estimates] diff --git a/src/estimint/storage.py b/src/estimint/storage.py deleted file mode 100644 index 0e4a2f6..0000000 --- a/src/estimint/storage.py +++ /dev/null @@ -1,717 +0,0 @@ -""" -Model storage and persistence functions for estiMINT package. - -Equivalent to: storage.R -""" - -import os -import json -import hashlib -import pickle -import zipfile -import tempfile -import warnings -from datetime import datetime -from pathlib import Path -from typing import Optional, Dict, Any, Union - -import numpy as np -import xgboost as xgb - - -def _get_package_data_dir() -> Path: - """Get the data directory inside the installed package.""" - return Path(__file__).parent / "data" - - -def _get_package_inst_dir() -> Path: - """Get the inst directory inside the installed package.""" - return Path(__file__).parent / "inst" - - -def _model_repo() -> str: - """ - Get the GitHub repository for model storage. - - Equivalent to R's .model_repo() function. - - Returns - ------- - str - Repository name in 'owner/repo' format - """ - return "CosmoNaught/estiMINT" - - -def _model_cache_dir() -> Path: - """ - Get the user cache directory for models. - - Equivalent to R's .model_cache_dir() function. - - Returns - ------- - Path - Path to cache directory - """ - try: - import appdirs - cache_dir = Path(appdirs.user_cache_dir("estiMINT")) - except ImportError: - # Fallback to home directory - cache_dir = Path.home() / ".cache" / "estiMINT" - - cache_dir.mkdir(parents=True, exist_ok=True) - return cache_dir - - -def _models_tag() -> str: - """ - Get the current models tag from package data. - - Equivalent to R's .models_tag() function. - - Returns - ------- - str - Model tag string - - Raises - ------ - FileNotFoundError - If models-tag.txt is not found - """ - # Try package inst/ location - inst_path = _get_package_inst_dir() / "models-tag.txt" - if inst_path.exists(): - return inst_path.read_text().strip() - - # Try importlib.resources - try: - import importlib.resources as pkg_resources - try: - with pkg_resources.files("estimint").joinpath("inst/models-tag.txt").open() as f: - return f.read().strip() - except (TypeError, FileNotFoundError): - pass - try: - with pkg_resources.files("estimint").joinpath("models-tag.txt").open() as f: - return f.read().strip() - except (TypeError, FileNotFoundError): - pass - except ImportError: - pass - - # Try local development location - for local_path in [Path("inst/models-tag.txt"), Path("src/estimint/inst/models-tag.txt")]: - if local_path.exists(): - return local_path.read_text().strip() - - raise FileNotFoundError( - "models-tag.txt missing. Publish a model and ship the tag." - ) - - -def _models_checksums() -> Optional[Dict[str, Any]]: - """ - Get model checksums from package data. - - Equivalent to R's .models_checksums() function. - - Returns - ------- - dict or None - Dictionary with 'path' and 'md5' keys, or None if not found - """ - import csv - - # Try package inst/ location - inst_path = _get_package_inst_dir() / "models-checksums.csv" - if inst_path.exists(): - with open(inst_path) as f: - reader = csv.DictReader(f) - rows = list(reader) - return rows if rows else None - - # Try importlib.resources - try: - import importlib.resources as pkg_resources - try: - with pkg_resources.files("estimint").joinpath("inst/models-checksums.csv").open() as f: - reader = csv.DictReader(f) - rows = list(reader) - return rows if rows else None - except (TypeError, FileNotFoundError): - pass - except ImportError: - pass - - # Try local development locations - for local_path in [Path("inst/models-checksums.csv"), Path("src/estimint/inst/models-checksums.csv")]: - if local_path.exists(): - with open(local_path) as f: - reader = csv.DictReader(f) - rows = list(reader) - return rows if rows else None - - return None - - -def _model_root(tag: Optional[str] = None) -> Path: - """ - Get the root directory for models. - - Equivalent to R's .model_root() function. - - Parameters - ---------- - tag : str, optional - Model tag (default: from _models_tag()) - - Returns - ------- - Path - Path to model root directory - """ - if tag is None: - tag = _models_tag() - - # Check for override environment variable - override = os.environ.get("ESTIMINT_MODELS_DIR", "") - if override: - path = Path(override) - if not path.exists(): - raise FileNotFoundError(f"ESTIMINT_MODELS_DIR does not exist: {override}") - return path - - return _model_cache_dir() / "models" / tag - - -def _ensure_models(tag: Optional[str] = None) -> Path: - """ - Ensure models are downloaded and verified. - - Equivalent to R's .ensure_models() function. - - Parameters - ---------- - tag : str, optional - Model tag (default: from _models_tag()) - - Returns - ------- - Path - Path to model root directory - - Raises - ------ - ImportError - If required packages are not installed - RuntimeError - If model checksum verification fails - """ - if tag is None: - tag = _models_tag() - - root = _model_root(tag) - ok_marker = root / ".ok" - - if ok_marker.exists(): - return root - - try: - import requests - except ImportError: - raise ImportError( - "Please install 'requests' to download published models: " - "pip install requests" - ) - - root.mkdir(parents=True, exist_ok=True) - - # Download from GitHub releases - repo = _model_repo() - zip_url = f"https://github.com/{repo}/releases/download/{tag}/{tag}.zip" - - with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_file: - tmp_zip = tmp_file.name - - response = requests.get(zip_url, stream=True) - response.raise_for_status() - - for chunk in response.iter_content(chunk_size=8192): - tmp_file.write(chunk) - - # Extract zip - with zipfile.ZipFile(tmp_zip, "r") as zip_ref: - zip_ref.extractall(root) - - os.unlink(tmp_zip) - - # Verify checksums - checksums = _models_checksums() - if checksums: - for entry in checksums: - file_path = root / entry["path"] - expected_md5 = entry["md5"] - - if not file_path.exists(): - raise RuntimeError(f"Model checksum verification failed: missing file {entry['path']}") - - # Calculate MD5 - with open(file_path, "rb") as f: - actual_md5 = hashlib.md5(f.read()).hexdigest() - - if actual_md5 != expected_md5: - raise RuntimeError( - f"Model checksum verification failed for {entry['path']}: " - f"expected {expected_md5}, got {actual_md5}" - ) - - # Mark as complete - ok_marker.touch() - - return root - - -_BUNDLED_MODELS = { - "prevalence": "estiMINT_model.pkl", - "hbr": "estiMINT_HBR_model.pkl", - "eir_to_hbr": "estiMINT_EIR_to_HBR_model.pkl", -} - - -def _find_installed_model(name: Optional[str] = None) -> Optional[str]: - """ - Find model files installed with package. - - Parameters - ---------- - name : str, optional - Model name: "prevalence" (default), "hbr", or "eir_to_hbr". - If None, returns the default prevalence model. - - Returns - ------- - str or None - Path to model directory if found, None otherwise - """ - data_dir = _get_package_data_dir() - - # Named model lookup - if name is not None: - filename = _BUNDLED_MODELS.get(name) - if filename is None: - raise ValueError( - f"Unknown model name '{name}'. " - f"Available: {', '.join(sorted(_BUNDLED_MODELS))}" - ) - path = data_dir / filename - if path.exists(): - return str(path) - return None - - # Default: check for JSON format (exported from R) - if (data_dir / "estiMINT_booster.json").exists(): - return str(data_dir) - - # Check for pickle format - if (data_dir / "estiMINT_model.pkl").exists(): - return str(data_dir / "estiMINT_model.pkl") - - return None - - -def _resolve_model_file(dir_or_file: Union[str, Path]) -> str: - """ - Resolve model file from directory or file path. - - Parameters - ---------- - dir_or_file : str or Path - Path to directory or file - - Returns - ------- - str - Path to model file or directory - - Raises - ------ - FileNotFoundError - If model file cannot be found - """ - path = Path(dir_or_file) - - # If it's a file that exists, return it - if path.is_file(): - return str(path) - - # Must be a directory - if not path.is_dir(): - raise FileNotFoundError(f"Path does not exist: {dir_or_file}") - - # Check for JSON format (new format from R export) - if (path / "estiMINT_booster.json").exists(): - return str(path) - - # Try candidate locations for pickle - candidates = [ - path / "estiMINT_model.pkl", - path / "eir_model" / "estiMINT_model.pkl", - ] - - for cand in candidates: - if cand.is_file(): - return str(cand) - - # Search recursively - for pattern in ["estiMINT_model.pkl", "estiMINT_booster.json"]: - hits = list(path.rglob(pattern)) - if hits: - if pattern.endswith(".json"): - return str(hits[0].parent) - return str(hits[0]) - - raise FileNotFoundError(f"Could not find estiMINT model under: {dir_or_file}") - - -def _load_from_json_dir(model_dir: Path) -> Dict[str, Any]: - """ - Load model from JSON files exported from R. - - Parameters - ---------- - model_dir : Path - Directory containing estiMINT_booster.json, estiMINT_calibrator.json, - estiMINT_metadata.json - - Returns - ------- - dict - estiMINT_model object - """ - booster_path = model_dir / "estiMINT_booster.json" - calibrator_path = model_dir / "estiMINT_calibrator.json" - metadata_path = model_dir / "estiMINT_metadata.json" - - # Load XGBoost booster - booster = xgb.Booster() - booster.load_model(str(booster_path)) - - # Load calibrator - with open(calibrator_path, "r") as f: - cal_data = json.load(f) - - calibrator = { - "kind": cal_data["kind"], - "qmap": { - "xq": np.array(cal_data["qmap"]["xq"]), - "yq": np.array(cal_data["qmap"]["yq"]), - }, - "scale": cal_data["scale"] - } - - # Load metadata - with open(metadata_path, "r") as f: - metadata = json.load(f) - - # Build model bundle - model_bundle = { - "class": metadata.get("class", "estiMINT_model"), - "booster": booster, - "calibrator": calibrator, - "features": metadata["features"], - "best_nrounds": metadata.get("best_nrounds"), - "preprocess": metadata.get("preprocess", {}), - } - - return model_bundle - - -def bundle_model( - model_path: Union[str, Path], - pkg_root: Optional[Union[str, Path]] = None, -) -> Path: - """ - Bundle a trained model into the package for distribution. - - This copies the model file into the package's data directory so that - `load_xgb_model()` can find it without any arguments. - - Parameters - ---------- - model_path : str or Path - Path to estiMINT_model.pkl file or directory containing it - pkg_root : str or Path, optional - Package root directory. If None, auto-detects from this file's location. - - Returns - ------- - Path - Path to the bundled model file - - Examples - -------- - >>> from estimint import bundle_model - >>> bundle_model("/path/to/estiMINT_model.pkl") - >>> # Now load_xgb_model() works without arguments - >>> model = load_xgb_model() - """ - import shutil - - model_path = Path(model_path) - - # Resolve to actual .pkl file - if model_path.is_dir(): - # Search for the model file - candidates = [ - model_path / "estiMINT_model.pkl", - model_path / "models" / "estiMINT_model.pkl", - ] - found = None - for c in candidates: - if c.exists(): - found = c - break - if found is None: - # Try recursive search - hits = list(model_path.rglob("estiMINT_model.pkl")) - if hits: - found = hits[0] - if found is None: - raise FileNotFoundError(f"Could not find estiMINT_model.pkl in {model_path}") - model_path = found - - if not model_path.exists(): - raise FileNotFoundError(f"Model file not found: {model_path}") - - if not model_path.suffix == ".pkl": - raise ValueError(f"Expected .pkl file, got: {model_path}") - - # Determine package root - if pkg_root is None: - # Auto-detect: this file is in src/estimint/storage.py - # Package root is 3 levels up - pkg_root = Path(__file__).parent.parent.parent - else: - pkg_root = Path(pkg_root) - - # Determine data directory - # Check if we're in src layout or flat layout - if (pkg_root / "src" / "estimint").is_dir(): - data_dir = pkg_root / "src" / "estimint" / "data" - elif (pkg_root / "estimint").is_dir(): - data_dir = pkg_root / "estimint" / "data" - else: - # Assume we're inside the package itself - data_dir = Path(__file__).parent / "data" - - # Create data directory - data_dir.mkdir(parents=True, exist_ok=True) - - # Copy model file - dest = data_dir / "estiMINT_model.pkl" - shutil.copy2(model_path, dest) - - # Compute checksum - with open(dest, "rb") as f: - md5 = hashlib.md5(f.read()).hexdigest() - - # Write checksum file - checksum_file = data_dir / "model_checksum.txt" - checksum_file.write_text(f"{md5} estiMINT_model.pkl\n") - - print(f"✓ Model bundled successfully!") - print(f" Source: {model_path}") - print(f" Destination: {dest}") - print(f" MD5: {md5}") - print(f"\nNow reinstall the package:") - print(f" cd {pkg_root} && pip install -e .") - print(f"\nThen load_xgb_model() will work without arguments.") - - return dest - - -def save_xgb_model( - model_dir: Union[str, Path], - tag: Optional[str] = None, - pkg_root: Union[str, Path] = ".", - repo: Optional[str] = None, - overwrite: bool = True, - wait_seconds: int = 90 -) -> str: - """ - Save XGBoost model to GitHub releases. - - Equivalent to R's save_xgb_model() function. - - Parameters - ---------- - model_dir : str or Path - Directory containing the model - tag : str, optional - Release tag (auto-generated if None) - pkg_root : str or Path, optional - Package root directory (default: ".") - repo : str, optional - GitHub repository (default: from _model_repo()) - overwrite : bool, optional - Whether to overwrite existing files (default: True) - wait_seconds : int, optional - Seconds to wait for release creation (default: 90) - - Returns - ------- - str - The tag used for the release - - Raises - ------ - FileNotFoundError - If model_dir does not exist - ImportError - If required packages are not installed - """ - model_dir = Path(model_dir) - if not model_dir.exists(): - raise FileNotFoundError(f"model_dir does not exist: {model_dir}") - - resolved = Path(_resolve_model_file(model_dir)) - - if repo is None: - repo = _model_repo() - - pkg_root = Path(pkg_root).resolve() - - # Generate tag if not provided - if tag is None: - # Hash the booster file - if resolved.is_dir(): - hash_file = resolved / "estiMINT_booster.json" - else: - hash_file = resolved - - with open(hash_file, "rb") as f: - md5 = hashlib.md5(f.read()).hexdigest() - short_md5 = md5[:8] - timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") - tag = f"models-{timestamp}-{short_md5}" - - # Write inst/ files - inst_dir = pkg_root / "inst" - inst_dir.mkdir(exist_ok=True) - - (inst_dir / "models-tag.txt").write_text(tag + "\n") - - print( - f"Model packaged under tag '{tag}'. " - f"To publish, upload to GitHub releases. " - f"Commit & reinstall to ship updated 'inst/models-*'." - ) - - return tag - - -def load_xgb_model(path: Optional[Union[str, Path]] = None) -> Dict[str, Any]: - """ - Load model onto memory for usage. - - Equivalent to R's load_xgb_model() function. - - Parameters - ---------- - path : str or Path, optional - Path to model directory (containing JSON files) or .pkl file. - Can also be a model name: "prevalence" (default), "hbr", or - "eir_to_hbr" to load bundled models by name. - If None, tries (in order): ESTIMINT_MODELS_DIR env var, - package data directory, then download using models-tag.txt. - - Returns - ------- - dict - An 'estiMINT_model' object (dictionary with model components) - - Raises - ------ - FileNotFoundError - If model file cannot be found - - Examples - -------- - >>> prev_model = load_xgb_model() # default prevalence model - >>> prev_model = load_xgb_model("prevalence") # same as above - >>> hbr_model = load_xgb_model("hbr") # HBR -> EIR model - >>> e2h_model = load_xgb_model("eir_to_hbr") # EIR -> HBR model - """ - # 0) Check for named model shortcut - if isinstance(path, str) and path in _BUNDLED_MODELS: - inst = _find_installed_model(name=path) - if inst is not None: - with open(inst, "rb") as f: - return pickle.load(f) - raise FileNotFoundError( - f"Bundled model '{path}' not found. " - f"Expected at: {_get_package_data_dir() / _BUNDLED_MODELS[path]}" - ) - - # 1) Explicit path - if path is not None: - path = Path(path) - resolved = Path(_resolve_model_file(path)) - - # Check if it's a directory with JSON files - if resolved.is_dir() and (resolved / "estiMINT_booster.json").exists(): - return _load_from_json_dir(resolved) - - # Otherwise assume pickle - if resolved.is_file() and resolved.suffix == ".pkl": - with open(resolved, "rb") as f: - obj = pickle.load(f) - valid_classes = {"estiMINT_model", "estiMINT_HBR_model", "estiMINT_EIR_to_HBR_model"} - if not isinstance(obj, dict) or obj.get("class") not in valid_classes: - warnings.warn("Loaded object does not appear to be an 'estiMINT_model'") - return obj - - raise FileNotFoundError(f"Could not load model from: {path}") - - # 2) Check environment variable override - override = os.environ.get("ESTIMINT_MODELS_DIR", "") - if override: - resolved = Path(_resolve_model_file(override)) - if resolved.is_dir() and (resolved / "estiMINT_booster.json").exists(): - return _load_from_json_dir(resolved) - if resolved.is_file(): - with open(resolved, "rb") as f: - return pickle.load(f) - - # 3) Check for installed model in package data/ - inst = _find_installed_model() - if inst is not None: - inst_path = Path(inst) - if inst_path.is_dir() and (inst_path / "estiMINT_booster.json").exists(): - return _load_from_json_dir(inst_path) - if inst_path.is_file(): - with open(inst_path, "rb") as f: - return pickle.load(f) - - # 4) Download from GitHub releases - tag = _models_tag() - root = _model_root(tag) - - if not (root / ".ok").exists(): - _ensure_models(tag) - - resolved = Path(_resolve_model_file(root)) - if resolved.is_dir() and (resolved / "estiMINT_booster.json").exists(): - return _load_from_json_dir(resolved) - - with open(resolved, "rb") as f: - obj = pickle.load(f) - - valid_classes = {"estiMINT_model", "estiMINT_HBR_model", "estiMINT_EIR_to_HBR_model"} - if not isinstance(obj, dict) or obj.get("class") not in valid_classes: - warnings.warn("Loaded object does not appear to be an 'estiMINT_model'") - - return obj diff --git a/src/estimint/train.py b/src/estimint/train.py deleted file mode 100644 index 447c74b..0000000 --- a/src/estimint/train.py +++ /dev/null @@ -1,433 +0,0 @@ -""" -Main training pipeline for estiMINT package. - -Equivalent to: train.R -""" - -import pickle -from pathlib import Path -from typing import Dict, Any, Optional - -import numpy as np -import pandas as pd -import xgboost as xgb - -from .utils import ( - ts, - r2, - rmse, - mse, - mae, - median_ae, - mae_rel, - rmsle, - safe_div, - smape, - fit_qmap_w, - predict_qmap_w, - scale_pos, -) -from .data_processing import ( - load_and_filter, - make_value_weights, - strata_and_split, -) -from .plotting import plot_obs_pred - - -def train_xgb_model( - in_parquet: str, - out_dir: str, - thr_lo: float = 0.02, - thr_hi: float = 0.95, - k_strata: int = 16, - K: int = 10, - seed: int = 42, - xgb_params: Optional[Dict[str, Any]] = None, - nrounds_max: int = 5000, - early_stopping_rounds: int = 100, - save_pkl: bool = True, - export_onnx: bool = False, - save_plots: bool = True, - save_artifacts: bool = True -) -> Dict[str, Any]: - """ - Train XGBoost with K-fold CV, QMAP+scale calibration, and optional artifacts. - - Equivalent to R's train_xgb_model() function. - - Parameters - ---------- - in_parquet : str - Path to input parquet file - out_dir : str - Base output directory (models/plots/metrics/predictions will be created) - thr_lo : float, optional - Lower prevalence filter, inclusive (default: 0.02) - thr_hi : float, optional - Upper prevalence filter, inclusive (default: 0.95) - k_strata : int, optional - Number of strata for k-means on log10(EIR) (default: 16) - K : int, optional - Number of CV folds (default: 10) - seed : int, optional - Random seed for reproducibility (default: 42) - xgb_params : dict, optional - XGBoost parameters (default: see function body) - nrounds_max : int, optional - Max rounds per fold for early-stopped training (default: 5000) - early_stopping_rounds : int, optional - Early stopping patience (default: 100) - save_pkl : bool, optional - Save a pickle bundle with model, calibrator, metadata (default: True) - export_onnx : bool, optional - Attempt ONNX export (default: False) - save_plots : bool, optional - Save diagnostic plots (default: True) - save_artifacts : bool, optional - Save CSV metrics and fold stats (default: True) - - Returns - ------- - dict - An 'estiMINT_model' object with booster, calibrator, features, metadata - """ - # Validate inputs - assert isinstance(in_parquet, str), "in_parquet must be a string" - assert isinstance(out_dir, str), "out_dir must be a string" - - # Default XGBoost parameters - if xgb_params is None: - xgb_params = { - "objective": "reg:squarederror", - "eval_metric": "rmse", - "tree_method": "hist", - "max_depth": 6, - "eta": 0.05, - "subsample": 0.8, - "colsample_bytree": 0.8, - "min_child_weight": 1.0, - "lambda": 1.0, - "seed": seed, - } - - # Create output directories - out_dir = Path(out_dir) - dir_models = out_dir / "models" - dir_plots = out_dir / "plots" - dir_metric = out_dir / "metrics" - dir_pred = out_dir / "predictions" - - for d in [dir_models, dir_plots, dir_metric, dir_pred]: - d.mkdir(parents=True, exist_ok=True) - - # Load and filter data - ts("Reading parquet & applying prevalence filters ...") - lf = load_and_filter(in_parquet, thr_lo=thr_lo, thr_hi=thr_hi) - DT = lf["DT"] - DT_excluded = lf["DT_excluded"] - - # Save excluded and kept data - DT_excluded.to_csv( - dir_metric / f"excluded_prev_outside_0p{int(thr_lo*100):02d}_0p{int(thr_hi*100):02d}.csv", - index=False - ) - DT.to_csv( - dir_metric / f"kept_after_prev_filters_0p{int(thr_lo*100):02d}_0p{int(thr_hi*100):02d}.csv", - index=False - ) - - # Define features - features = ["dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use", "prev_y9"] - - # Transform EIR to log10 - DT["eir_log10"] = np.log10(DT["eir"]) - - assert len(DT) > 0, "No data remaining after filters" - - # Create strata and split - np.random.seed(seed) - ts("Creating %d strata on log10(EIR) and 70/15/15 split ...", k_strata) - DT = strata_and_split(DT, k_strata=k_strata, seed=seed) - - # Hold-out test set - DT_test = DT[DT["split"] == "test"] - X_test = DT_test[features].values.astype(np.float64) - y_test = DT_test["eir_log10"].values - obs_eir_test = np.power(10, y_test) - - assert np.all(np.isfinite(X_test)), "X_test contains non-finite values" - assert np.all(np.isfinite(y_test)), "y_test contains non-finite values" - - # CV folds on train+val - ts("Assigning %d-fold CV within TRAIN+VAL strata ...", K) - DTcv = DT[DT["split"] != "test"].copy() - - np.random.seed(seed + 1) - - # Assign folds within each stratum - DTcv["fold"] = 0 - for b in DTcv["strat_bin"].unique(): - mask = DTcv["strat_bin"] == b - n_b = mask.sum() - idx = DTcv.index[mask].tolist() - np.random.shuffle(idx) - folds = np.tile(np.arange(1, K + 1), int(np.ceil(n_b / K)))[:n_b] - np.random.shuffle(folds) - DTcv.loc[idx, "fold"] = folds - - # K-fold CV training - ts("Running %d-fold CV with early stopping ...", K) - oof_pred_raw = np.full(len(DTcv), np.nan) - best_iters = np.zeros(K, dtype=int) - - for k in range(1, K + 1): - ts(" Fold %d / %d", k, K) - - idx_val = DTcv["fold"] == k - idx_tr = DTcv["fold"] != k - - X_tr = DTcv.loc[idx_tr, features].values.astype(np.float64) - y_tr = DTcv.loc[idx_tr, "eir_log10"].values - X_va = DTcv.loc[idx_val, features].values.astype(np.float64) - y_va = DTcv.loc[idx_val, "eir_log10"].values - - # Compute weights - w_tr = make_value_weights(np.power(10, y_tr), digits=3) - w_va = make_value_weights(np.power(10, y_va), digits=3) - - # Create DMatrix objects - dtr = xgb.DMatrix(X_tr, label=y_tr, weight=w_tr) - dva = xgb.DMatrix(X_va, label=y_va, weight=w_va) - - # Train model - mdl = xgb.train( - params=xgb_params, - dtrain=dtr, - num_boost_round=nrounds_max, - evals=[(dtr, "train"), (dva, "val")], - early_stopping_rounds=early_stopping_rounds, - verbose_eval=False, - ) - - best_iters[k - 1] = mdl.best_iteration - - # Predict on validation fold - pred_log10_va = mdl.predict(dva) - oof_pred_raw[idx_val.values] = np.power(10, pred_log10_va) - - assert np.all(np.isfinite(oof_pred_raw)), "OOF predictions contain non-finite values" - obs_cv_raw = np.power(10, DTcv["eir_log10"].values) - - # Save fold statistics - if save_artifacts: - fold_stats = pd.DataFrame({ - "fold": np.arange(1, K + 1), - "best_iteration": best_iters - }) - fold_stats.to_csv(dir_metric / f"cv_fold_best_iterations_K{K}.csv", index=False) - - # Fit calibrator on OOF predictions - ts("Fitting final calibrator (QMAP + positive scale) on OOF ...") - cal_oof = fit_qmap_w(oof_pred_raw, obs_cv_raw, ngrid=1024, round_digits=8) - oof_pred_cal = predict_qmap_w(oof_pred_raw, cal_oof) - a_oof = scale_pos(obs_cv_raw, oof_pred_cal) - oof_pred_final = np.maximum(0, a_oof * oof_pred_cal) - - # Save OOF metrics - if save_artifacts: - oof_metrics = pd.DataFrame({ - "set": ["OOF_uncalibrated", "OOF_calibrated"], - "R2": [r2(obs_cv_raw, oof_pred_raw), r2(obs_cv_raw, oof_pred_final)], - "bias": [np.mean(oof_pred_raw - obs_cv_raw), np.mean(oof_pred_final - obs_cv_raw)], - "MSE": [mse(obs_cv_raw, oof_pred_raw), mse(obs_cv_raw, oof_pred_final)], - "RMSE": [rmse(obs_cv_raw, oof_pred_raw), rmse(obs_cv_raw, oof_pred_final)], - "MAE": [mae(obs_cv_raw, oof_pred_raw), mae(obs_cv_raw, oof_pred_final)], - "MedianAE": [median_ae(obs_cv_raw, oof_pred_raw), median_ae(obs_cv_raw, oof_pred_final)], - "MAE_rel": [mae_rel(obs_cv_raw, oof_pred_raw), mae_rel(obs_cv_raw, oof_pred_final)], - "RMSLE": [rmsle(obs_cv_raw, oof_pred_raw), rmsle(obs_cv_raw, oof_pred_final)], - "NRMSE_mean": [ - safe_div(rmse(obs_cv_raw, oof_pred_raw), np.mean(obs_cv_raw)), - safe_div(rmse(obs_cv_raw, oof_pred_final), np.mean(obs_cv_raw)) - ], - "RelRMSE_p1": [ - np.sqrt(np.mean(safe_div(oof_pred_raw - obs_cv_raw, np.maximum(1, obs_cv_raw)) ** 2)), - np.sqrt(np.mean(safe_div(oof_pred_final - obs_cv_raw, np.maximum(1, obs_cv_raw)) ** 2)) - ], - "sMAPE": [smape(obs_cv_raw, oof_pred_raw), smape(obs_cv_raw, oof_pred_final)], - }) - oof_metrics.to_csv(dir_metric / f"eir_OOF_metrics_K{K}CV.csv", index=False) - - # Train final model on TRAIN+VAL - ts("Training final model on TRAIN+VAL with nrounds = median(best_iteration) ...") - best_nrounds = int(np.round(np.median(best_iters))) - - DT_trcv = DT[DT["split"] != "test"] - X_trcv = DT_trcv[features].values.astype(np.float64) - y_trcv = DT_trcv["eir_log10"].values - w_trcv = make_value_weights(np.power(10, y_trcv), digits=3) - - dtrcv = xgb.DMatrix(X_trcv, label=y_trcv, weight=w_trcv) - - xgb_cvfit = xgb.train( - params=xgb_params, - dtrain=dtrcv, - num_boost_round=best_nrounds, - verbose_eval=False, - ) - xgb_cvfit.save_model(str(dir_models / "eir_xgb_KCV.model")) - - # Predict on TEST and calibrate - dtest = xgb.DMatrix(X_test, label=y_test) - pred_log10_test_raw = xgb_cvfit.predict(dtest) - pred_raw_test = np.power(10, pred_log10_test_raw) - pred_eir_test = predict_qmap_w(pred_raw_test, cal_oof) - pred_eir_test = np.maximum(0, a_oof * pred_eir_test) - - # Save test predictions - test_preds = pd.DataFrame({ - "obs": obs_eir_test, - "pred_xgb": pred_eir_test - }) - test_preds.to_csv(dir_pred / "eir_test_predictions_xgb_QMAP_SCALE.csv", index=False) - - # Save diagnostic plots - if save_plots: - plot_obs_pred( - obs_eir_test, pred_eir_test, - f"EIR — Observed vs Predicted (XGBoost, K={K} CV, QMAP+Scale, test)", - str(dir_plots / "eir_obs_vs_pred_xgb_QMAP_SCALE_test.png"), - xlab="Observed EIR", ylab="Predicted EIR" - ) - plot_obs_pred( - y_test, np.log10(np.maximum(1e-12, pred_eir_test)), - f"EIR (log10) — Observed vs Predicted (XGBoost, K={K} CV after QMAP+Scale, test)", - str(dir_plots / "eir_log10_obs_vs_pred_xgb_after_QMAP_SCALE_test.png"), - xlab="Observed log10(EIR)", ylab="Predicted log10(EIR)" - ) - - # Calculate range-based metrics - bins = [0, 10, 50, 100, 200, np.inf] - labels = ["[0,10]", "(10,50]", "(50,100]", "(100,200]", "(200,Inf]"] - DTm = pd.DataFrame({ - "range": pd.cut(obs_eir_test, bins=bins, labels=labels, include_lowest=True), - "obs": obs_eir_test, - "pred": pred_eir_test, - "err": pred_eir_test - obs_eir_test - }) - - per_range_list = [] - for rng in labels: - subset = DTm[DTm["range"] == rng] - if len(subset) == 0: - continue - obs_s = subset["obs"].values - pred_s = subset["pred"].values - err_s = subset["err"].values - - per_range_list.append({ - "range": rng, - "N": len(subset), - "obs_mean": np.mean(obs_s), - "obs_median": np.median(obs_s), - "obs_sd": np.std(obs_s, ddof=1) if len(obs_s) > 1 else np.nan, - "pred_mean": np.mean(pred_s), - "bias": np.mean(err_s), - "MAE": mae(obs_s, pred_s), - "MedianAE": median_ae(obs_s, pred_s), - "RMSE": rmse(obs_s, pred_s), - "RMSLE": rmsle(obs_s, pred_s), - "NRMSE_mean": safe_div(rmse(obs_s, pred_s), np.mean(obs_s)), - "RelRMSE_p1": np.sqrt(np.mean(safe_div(err_s, np.maximum(1, obs_s)) ** 2)), - "sMAPE": smape(obs_s, pred_s), - }) - - per_range = pd.DataFrame(per_range_list) - - if save_artifacts: - per_range[["range", "RMSE"]].assign(model="xgboost_KCV").to_csv( - dir_metric / "eir_RMSE_by_range_test_QMAP_SCALE.csv", index=False - ) - per_range.assign(model="xgboost_KCV").to_csv( - dir_metric / "eir_metrics_by_range_test_QMAP_SCALE.csv", index=False - ) - - # Train deployment model on ALL filtered data - ts("Training deployment booster on ALL filtered data ...") - X_all = DT[features].values.astype(np.float64) - y_all = DT["eir_log10"].values - dall = xgb.DMatrix(X_all, label=y_all) - - xgb_final = xgb.train( - params=xgb_params, - dtrain=dall, - num_boost_round=best_nrounds, - verbose_eval=False, - ) - xgb_final.save_model(str(dir_models / "eir_xgb_FINAL.model")) - - # Create calibration bundle - cal_bundle = { - "kind": "qmap+scale", - "qmap": {"xq": cal_oof["xq"], "yq": cal_oof["yq"]}, - "scale": a_oof - } - - # Create preprocessing metadata - preprocess = { - "features": features, - "target": "eir", - "transform": "log10", - "inverse": "pow10", - "prevalence_filter": { - "min_prev_input": thr_lo, - "avg_prev_years_1_to_8_ge": thr_lo, - "year9_prev_le": thr_hi - }, - "reweighting": { - "scheme": "inverse_frequency_by_raw_EIR_value", - "digits": 3, - "applied_to": ["train", "val", "cv_folds"] - }, - "cv": { - "K": K, - "stratify_by": f"strat_bin (k-means on log10(EIR), centers={k_strata})", - "best_iteration_median": best_nrounds - }, - "calibration": { - "final": "QMAP then positive scale", - "final_pred": "pmax(0, a * QMAP(10^pred_log10))" - } - } - - # Create model bundle - model_bundle = { - "class": "estiMINT_model", - "booster": xgb_final, - "calibrator": cal_bundle, - "features": features, - "best_nrounds": best_nrounds, - "preprocess": preprocess, - "artifacts": { - "dir_models": str(dir_models), - "dir_plots": str(dir_plots), - "dir_metric": str(dir_metric), - "dir_pred": str(dir_pred) - } - } - - # Save pickle bundle - if save_pkl: - with open(dir_models / "estiMINT_model.pkl", "wb") as f: - pickle.dump(model_bundle, f, protocol=pickle.HIGHEST_PROTOCOL) - - # ONNX export not supported - if export_onnx: - raise NotImplementedError( - "ONNX export is not implemented in this version. " - "Use onnxmltools or skl2onnx externally to convert the model." - ) - - ts("Done. Artifacts saved under: %s", out_dir) - - return model_bundle diff --git a/src/estimint/types.py b/src/estimint/types.py index 10bc61a..80b01f6 100644 --- a/src/estimint/types.py +++ b/src/estimint/types.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Any, Literal from dataclasses import dataclass Input_Mode = Literal["prevalence", "eir", "hbr"] @@ -29,3 +29,13 @@ class Scenario: irs_future: float = 0.0 routine: float = 0.0 lsm: float = 0.0 + +@dataclass +class PreparedScenario: + """A ``Scenario`` resolved into the inputs each downstream model needs.""" + + eir_target: EirTarget + mosquito_density_change: float + eir_model_features: dict[str, float] + summary_values: dict[str, Any] + emulator_covariates: dict[str, float] diff --git a/src/estimint/utils.py b/src/estimint/utils.py index 8e2a5e3..71df065 100644 --- a/src/estimint/utils.py +++ b/src/estimint/utils.py @@ -6,8 +6,7 @@ import sys from datetime import datetime -from typing import Optional, Dict, Any, Union -from pathlib import Path +from typing import Dict, Any import numpy as np from numpy.typing import ArrayLike @@ -16,9 +15,9 @@ def ts(*args) -> None: """ Print timestamped message to console. - + Equivalent to R's ts() function. - + Parameters ---------- *args : str @@ -36,16 +35,16 @@ def ts(*args) -> None: def r2(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate R-squared (coefficient of determination). - + Equivalent to R's r2() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -61,16 +60,16 @@ def r2(y: ArrayLike, yhat: ArrayLike) -> float: def rmse(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate Root Mean Squared Error. - + Equivalent to R's rmse() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -84,16 +83,16 @@ def rmse(y: ArrayLike, yhat: ArrayLike) -> float: def mse(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate Mean Squared Error. - + Equivalent to R's mse() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -107,16 +106,16 @@ def mse(y: ArrayLike, yhat: ArrayLike) -> float: def mae(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate Mean Absolute Error. - + Equivalent to R's mae() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -126,20 +125,61 @@ def mae(y: ArrayLike, yhat: ArrayLike) -> float: yhat = np.asarray(yhat) return np.mean(np.abs(y - yhat)) +def medape(y: ArrayLike, yhat: ArrayLike) -> float: + """ + Calculate Median Absolute Percentage Error. + + Equivalent to R's medape() function. + + Parameters + ---------- + y : array-like + True values + yhat : array-like + Predicted values + + Returns + ------- + float + Median APE value + """ + y = np.asarray(y) + yhat = np.asarray(yhat) + return np.median(np.abs((y - yhat) / np.maximum(1, y))) * 100 + +def bias(y: ArrayLike, yhat: ArrayLike) -> float: + """ + Calculate bias (mean error). + + Parameters + ---------- + y : array-like + True values + yhat : array-like + Predicted values + + Returns + ------- + float + Bias value (mean of yhat - y) + """ + y = np.asarray(y) + yhat = np.asarray(yhat) + return np.mean(yhat - y) def median_ae(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate Median Absolute Error. - + Equivalent to R's median_ae() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -153,16 +193,16 @@ def median_ae(y: ArrayLike, yhat: ArrayLike) -> float: def mae_rel(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate Relative Median Absolute Error. - + Equivalent to R's mae_rel() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -176,16 +216,16 @@ def mae_rel(y: ArrayLike, yhat: ArrayLike) -> float: def rmsle(y: ArrayLike, yhat: ArrayLike) -> float: """ Calculate Root Mean Squared Log Error. - + Equivalent to R's rmsle() function. - + Parameters ---------- y : array-like True values yhat : array-like Predicted values - + Returns ------- float @@ -199,9 +239,9 @@ def rmsle(y: ArrayLike, yhat: ArrayLike) -> float: def safe_div(num: ArrayLike, den: ArrayLike, eps: float = 1e-12) -> np.ndarray: """ Safe division with epsilon floor on denominator. - + Equivalent to R's safe_div() function. - + Parameters ---------- num : array-like @@ -210,7 +250,7 @@ def safe_div(num: ArrayLike, den: ArrayLike, eps: float = 1e-12) -> np.ndarray: Denominator eps : float, optional Minimum value for denominator (default: 1e-12) - + Returns ------- np.ndarray @@ -224,9 +264,9 @@ def safe_div(num: ArrayLike, den: ArrayLike, eps: float = 1e-12) -> np.ndarray: def smape(y: ArrayLike, yhat: ArrayLike, eps: float = 1e-12) -> float: """ Calculate Symmetric Mean Absolute Percentage Error. - + Equivalent to R's smape() function. - + Parameters ---------- y : array-like @@ -235,7 +275,7 @@ def smape(y: ArrayLike, yhat: ArrayLike, eps: float = 1e-12) -> float: Predicted values eps : float, optional Epsilon for numerical stability (default: 1e-12) - + Returns ------- float @@ -254,9 +294,9 @@ def fit_qmap_w( ) -> Dict[str, Any]: """ Fit weighted quantile mapping calibration. - + Equivalent to R's fit_qmap_w() function. - + Parameters ---------- pred_raw : array-like @@ -267,7 +307,7 @@ def fit_qmap_w( Number of grid points for quantile mapping (default: 1024) round_digits : int, optional Digits for rounding observed values (default: 8) - + Returns ------- dict @@ -275,26 +315,26 @@ def fit_qmap_w( """ pred_raw = np.asarray(pred_raw) obs_raw = np.asarray(obs_raw) - + # Keep only finite values keep = np.isfinite(pred_raw) & np.isfinite(obs_raw) x = pred_raw[keep] y = obs_raw[keep] - + # Sort predictions and compute empirical CDF o1 = np.argsort(x) x1 = x[o1] F1 = (np.arange(1, len(x1) + 1) - 0.5) / len(x1) - + # Weighted CDF for observations y_key = np.round(y, round_digits) unique_y, counts = np.unique(y_key, return_counts=True) - + o2 = np.argsort(unique_y) y2 = unique_y[o2] w2 = counts[o2] F2 = np.cumsum(w2) / np.sum(w2) - + # Interpolate quantiles q = np.linspace(0, 1, ngrid) xq = np.interp(q, F1, x1) @@ -358,16 +398,16 @@ def predict_qmap_w(newx_raw: ArrayLike, cal: Dict[str, Any]) -> np.ndarray: def scale_pos(obs: ArrayLike, pred: ArrayLike) -> float: """ Calculate positive scaling factor. - + Equivalent to R's scale_pos() function. - + Parameters ---------- obs : array-like Observed values pred : array-like Predicted values - + Returns ------- float @@ -379,92 +419,3 @@ def scale_pos(obs: ArrayLike, pred: ArrayLike) -> float: if not np.isfinite(a) or a <= 0: a = 1.0 return a - - -def _find_installed_model() -> Optional[str]: - """ - Find model file installed with package. - - Equivalent to R's .find_installed_model() function. - - Returns - ------- - str or None - Path to model file if found, None otherwise - """ - import importlib.resources as pkg_resources - - try: - # Try different possible locations - candidates = [] - - # Check package data directories - try: - with pkg_resources.files("estimint") as pkg_path: - candidates.extend([ - pkg_path / "extdata" / "eir_model" / "estiMINT_model.pkl", - pkg_path / "extdata" / "estiMINT_model.pkl", - pkg_path / "estiMINT_model.pkl", - ]) - except (TypeError, AttributeError): - pass - - for cand in candidates: - if hasattr(cand, 'is_file') and cand.is_file(): - return str(cand) - elif isinstance(cand, (str, Path)) and Path(cand).exists(): - return str(cand) - - except Exception: - pass - - return None - - -def _resolve_model_file(dir_or_file: Union[str, Path]) -> str: - """ - Resolve model file from directory or file path. - - Equivalent to R's .resolve_model_file() function. - - Parameters - ---------- - dir_or_file : str or Path - Path to directory or file - - Returns - ------- - str - Path to model file - - Raises - ------ - FileNotFoundError - If model file cannot be found - """ - path = Path(dir_or_file) - - # If it's a file that exists, return it - if path.is_file(): - return str(path) - - # Must be a directory - if not path.is_dir(): - raise FileNotFoundError(f"Path does not exist: {dir_or_file}") - - # Try candidate locations - candidates = [ - path / "estiMINT_model.pkl", - path / "eir_model" / "estiMINT_model.pkl", - ] - - for cand in candidates: - if cand.is_file(): - return str(cand) - - # Search recursively for .pkl files - hits = list(path.rglob("estiMINT_model.pkl")) - if hits: - return str(hits[0]) - - raise FileNotFoundError(f"Could not find 'estiMINT_model.pkl' under: {dir_or_file}") diff --git a/src/estimint/v2/common/types.py b/src/estimint/v2/common/types.py new file mode 100644 index 0000000..1e0c0ae --- /dev/null +++ b/src/estimint/v2/common/types.py @@ -0,0 +1,23 @@ + +from typing import Callable, Literal, Protocol + +from flax import nnx +from jaxtyping import Array +from omegaconf import DictConfig +import numpy as np + + +class ModelFactory(Protocol): + @classmethod + def from_cfg(cls, cfg: DictConfig, input_size: int) -> nnx.Module: ... + +class ModelArtifact(Protocol): + def predict(self, X_raw: np.ndarray) -> np.ndarray: ... + +PredictorType = Literal["prev_y9", "eir", "hbr_y9"] +TargetType = Literal["eir", "hbr_y9"] + +# A loss returns (sum_loss, normalization): the unnormalized objective and its normalizer. +# The scalar loss is sum_loss / normalization; both terms sum across batches. +LossFn = Callable[[nnx.Module, Array, Array, Array], tuple[Array, Array]] + diff --git a/src/estimint/v2/conf/export_config.yaml b/src/estimint/v2/conf/export_config.yaml new file mode 100644 index 0000000..d178340 --- /dev/null +++ b/src/estimint/v2/conf/export_config.yaml @@ -0,0 +1,27 @@ +predictor: "eir" +target: "hbr_y9" +name: "${predictor}-${target}" + +# Data +features_scaler_file: ${output_dir}/features_scaler.pkl +target_scaler_file: ${output_dir}/target_scaler.pkl + + +# Model - ensure these match the checkpointed model's parameters +model_name: "RQS" +width: 256 +depth: 4 +mlp_residual: true +n_bins: 32 +rqs_bounds: 6 +dropout_rate: 0.0 +conformal_file: "${output_dir}/conformal-${timestamp}.json" + +# Checkpointing +checkpoint_dir: "${output_dir}/ckpts-${timestamp}" + +# General +output_dir: "train_outputs/${name}" +timestamp: ??? +seed: 42 +artifact_dir: "artifacts/${name}" \ No newline at end of file diff --git a/src/estimint/v2/conf/sweeps/sweep.yaml b/src/estimint/v2/conf/sweeps/sweep.yaml new file mode 100644 index 0000000..e5d4759 --- /dev/null +++ b/src/estimint/v2/conf/sweeps/sweep.yaml @@ -0,0 +1,44 @@ +program: estimint/v2/train_base.py +project: estimint-sweep +name: estimint-sweep +method: bayes +metric: + goal: minimize + name: val/loss + +parameters: + lr: + distribution: log_uniform_values + min: 1e-4 + max: 1e-2 + batch_size: + values: [256] + dropout_rate: + values: [0.0, 0.05, 0.1] + + # model-specific + width: + values: [256, 512] + depth: + values: [2, 3, 4] + n_bins: + values: [24, 32] + rqs_bounds: + values: [6, 8] + mlp_residual: + values: [true, false] + + +command: + - ${env} + - python + - -m + - estimint.v2.train_base + - "hydra.output_subdir=null" + - "hydra.run.dir=." + - ${args_no_hyphens} + - use_wandb=true + - num_epochs=120 + - min_epochs=120 + - predictor=eir + - target=hbr_y9 diff --git a/src/estimint/v2/conf/train_config.yaml b/src/estimint/v2/conf/train_config.yaml new file mode 100644 index 0000000..81fe26a --- /dev/null +++ b/src/estimint/v2/conf/train_config.yaml @@ -0,0 +1,40 @@ +predictor: "prev_y9" +target: "eir" +name: "${predictor}-${target}" + +# data +data_file: "datasets/estimint_simulations_y9.parquet" +split_file: "datasets/split_${name}.csv" +num_workers: 0 +use_existing_split: false +stratify: false +calib_frac: 0.06 + +# model +width: 256 +depth: 2 +mlp_residual: false +dropout_rate: 0.0 +n_bins: 24 +rqs_bounds: 6 +conformal_file: "${output_dir}/conformal-${cur_time}.json" + +# Hyperparameters +num_epochs: 120 +min_epochs: 100 +patience: 30 +lr: 1e-3 +batch_size: 256 +weight_decay: 1e-4 + +# Checkpoint +checkpoint_dir: "${output_dir}/ckpts-${cur_time}" + +# General +cur_time: ${now:%Y-%m-%dT%H:%M:%S} +seed: 42 +use_wandb: false +output_dir: "train_outputs/${name}" +wandb: + project: "estimint-training-${name}" + name: "train-${cur_time}" \ No newline at end of file diff --git a/src/estimint/v2/data/dataset.py b/src/estimint/v2/data/dataset.py new file mode 100644 index 0000000..3b3b3d5 --- /dev/null +++ b/src/estimint/v2/data/dataset.py @@ -0,0 +1,77 @@ +import grain.python as grain + + +class DataSource(grain.RandomAccessDataSource): + def __init__(self, data: list[dict]): + """ + Store records for random access. + + Args: + data: Sequence records. + + Returns: + None. + """ + self.data = data + + def __len__(self) -> int: + """ + Return the number of records. + + Returns: + Number of records. + """ + return len(self.data) + + def __getitem__(self, idx: int) -> dict: + """ + Return one record by index. + + Args: + idx: Record index. + + Returns: + Record dictionary. + """ + return self.data[idx] + + +def make_loader( + data: list[dict], + batch_size: int, + shuffle: bool = False, + seed: int = 42, + num_workers: int = 0, + drop_remainder: bool = True, +) -> grain.DataLoader: + """ + Build a Grain data loader. + + Args: + data: Sequence records. + batch_size: Batch size. + shuffle: Whether to shuffle indices. + seed: Sampler seed. + num_workers: Worker count. + drop_remainder: Whether to drop partial batches. + + Returns: + Configured Grain data loader. + """ + data_source = DataSource(data) + + sampler = grain.IndexSampler( + num_records=len(data_source), + num_epochs=1, + shard_options=grain.NoSharding(), + shuffle=shuffle, + seed=seed, + ) + + loader = grain.DataLoader( + data_source=data_source, + sampler=sampler, + operations=[grain.Batch(batch_size=batch_size, drop_remainder=drop_remainder)], + worker_count=num_workers, + ) + return loader diff --git a/src/estimint/v2/data/features.py b/src/estimint/v2/data/features.py new file mode 100644 index 0000000..1554a9f --- /dev/null +++ b/src/estimint/v2/data/features.py @@ -0,0 +1,114 @@ +import numpy as np +from ..common.types import PredictorType + +FEATURES_BASE = ["dn0_use", "Q0", "phi_bednets", "seasonal", "itn_use", "irs_use"] +LOG_FEATURES = ("eir", "hbr_y9") + +def get_features(predictor: PredictorType) -> list[str]: + """ + Get the list of features based on the predictor and target. + The predictor is is inserted at the beginning of the list of features. + + Args: + predictor: The predictor type. + Returns: + A list of feature names. + """ + return [predictor] + FEATURES_BASE + + +class StandardScaler: + def __init__(self): + """ + Initialize an unfitted scaler. + + Returns: + None. + """ + self.mean_: np.ndarray | None = None + self.scale_: np.ndarray | None = None + + @property + def is_fitted(self) -> bool: + return self.mean_ is not None and self.scale_ is not None + + def fit(self, X: np.ndarray) -> "StandardScaler": + """ + Fit feature means and scales. + + Args: + X: Feature matrix. + + Returns: + Fitted scaler. + """ + self.mean_ = np.mean(X, axis=0) + # To avoid division by zero, set scale to 1.0 for any feature with zero variance + scale = np.std(X, axis=0) + scale[scale == 0] = 1.0 + self.scale_ = scale + return self + + def transform(self, X: np.ndarray) -> np.ndarray: + """ + Standardize features. + + Args: + X: Feature matrix. + + Returns: + Standardized features. + """ + if not self.is_fitted: + raise ValueError("StandardScaler instance is not fitted yet.") + return (X - self.mean_) / self.scale_ + + def fit_transform(self, X: np.ndarray) -> np.ndarray: + """ + Fit and standardize features. + + Args: + X: Feature matrix. + + Returns: + Standardized features. + """ + return self.fit(X).transform(X) + + def inverse_transform(self, X: np.ndarray) -> np.ndarray: + """ + Restore standardized features. + + Args: + X: Standardized feature matrix. + + Returns: + Features in the original scale. + """ + if not self.is_fitted: + raise ValueError("StandardScaler instance is not fitted yet.") + return X * self.scale_ + self.mean_ + +class FeatureScaler(StandardScaler): + """StandardScaler that log10s the features that are in LOG_FEATURES before standardizing them.""" + def __init__(self, log_idx: list[int] = []): + super().__init__() + self.log_idx = log_idx + + @classmethod + def for_features(cls, features: list[str]) -> "FeatureScaler": + return cls([i for i, f in enumerate(features) if f in LOG_FEATURES]) + + def _pre(self, X: np.ndarray) -> np.ndarray: + if not self.log_idx: + return X + X = np.array(X, copy=True) + X[..., self.log_idx] = np.log10(np.maximum(X[..., self.log_idx], 1e-12)) + return X + + def fit(self, X: np.ndarray) -> "FeatureScaler": + super().fit(self._pre(X)) + return self + + def transform(self, X: np.ndarray) -> np.ndarray: + return super().transform(self._pre(X)) diff --git a/src/estimint/v2/data/preprocess.py b/src/estimint/v2/data/preprocess.py new file mode 100644 index 0000000..f951b07 --- /dev/null +++ b/src/estimint/v2/data/preprocess.py @@ -0,0 +1,422 @@ + + +from omegaconf import DictConfig +import random +import pandas as pd +import logging +from pathlib import Path +import numpy as np +from .features import StandardScaler, FeatureScaler, get_features +from estimint.data_processing import make_value_weights +import pickle +from dataclasses import dataclass, field +from typing import cast + +log = logging.getLogger(__name__) + +@dataclass +class PreparedData: + train_data: list + val_data: list + test_data: list + input_size: int + feature_scaler: StandardScaler + target_scaler: StandardScaler + calib_data: list = field(default_factory=list) + train_param_sims: set[tuple[int, int]] = field(default_factory=set) + val_param_sims: set[tuple[int, int]] = field(default_factory=set) + test_param_sims: set[tuple[int, int]] = field(default_factory=set) + calib_param_sims: set[tuple[int, int]] = field(default_factory=set) + +@dataclass +class SplitParamSims: + train: set[tuple[int, int]] = field(default_factory=set) + val: set[tuple[int, int]] = field(default_factory=set) + calib: set[tuple[int, int]] = field(default_factory=set) + test: set[tuple[int, int]] = field(default_factory=set) + +def _filter_by_threshold(df: pd.DataFrame) -> pd.DataFrame: + """ + Filter parameter-simulation pairs where the mean target value is below the threshold. + + Args: + df: Input dataframe. + + Returns: + Filtered dataframe. + """ + prev_threshold = 0.01 # TODO: make this configurable + group_means = df.groupby(["parameter_index", "simulation_index"])["prev_y9"].mean() + valid = set(map(tuple, group_means[group_means >= prev_threshold].index.tolist())) + df["_ps"] = list(zip(df["parameter_index"], df["simulation_index"])) + + log.info( + f"Filtering with prev_threshold {prev_threshold} on {"prevalence"}: {len(valid)} valid parameter-simulation pairs out of {len(group_means)}" + ) + + return df[df["_ps"].isin(valid)] + + +def _load_split( + split_file: str, df: pd.DataFrame +) -> SplitParamSims: + """ + Load an existing train/val/test split. + + Args: + split_file: Split CSV path. + df: Filtered dataframe. + + Returns: + Train, validation, and test parameter-simulation sets. + """ + split_df = pd.read_csv(split_file) + present = set(df[["parameter_index", "simulation_index"]].itertuples(index=False, name=None)) + + def _pairs(df: pd.DataFrame, split_name: str) -> set[tuple[int, int]]: + return { + + (r.parameter_index, r.simulation_index) for r in df[df["split"] == split_name].itertuples() + } & present # type: ignore + + return SplitParamSims( + train=_pairs(split_df, "train"), + val=_pairs(split_df, "validate"), + test=_pairs(split_df, "test"), + calib=_pairs(split_df, "calibrate") + ) + +# TODO: check if need to strata or not. +def _parameter_strata( + df: pd.DataFrame, + target: str, + n_bins: int, + stratify: bool, +) -> list[np.ndarray]: + """ + Build parameter-index groups used for stratified split assignment. + + Args: + df: Input dataframe containing ``parameter_index`` and ``target`` columns. + target: Column whose per-parameter mean defines the strata. + n_bins: Maximum number of quantile bins to create when stratifying. + stratify: If False, return all parameters as a single group. + + Returns: + A list of parameter-index arrays. Each array is assigned to train, + validation, calibration, and test splits independently. + """ + param_target = df.groupby("parameter_index")[target].mean() + param_indices = param_target.index.to_numpy() + if not stratify: + return [param_indices] + + target_values = np.log10(param_target.to_numpy(dtype=np.float64)) + unique_targets = np.unique(target_values) + + quantile_bins = pd.qcut( + target_values, + q=min(n_bins, len(unique_targets)), + labels=False, + duplicates="drop", + ) + + return [param_indices[quantile_bins == bucket] for bucket in np.unique(quantile_bins)] + + +def _param_sims_from_assignment(df: pd.DataFrame, assign: dict[int, str]) -> SplitParamSims: + """Build parameter-simulation sets from a parameter-level split assignment.""" + all_ps = set(df[["parameter_index", "simulation_index"]].itertuples(index=False, name=None)) + split_params = {name: {p for p, split in assign.items() if split == name} for name in ["train", "val", "calib", "test"]} + return SplitParamSims( + train={ps for ps in all_ps if ps[0] in split_params["train"]}, + val={ps for ps in all_ps if ps[0] in split_params["val"]}, + calib={ps for ps in all_ps if ps[0] in split_params["calib"]}, + test={ps for ps in all_ps if ps[0] in split_params["test"]}, + ) + +def _assign_param_group( + params: np.ndarray, + rng: np.random.Generator, + train_frac: float, + val_frac: float, + calib_frac: float, +) -> dict[int, str]: + shuffled = np.array(params, copy=True) + rng.shuffle(shuffled) + n = len(shuffled) + train_end, val_end, calib_end = np.cumsum([np.array([train_frac, val_frac, calib_frac]) * n]).astype(int) + + assigned: dict[int, str] = {} + for split_name, split_param in ( + ("train", shuffled[:train_end]), + ("val", shuffled[train_end: val_end]), + ("calib", shuffled[val_end: calib_end]), + ("test", shuffled[calib_end:]), + ): + assigned.update({param: split_name for param in split_param}) + return assigned + +def _assign_param_splits( + df: pd.DataFrame, + *, + seed: int, + calib_frac: float = 0.0, + stratify: bool = True, + target: str = "eir", + n_bins: int = 10, +) -> dict[int, str]: + """Assign each parameter_index to a split, grouping whole parameters. + If stratify is True, the assignment is balanced across quantile bins of the mean target value per parameter. + """ + rng = np.random.default_rng(seed) + train_frac = 0.7 + val_frac = (1.0 - train_frac - calib_frac) / 2.0 + if val_frac <= 0: + raise ValueError("Validation fraction must be positive; check calib_frac.") + + assign: dict[int, str] = {} + for params in _parameter_strata(df, target, n_bins, stratify): + assign.update(_assign_param_group(params, rng, train_frac, val_frac, calib_frac)) + return assign + +def _create_split( + df: pd.DataFrame, + seed: int, + *, + stratify: bool = True, + target: str = "eir", + n_bins: int = 10, + calib_frac: float = 0.0 +) -> SplitParamSims: + """ + Create grouped parameter-simulation splits. + + This is the pair-set version of :func:`group_split`. Both functions share + the same parameter-level assignment, so train/val/calib/test never contain + different simulations from the same ``parameter_index``. + + Args: + df: Filtered dataframe. + seed: Shuffle seed. + calib_frac: Fraction to split calibration + stratify: Balance the split across target-magnitude quantile bins. + target: Column used for stratification. + n_bins: Number of quantile strata when ``stratify`` is True. + + Returns: + Grouped train, validation, optional calibration, and test + parameter-simulation sets. + """ + assign = _assign_param_splits( + df, + seed=seed, + calib_frac=calib_frac, + stratify=stratify, + target=target, + n_bins=n_bins, + ) + return _param_sims_from_assignment(df, assign) + + +def _save_split(path, split_ps: SplitParamSims): + """ + Save parameter-simulation split assignments. + + Args: + path: Output CSV path. + split_ps: Split parameter-simulation sets. + df: Source dataframe. + + Returns: + None. + """ + rows = [ + (param_idx, sim_idx, split) + for split, ps in ( + ("train", split_ps.train), + ("validate", split_ps.val), + ("test", split_ps.test), + ("calibrate", split_ps.calib), + ) + for param_idx, sim_idx in ps + ] + pd.DataFrame(rows, columns=["parameter_index", "simulation_index", "split"]).to_csv(path, index=False) + log.info(f"Split saved to {path}") + +def _fit_features_scaler(df: pd.DataFrame, train_ps: set[tuple[int, int]], output_dir: str, features: list[str]) -> FeatureScaler: + """ + Fit and save the static feature scaler. + + Args: + df: Filtered dataframe. + train_ps: Training pairs. + output_dir: Directory for scaler output. + features: List of feature columns to use. + + Returns: + Fitted scaler. + """ + train_mask = df["_ps"].isin(train_ps) + train_static = ( + df.loc[train_mask, ["_ps"] + features] + .drop_duplicates(subset=["_ps"])[features] + .to_numpy(dtype=np.float32) + ) + scaler = FeatureScaler.for_features(features) + scaler.fit(train_static) + + save_path = Path(output_dir) / "features_scaler.pkl" + save_path.parent.mkdir(parents=True, exist_ok=True) + with open(save_path, "wb") as f: + pickle.dump(scaler, f) + + return scaler + +def _fit_target_scaler(df: pd.DataFrame, train_ps: set[tuple[int, int]], output_dir: str, target: str = "eir") -> StandardScaler: + """ + Fit and save the target scaler. + + Args: + df: Filtered dataframe. + train_ps: Training pairs. + output_dir: Directory for scaler output. + target: Target column to scale. + Returns: + Fitted scaler. + """ + train_mask = df["_ps"].isin(train_ps) + train_y = ( + df.loc[train_mask, ["_ps", target]] + .drop_duplicates(subset=["_ps"])[target] + .to_numpy(dtype=np.float32) + ) + scaler = StandardScaler() + scaler.fit(np.log10(train_y)[:, None]) # Fit on log10 of target + + save_path = Path(output_dir) / "target_scaler.pkl" + save_path.parent.mkdir(parents=True, exist_ok=True) + with open(save_path, "wb") as f: + pickle.dump(scaler, f) + return scaler + +def _build_data( + df: pd.DataFrame, + param_sims: set[tuple[int, int]], + scaler: StandardScaler, + target_scaler: StandardScaler, + features: list[str], + target +) -> list[dict[str, np.ndarray]]: + """ + Build scaled per-parameter-simulation training records. + + Args: + df: Filtered dataframe with feature, target, and ``_weight`` columns. + param_sims: Parameter-simulation pairs to include. + scaler: Fitted static feature scaler. + target_scaler: Fitted target scaler. + features: Feature columns to scale and include as model inputs. + target: Positive target column to log-transform. + + Returns: + A list of sequence dictionaries containing scaled features, raw + features, log10 targets, raw targets, sample weights, and pair IDs. + """ + # One row per (parameter_index, simulation_index): features/target are + # static per simulation, so any row in the group carries the same values. + rows = df.groupby(["parameter_index", "simulation_index"]).first() + data = [] + + for ps in param_sims: + if ps not in rows.index: + continue + + row = cast(pd.Series, rows.loc[ps]) + X_raw = row[features].to_numpy(dtype=np.float32) + X = scaler.transform(X_raw) + Y_raw = np.float32(row[target]) + Y = np.log10(Y_raw) + Y_std = target_scaler.transform(np.array([[Y]], dtype=np.float32))[0, 0] + W = np.float32(row["_weight"]) + data.append( + { + "x_raw": X_raw, + "x": X, + "y_raw": Y_raw, + "y": Y, # log10 target + "y_std": Y_std, # standardized log10 target + "w": W, + "ps": np.asarray(ps, dtype=np.int32), # (2,) parameter_index, simulation_index + } + ) + + return data + + + +def prepare_data(df: pd.DataFrame, cfg: DictConfig, calib_frac: float = 0.0) -> PreparedData: + """ + Split and transform raw simulation data. + + Filters out low-signal parameter-simulation pairs, creates or loads the + train/val/test split, fits static feature and target scaling on the train split only, + and builds per-sequence records for each split. + + Args: + df: Raw simulation dataframe. + cfg: Data preparation config. + + Returns: + Prepared train, validation, and test data. + """ + random.seed(cfg.seed) + + # Filter by threshold + df = _filter_by_threshold(df) + df["_weight"] = make_value_weights(df[cfg.target].to_numpy(dtype=np.float32)) + # split data + if cfg.use_existing_split and Path(cfg.split_file).exists(): + log.info(f"Loading existing split from {cfg.split_file}") + split_ps = _load_split(cfg.split_file, df) + else: + log.info("Creating new train/val/calib/test split") + split_ps = _create_split(df, cfg.seed, stratify=cfg.stratify, target=cfg.target, calib_frac=calib_frac) + if cfg.split_file: + log.info(f"Saving split to {cfg.split_file}") + _save_split(cfg.split_file, split_ps) + + log.info( + "Split — train: %s, val: %s, calib: %s, test: %s", + len(split_ps.train), + len(split_ps.val), + len(split_ps.calib), + len(split_ps.test), + ) + + features = get_features(cfg.predictor) + features_scaler = _fit_features_scaler(df, split_ps.train, cfg.output_dir, features) + target_scaler = _fit_target_scaler(df, split_ps.train, cfg.output_dir, target=cfg.target) + + def build_data(param_sims): + return _build_data(df, param_sims, features_scaler, target_scaler, features, cfg.target) + + train_data = build_data(split_ps.train) + val_data = build_data(split_ps.val) + test_data = build_data(split_ps.test) + calib_data = build_data(split_ps.calib) + + return PreparedData( + train_data=train_data, + val_data=val_data, + test_data=test_data, + calib_data=calib_data, + input_size=len(features), + feature_scaler=features_scaler, + target_scaler=target_scaler, + train_param_sims=split_ps.train, + val_param_sims=split_ps.val, + test_param_sims=split_ps.test, + calib_param_sims=split_ps.calib + ) + diff --git a/src/estimint/v2/eval/metrics.py b/src/estimint/v2/eval/metrics.py new file mode 100644 index 0000000..444f662 --- /dev/null +++ b/src/estimint/v2/eval/metrics.py @@ -0,0 +1,45 @@ +import numpy as np +from grain.python import DataLoader +from estimint.utils import mse, r2, rmse, mae, bias, medape +from dataclasses import dataclass +from estimint.v2.common.types import ModelArtifact + +@dataclass +class Metrics: + mse: float + r2: float + rmse: float + log10_mse: float + mae: float + bias: float + medape: float + +def compute_metrics( + model_artifact: ModelArtifact, + loader: DataLoader, +): + preds, targets = get_preds_targets(model_artifact, loader) + + return Metrics( + mse=mse(targets, preds), + r2=r2(targets, preds), + rmse=rmse(targets, preds), + mae=mae(targets, preds), + bias=bias(targets, preds), + medape=medape(targets, preds), + log10_mse=mse(np.log10(targets), np.log10(preds)) + ) + + +def get_preds_targets(model_artifact: ModelArtifact, data_loader: DataLoader) -> tuple[np.ndarray, np.ndarray]: + all_preds, all_targets = [], [] + for batch in data_loader: + preds = model_artifact.predict(batch["x_raw"]) + all_preds.append(preds) + all_targets.append(batch["y_raw"]) + + all_preds = np.concat(all_preds) + all_targets = np.concat(all_targets) + + return all_preds, all_targets + diff --git a/src/estimint/v2/model_export.py b/src/estimint/v2/model_export.py new file mode 100644 index 0000000..b3a5d68 --- /dev/null +++ b/src/estimint/v2/model_export.py @@ -0,0 +1,67 @@ +import logging +import hydra +import pickle +from omegaconf import DictConfig, OmegaConf +from pathlib import Path +from .models.rqs import ConditionalRQS +from .training.checkpoint import restore_model, save_checkpoint +from .data.preprocess import StandardScaler, FeatureScaler, get_features +import json + +log = logging.getLogger(__name__) + +@hydra.main(version_base=None, config_path="conf", config_name="export_config") +def main(cfg: DictConfig): + """ + Export the trained model for sharing to other users. + + Args: + cfg: Hydra config for exporting the model. + """ + log.info(OmegaConf.to_yaml(cfg)) + artifact_dir = Path(cfg.artifact_dir) + artifact_dir.mkdir(parents=True, exist_ok=True) + + with open(cfg.features_scaler_file, "rb") as f: + feature_scaler: FeatureScaler = pickle.load(f) + with open(cfg.target_scaler_file, "rb") as f: + target_scaler: StandardScaler = pickle.load(f) + if not feature_scaler.is_fitted or not target_scaler.is_fitted: + raise ValueError("Feature or target scaler is not fitted. Please fit the scalers before exporting the model.") + + with open(cfg.conformal_file, "r") as f: + conformal = json.load(f) + features = get_features(cfg.predictor) + model = ConditionalRQS.from_cfg(cfg, n_context=len(features)) + model = restore_model(cfg.checkpoint_dir, cfg.model_name, model) + model.eval() + save_checkpoint(f"{cfg.artifact_dir}/checkpoint", cfg.model_name, model) + + + config = dict( + model_name=cfg.model_name, + predictor=cfg.predictor, + target=cfg.target, + width=cfg.width, + depth=cfg.depth, + n_bins=cfg.n_bins, + rqs_bounds=cfg.rqs_bounds, + mlp_residual=cfg.mlp_residual, + dropout_rate=cfg.dropout_rate, + features=features, + feature_scalar_mean=feature_scaler.mean_.tolist(), + feature_scalar_scale=feature_scaler.scale_.tolist(), + feature_log_idx=list(feature_scaler.log_idx), + target_scalar_mean=target_scaler.mean_.tolist(), + target_scalar_scale=target_scaler.scale_.tolist(), + conformal=conformal, + ) + + with (artifact_dir / "config.json").open("w") as f: + json.dump(config, f, indent=2) + + log.info(f"Exported model and config to {artifact_dir}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/estimint/v2/models/hub.py b/src/estimint/v2/models/hub.py new file mode 100644 index 0000000..eceb92e --- /dev/null +++ b/src/estimint/v2/models/hub.py @@ -0,0 +1,95 @@ +import json +from pathlib import Path +from typing import Any + +import numpy as np +from huggingface_hub import snapshot_download +from omegaconf import OmegaConf +from ..common.types import PredictorType, TargetType + +from ..data.features import StandardScaler, FeatureScaler +from .rqs import ConditionalRQS, RQSArtifact +from ..training.checkpoint import restore_model + + +def repo_id(predictor: PredictorType, target: TargetType) -> str: + return f"{predictor}-{target}" + +def _load_json(path: Path) -> dict[str, Any]: + with path.open("r") as f: + return json.load(f) + + +def _load_scaler(mean: list[float], scale: list[float], log_idx: list[int] | None = None) -> StandardScaler: + scaler = FeatureScaler(log_idx) if log_idx is not None else StandardScaler() + scaler.mean_ = np.array(mean, dtype=np.float32) + scaler.scale_ = np.array(scale, dtype=np.float32) + return scaler + + +def _download_from_hf( + repo_id: str, + name: str, + *, + revision: str | None = None, + cache_dir: str | Path | None = None, + local_dir: str | Path | None = None, +) -> Path: + root = snapshot_download( + repo_id=repo_id, + allow_patterns=[f"{name}/*", f"{name}/**"], + revision=revision, + cache_dir=cache_dir, + local_dir=local_dir, + ) + return Path(root) / name + + +def load_model_artifact( + path_or_repo_id: str, + predictor: PredictorType, + target: TargetType, + *, + revision: str | None = None, + cache_dir: str | Path | None = None, + local_dir: str | Path | None = None, +) -> RQSArtifact: + """ + Load an RQS inference artifact from a local folder or Hugging Face repo. + + Expected artifact layout, matching estimint.v2.model_export: + /config.json + /checkpoint/ + + Args: + path_or_repo_id: Hugging Face repo ID or local folder path. + name: Artifact subfolder name, e.g. "prev_y9-eir". + revision: Optional revision of the model to load from the repo. + cache_dir: Optional cache directory for the Hugging Face repo. + local_dir: Optional local directory to download the repo into. + + Returns: + RQSArtifact wrapping the restored model and fitted scalers. + """ + if Path(path_or_repo_id).exists(): + root = Path(path_or_repo_id) + model_dir = root / repo_id(predictor, target) + artifact_dir = model_dir if model_dir.is_dir() else root + else: + artifact_dir = _download_from_hf( + path_or_repo_id, repo_id(predictor, target), revision=revision, cache_dir=cache_dir, local_dir=local_dir + ) + + config = _load_json(artifact_dir / "config.json") + + features = config["features"] + model = ConditionalRQS.from_cfg(OmegaConf.create(config), n_context=len(features)) + model = restore_model(str(artifact_dir / "checkpoint"), config["model_name"], model) + model.eval() + + feature_scaler = _load_scaler(config["feature_scalar_mean"], config["feature_scalar_scale"], config["feature_log_idx"]) + target_scaler = _load_scaler(config["target_scalar_mean"], config["target_scalar_scale"]) + conformal = {float(k): v for k, v in config.get("conformal", {}).items()} + + return RQSArtifact(model=model, feature_scaler=feature_scaler, target_scaler=target_scaler, features=features, conformal=conformal) + diff --git a/src/estimint/v2/models/mlp.py b/src/estimint/v2/models/mlp.py new file mode 100644 index 0000000..3071318 --- /dev/null +++ b/src/estimint/v2/models/mlp.py @@ -0,0 +1,27 @@ +from flax import nnx + +# TODO: can props remove residual if block +class MLP(nnx.Module): + def __init__(self, din, dout, *, width=128, depth=4, residual=False, dropout_rate=0.1, rngs): + self.residual = residual + self.inp = nnx.Linear(din, width, rngs=rngs) + if residual: + self.norms = nnx.List([nnx.LayerNorm(width, rngs=rngs) for _ in range(depth)]) + self.fc1 = nnx.List([nnx.Linear(width, width, rngs=rngs) for _ in range(depth)]) + self.fc2 = nnx.List([nnx.Linear(width, width, rngs=rngs) for _ in range(depth)]) + else: + self.hidden = nnx.List([nnx.Linear(width, width, rngs=rngs) for _ in range(depth)]) + self.dropout = nnx.Dropout(dropout_rate, rngs=rngs) + self.out = nnx.Linear(width, dout, rngs=rngs) + + def __call__(self, x): + x = self.inp(x) + if self.residual: + for ln, f1, f2 in zip(self.norms, self.fc1, self.fc2): + x = x + self.dropout(f2(nnx.gelu(f1(ln(x))))) # pre-LN residual block + x = nnx.gelu(x) + else: + x = nnx.gelu(x) + for h in self.hidden: + x = self.dropout(nnx.gelu(h(x))) + return self.out(x) \ No newline at end of file diff --git a/src/estimint/v2/models/rqs.py b/src/estimint/v2/models/rqs.py new file mode 100644 index 0000000..d2ac10a --- /dev/null +++ b/src/estimint/v2/models/rqs.py @@ -0,0 +1,273 @@ +from flax import nnx +from .mlp import MLP +import jax.numpy as jnp +import jax +from estimint.v2.data.features import StandardScaler, FeatureScaler +import numpy as np +from omegaconf import DictConfig +from ..common.types import PredictorType, TargetType + +class ConditionalRQS(nnx.Module): + """Conditional rational-quadratic spline flow. + + Note: The inverse is flipped as compaerd to standard normalizing flow convention. + Usually it is as x = T(z) where z ~ N(0, 1). Here, we have defined z = T(x) where z ~ N(0, 1). + And x = T^{-1}(z) + + """ + def __init__(self, n_context, *, width=128, depth=4, n_bins=12, bounds=6, residual=False, dropout_rate=0.0, rngs: nnx.Rngs): + self.K = n_bins + self.bounds = bounds + # The net maps context -> all spline params: K widths + K heights + + # (K+1) derivatives = 3K+1 numbers. + self.net = MLP(n_context, 3 * n_bins + 1, width=width, depth=depth, residual=residual, dropout_rate=dropout_rate, rngs=rngs) + + def _params(self, context): + params = self.net(context) # (B, 3K + 1) + return jnp.split(params, [self.K, 2 * self.K], axis=-1) # widths(K), heights(K), derivatives(K+1) + + def log_prob(self, y0, context): + """standardized target y0 -> base z; density = Normal(z) * |dT/dy|.""" + widths, heights, derivatives = self._params(context) + z, log_det = _rqs(y0, widths, heights, derivatives, self.bounds, inverse=False) + return jax.scipy.stats.norm.logpdf(z) + log_det + + def quantile(self, context, quantile): + """Flow base z -> target y.The q-quantile of y is flow T^{-1}(Phi^{-1}(q))""" + widths, heights, derivatives = self._params(context) + z = jax.scipy.stats.norm.ppf(quantile) # Phi^{-1}(q) + y0, _ = _rqs(z, widths, heights, derivatives, self.bounds, inverse=True) + return y0 + + def quantiles(self, context, probs): + """Evaluate many quantile levels at once, reusing the per-row spline params. + + Args: + context: (B, C) conditioning features. + probs: (Q,) probability levels in (0, 1). + + Returns: + (Q, B) standardized targets y0, one row per probability level. + """ + widths, heights, derivatives = self._params(context) # (B,K),(B,K),(B,K+1) + zs = jax.scipy.stats.norm.ppf(probs) # (Q,) + + def _invert(z_scalar): + z_col = jnp.full((context.shape[0],), z_scalar) + y0, _ = _rqs(z_col, widths, heights, derivatives, self.bounds, inverse=True) + return y0 + + return jax.vmap(_invert)(zs) # (Q, B) + + @classmethod + def from_cfg(cls, cfg: DictConfig, n_context: int) -> "ConditionalRQS": + return cls( + n_context, + rngs=nnx.Rngs(cfg.get("seed", 0)), + width=cfg.width, + depth=cfg.depth, + n_bins=cfg.n_bins, + bounds=cfg.rqs_bounds, + residual=cfg.mlp_residual, + dropout_rate=cfg.dropout_rate) + + @classmethod + def from_pretrained( + cls, + path_or_repo_id: str, + predictor: PredictorType, + target: TargetType, + *, + revision: str | None = None, + cache_dir: str | None = None, + local_dir: str | None = None, + ) -> "RQSArtifact": + """ + Load a pretrained RQS artifact from a local folder or Hugging Face repo. + + Example usage: + ``` + ConditionalRQS.from_pretrained("dide-ic/estiMINT", predictor="prev_y9", target="eir") + ConditionalRQS.from_pretrained("dide-ic/estiMINT", predictor="prev_y9", target="eir", revision="v0.1.0") + ``` + + Args: + path_or_repo_id: Hugging Face repo ID or local folder path. + predictor: Predictor type. + target: Target type. + revision: Optional revision of the model to load from the repo. + cache_dir: Optional cache directory for the Hugging Face repo. + local_dir: Optional local directory to download the repo into. + + Returns: + RQSArtifact containing the restored model and fitted scalers. + """ + from .hub import load_model_artifact + + + return load_model_artifact( + path_or_repo_id, predictor, target, revision=revision, cache_dir=cache_dir, local_dir=local_dir, + ) + +@nnx.jit +def _forward(model: nnx.Module, context: jnp.ndarray, quantile: float): + return model.quantile(context, quantile) # type: ignore + +FeatureInput = np.ndarray | dict[str, float] | list[dict[str, float]] +class RQSArtifact: + def __init__(self, model: nnx.Module, feature_scaler: FeatureScaler | StandardScaler, target_scaler: StandardScaler, features: list[str], conformal: dict[float, float] | None = None): + self.model = model + self.feature_scaler = feature_scaler + self.target_scaler = target_scaler + self.conformal = dict(conformal or {}) # alpha -> offset Q + self.feature_names = features + + if self.feature_scaler.mean_.shape[0] != len(self.feature_names): + raise ValueError(f"Feature scaler has {self.feature_scaler.mean_.shape[0]} features, but expected {len(self.feature_names)} features for features {self.feature_names}.") + + def _prepare_inputs(self, X_raw: FeatureInput) -> np.ndarray: + """Normalize user input to a (B, C) float32 array in training feature order.""" + if isinstance(X_raw, dict): + X_raw = [X_raw] + + if isinstance(X_raw, list): + if not X_raw: + raise ValueError("Input list is empty.") + + names = self.feature_names + rows = [] + for i, row in enumerate(X_raw): + missing = [f for f in names if f not in row] + extra = [f for f in row if f not in names] + if missing or extra: + raise KeyError(f"row {i}: missing={missing}, unexpected={extra}. Expected exactly {names}.") + rows.append([row[f] for f in names]) # preserve feature order + X = np.array(rows, dtype=np.float32) + else: + X = np.asarray(X_raw, dtype=np.float32) + if X.ndim == 1: + X = X[None, :] # add batch dimension + return X + + + def _quantile(self, X_raw: FeatureInput, quantile: float) -> np.ndarray: + X = self._prepare_inputs(X_raw) + context = jnp.array(self.feature_scaler.transform(X)) + y0 = _forward(self.model, context, quantile) + return np.maximum(0, np.power(10, self.target_scaler.inverse_transform(y0))) + + def predict(self, X_raw: FeatureInput) -> np.ndarray: + return self._quantile(X_raw, 0.5) # median prediction + + def quantile(self, X_raw: FeatureInput, quantile: float) -> np.ndarray: + return self._quantile(X_raw, quantile) + + def interval(self, X_raw: FeatureInput, alpha: float = 0.10) -> tuple[np.ndarray, np.ndarray]: + """Conformal (1-alpha) band with guaranteed coverage on calibration set. Returns (lower, upper) bounds.""" + lower = self._quantile(X_raw, alpha / 2) + upper = self._quantile(X_raw, 1 - alpha / 2) + Q = self.conformal.get(alpha, 0.0) + return np.maximum(0, lower - Q), upper + Q + +# ------------ RQS loss ---------------- +def rqs_loss(model, X, y0, w) -> tuple[jax.Array, jax.Array]: + """"Weighted NLL as (total_loss, normalization) rather than a pre-divided mean.""" + log_prob = model.log_prob(y0, X) + return -jnp.sum(w * log_prob), jnp.sum(w) + +# ------------ RQS utils ---------------- +def _spline_knots(raw_widths: jax.Array, raw_heights: jax.Array, raw_derivatives: jax.Array, bounds: int, n_points: int): + """Map unconstrained net outputs to positive bin sizes/derivatives and knot coordinates. + + Widths and heights are softmax'd to sum to 2*bounds, so their cumulative sums + (starting at -bounds) land exactly on +bounds. Derivatives are softplus'd to + stay positive. + """ + widths = jax.nn.softmax(raw_widths, axis=-1) * (2 * bounds) + heights = jax.nn.softmax(raw_heights, axis=-1) * (2 * bounds) + derivatives = jax.nn.softplus(raw_derivatives) + 1e-3 + + knot_x = jnp.concatenate([jnp.full((n_points, 1), -bounds), -bounds + jnp.cumsum(widths, axis=-1)], axis=-1) # (n_points, K+1) + knot_y = jnp.concatenate([jnp.full((n_points, 1), -bounds), -bounds + jnp.cumsum(heights, axis=-1)], axis=-1) # (n_points, K+1) + return knot_x, knot_y, derivatives + + +def _locate_bin(x: jax.Array, knots: jax.Array, n_bins: int) -> jax.Array: + """Index k of the bin containing x, i.e. knots[k] <= x < knots[k+1].""" + bin_idx = jnp.sum((x[..., None] >= knots[:, :-1]).astype(jnp.int32), axis=-1) - 1 + return jnp.clip(bin_idx, 0, n_bins - 1) + + +def _gather_bin(knot: jax.Array, bin_idx: jax.Array): + """Return (knot[i, bin_idx[i]], knot[i, bin_idx[i] + 1]) for every row i.""" + lo = jnp.take_along_axis(knot, bin_idx[:, None], axis=1)[:, 0] + hi = jnp.take_along_axis(knot, bin_idx[:, None] + 1, axis=1)[:, 0] + return lo, hi + + +def _rqs_logdet(theta: jax.Array, s: jax.Array, d_lo: jax.Array, d_hi: jax.Array): + """log|dz/dx| at spline-local parameter theta in [0, 1] (Durkan et al. 2019, eq. 5). + + Also returns the shared denominator and (1 - theta), which the forward pass reuses. + """ + theta_comp = 1.0 - theta + denom = s + (d_hi + d_lo - 2 * s) * theta * theta_comp + deriv_numer = s**2 * (d_hi * theta**2 + 2 * s * theta * theta_comp + d_lo * theta_comp**2) + log_abs_det = jnp.log(deriv_numer) - 2 * jnp.log(denom) + return log_abs_det, denom, theta_comp + + +def _solve_theta(z: jax.Array, y_lo: jax.Array, dy: jax.Array, s: jax.Array, d_lo: jax.Array, d_hi: jax.Array): + """Invert eq. for theta given a target z: solve a*theta^2 + b*theta + c = 0. + + """ + dz = z - y_lo + slope_term = d_hi + d_lo - 2 * s + a = dy * (s - d_lo) + dz * slope_term + b = dy * d_lo - dz * slope_term + c = -s * dz + return 2 * c / (-b - jnp.sqrt(jnp.maximum(b**2 - 4 * a * c, 0.0))) + + +def _rqs(x: jax.Array, raw_widths: jax.Array, raw_heights: jax.Array, raw_derivatives: jax.Array, bounds: int, inverse=False): + """Evaluate the monotone rational-quadratic spline (Durkan et al. 2019, "Neural Spline Flows"). + + Args: + x : (B,) points to transform. Will be z for inverse=False and y for inverse=True. + raw_widths : (B,K) unconstrained bin widths. + raw_heights : (B,K) unconstrained bin heights. + raw_derivatives : (B,K+1) unconstrained knot derivatives. + bounds : the spline is the identity outside [-bounds, bounds]. + inverse : False computes z = T(x); True computes x = T^{-1}(z). + + Returns: + (transformed, log|d(transformed)/dx|). + """ + n_points, n_bins = raw_widths.shape + knot_x, knot_y, derivatives = _spline_knots(raw_widths, raw_heights, raw_derivatives, bounds, n_points) + + in_domain = (x > -bounds) & (x < bounds) + x_clamped = jnp.clip(x, -bounds + 1e-6, bounds - 1e-6) + + # forward looks up knot_x for x; inverse looks up knot_y for z. + bin_idx = _locate_bin(x_clamped, knot_y if inverse else knot_x, n_bins) + x_lo, x_hi = _gather_bin(knot_x, bin_idx) + y_lo, y_hi = _gather_bin(knot_y, bin_idx) + d_lo, d_hi = _gather_bin(derivatives, bin_idx) + dx, dy = x_hi - x_lo, y_hi - y_lo + s = dy / dx # bin slope + + if inverse: + theta = _solve_theta(x_clamped, y_lo, dy, s, d_lo, d_hi) + log_dzdx, _, _ = _rqs_logdet(theta, s, d_lo, d_hi) + out = theta * dx + x_lo # x = T^{-1}(z) + log_abs_det = -log_dzdx # d(T^{-1})/dz = 1 / (dz/dx) + else: + theta = (x_clamped - x_lo) / dx + log_dzdx, denom, theta_comp = _rqs_logdet(theta, s, d_lo, d_hi) + numer = dy * (s * theta**2 + d_lo * theta * theta_comp) + out = y_lo + numer / denom # z = T(x) + log_abs_det = log_dzdx + + return jnp.where(in_domain, out, x), jnp.where(in_domain, log_abs_det, 0.0) + diff --git a/src/estimint/v2/train_base.py b/src/estimint/v2/train_base.py new file mode 100644 index 0000000..9c7120b --- /dev/null +++ b/src/estimint/v2/train_base.py @@ -0,0 +1,107 @@ +import logging +from pathlib import Path + +import duckdb +import hydra +import jax +from omegaconf import DictConfig, OmegaConf + +from .data.preprocess import PreparedData, prepare_data +from .data.dataset import make_loader +import wandb +from .data.features import get_features +from .training.train_step import train_model +import numpy as np +from estimint.v2.eval.metrics import compute_metrics +from .models.rqs import ConditionalRQS, rqs_loss, RQSArtifact +from .training.calibrate import conformal_offset +import json + +log = logging.getLogger(__name__) + +def train_rqs(cfg: DictConfig, prepared_data: PreparedData): + features = get_features(cfg.predictor) + model = ConditionalRQS.from_cfg(cfg, n_context=len(features)) + model = train_model(model, cfg, prepared_data, rqs_loss, name="RQS", use_standardized_y=True) + + rqs_artifact = RQSArtifact(model, prepared_data.feature_scaler, prepared_data.target_scaler, features=features) + + # ------------ calibration ------------------- + calib_loader = make_loader( + data=prepared_data.calib_data, + batch_size=len(prepared_data.calib_data), # load all at once + seed=cfg.seed, + shuffle=False, + num_workers=cfg.num_workers, + drop_remainder=True, + ) + for batch in calib_loader: + calib_x_raw, calib_y_raw = batch["x_raw"], batch["y_raw"] + lower, upper = rqs_artifact.quantile(calib_x_raw, 0.05), rqs_artifact.quantile(calib_x_raw, 0.95) + rqs_artifact.conformal[0.10] = conformal_offset(lower, upper, calib_y_raw, alpha=0.10) + + with open(cfg.conformal_file, "w") as f: + json.dump(rqs_artifact.conformal, f, indent=2) + log.info(f"Saved conformal offsets to {cfg.conformal_file}") + + + # ------------ test evaluation ---------------- + test_loader = make_loader( + data=prepared_data.test_data, + batch_size=cfg.batch_size, + seed=cfg.seed, + shuffle=False, + num_workers=cfg.num_workers, + drop_remainder=False, + ) + metrics = compute_metrics(rqs_artifact, test_loader) + log.info(f"test R2={metrics.r2:.4f} RMSE={metrics.rmse:.2f} MAE={metrics.mae:.2f} MSE={metrics.mse:.2f} Bias={metrics.bias:.2f} Median APE={metrics.medape:.2f} Log10 MSE={metrics.log10_mse:.2f}") + if cfg.use_wandb: + wandb.log({"test/r2": metrics.r2, "test/rmse": metrics.rmse, "test/mae": metrics.mae, "test/mse": metrics.mse, "test/bias": metrics.bias, "test/medape": metrics.medape, "test/log10_mse": metrics.log10_mse}) + + # confidence interval evaluation + test_loader = make_loader( + data=prepared_data.test_data, + batch_size=len(prepared_data.test_data), # load all at once + seed=cfg.seed, + shuffle=False, + num_workers=cfg.num_workers, + drop_remainder=True, + ) + for batch in test_loader: + test_x_raw, test_y_raw = batch["x_raw"], batch["y_raw"] + raw_lower, raw_upper = rqs_artifact.quantile(test_x_raw, 0.05), rqs_artifact.quantile(test_x_raw, 0.95) + conformal_lower, conformal_upper = rqs_artifact.interval(test_x_raw, alpha=0.10) + + coverage_raw = np.mean((test_y_raw >= raw_lower) & (test_y_raw <= raw_upper)) + coverage_conformal = np.mean((test_y_raw >= conformal_lower) & (test_y_raw <= conformal_upper)) + log.info(f"Raw 90% interval coverage: {coverage_raw:.4f}") + log.info(f"Conformal 90% interval coverage: {coverage_conformal:.4f}") + if cfg.use_wandb: + wandb.log({"test/raw_coverage": coverage_raw, "test/conformal_coverage": coverage_conformal}) + return rqs_artifact + +@hydra.main(version_base=None, config_path="conf", config_name="train_config") +def main(cfg: DictConfig) -> None: + log.info(OmegaConf.to_yaml(cfg)) + log.info("JAX devices: %s", jax.devices()) + if cfg.use_wandb: + wandb.init( + project=cfg.wandb.project, + name=cfg.wandb.name, + config=OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True), # type: ignore + settings=wandb.Settings(start_method="thread"), + ) + + Path(cfg.output_dir).mkdir(parents=True, exist_ok=True) + raw_df = duckdb.read_parquet(cfg.data_file).df() + prepared_data = prepare_data(raw_df, cfg, calib_frac=cfg.calib_frac) + + train_rqs(cfg, prepared_data) + + if cfg.use_wandb: + wandb.finish() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/estimint/v2/training/calibrate.py b/src/estimint/v2/training/calibrate.py new file mode 100644 index 0000000..53be3cb --- /dev/null +++ b/src/estimint/v2/training/calibrate.py @@ -0,0 +1,14 @@ +import numpy as np + +def conformal_offset(lower: np.ndarray, upper: np.ndarray, y_true: np.ndarray, alpha: float = 0.10) -> float: + """Split-conformal (CQR, Romano 2019) width correction. + + Returns offset Q so that widening to [lo-Q, hi+Q] gives >= (1-alpha) + marginal coverage. + """ + if not (0.0 < alpha < 1.0): + raise ValueError(f"alpha must be in (0,1), got {alpha}") + scores = np.maximum(lower - y_true, y_true - upper) + n = len(scores) + k = np.ceil((n + 1) * (1 - alpha)).astype(int) + return float(np.sort(scores)[k - 1]) diff --git a/src/estimint/v2/training/checkpoint.py b/src/estimint/v2/training/checkpoint.py new file mode 100644 index 0000000..a4598fd --- /dev/null +++ b/src/estimint/v2/training/checkpoint.py @@ -0,0 +1,50 @@ +import logging +import flax.nnx as nnx +from orbax.checkpoint import v1 as ocp +from etils import epath + +log = logging.getLogger(__name__) + +# Orbax checkpointing logs verbosely via the absl logger; silence INFO-level noise. +logging.getLogger("absl").setLevel(logging.WARNING) + + +def _resolve_checkpoint_dir(checkpoint_dir: str, model_name: str) -> epath.Path: + return (epath.Path(checkpoint_dir) / model_name).resolve() + +def restore_model( + checkpoint_dir: str, + model_name: str, + model: nnx.Module, +) -> nnx.Module: + """ + Restore model from checkpoint. + + Args: + ckptr: Orbax checkpointer. + model: Model to restore. + step: Checkpoint step to restore. If None, restore the latest checkpoint. + + Returns: + Restored model. + """ + ckpt_dir = _resolve_checkpoint_dir(checkpoint_dir, model_name) + with ocp.training.Checkpointer(ckpt_dir) as ckptr: + loaded = ckptr.load_checkpointables( + abstract_checkpointables={"model": nnx.state(model)}, + ) + nnx.update(model, loaded["model"]) + return model + + +def save_checkpoint(checkpoint_dir: str, model_name: str, model: nnx.Module): + ckpt_dir = _resolve_checkpoint_dir(checkpoint_dir, model_name) + preservation_policy = ocp.training.preservation_policies.LatestN(n=1) + with ocp.training.Checkpointer(ckpt_dir, preservation_policy=preservation_policy) as ckptr: # type: ignore[arg-type] + ckptr.save_checkpointables( + 0, + { + "model": nnx.state(model), + }, + overwrite=True + ) diff --git a/src/estimint/v2/training/train_step.py b/src/estimint/v2/training/train_step.py new file mode 100644 index 0000000..df62c35 --- /dev/null +++ b/src/estimint/v2/training/train_step.py @@ -0,0 +1,154 @@ +import logging + +from flax import nnx +import optax +from typing import Callable +from jaxtyping import Array +import jax +import numpy as np +from ..data.preprocess import PreparedData +from omegaconf import DictConfig +import tqdm +import jax.numpy as jnp +import wandb +from ..data.dataset import make_loader +from tqdm import tqdm +from .checkpoint import save_checkpoint +from ..common.types import LossFn + + +log = logging.getLogger(__name__) + + +def create_optimizer( + model: nnx.Module, learning_rate: float, total_steps: int, weight_decay: float = 1e-4 +) -> nnx.Optimizer: + """ + Create the training optimizer. + + Args: + model: Model to optimize. + learning_rate: Peak learning rate. + total_steps: Total training steps. + + Returns: + Configured optimizer. + """ + scheduler = optax.warmup_cosine_decay_schedule( + init_value=0.0, + peak_value=learning_rate, + warmup_steps=int(0.03 * total_steps), # warmup for 3% of training + decay_steps=total_steps, + end_value=0.01 * learning_rate, # decay to 1% of initial LR + ) + tx = optax.chain(optax.clip_by_global_norm(1.0), optax.adamw(learning_rate=scheduler, weight_decay=weight_decay)) + return nnx.Optimizer(model, tx, wrt=nnx.Param) + +def get_total_params(model: nnx.Module) -> int: + """ + Get the total number of parameters in the model. + + Args: + model: Flax module. + + Returns: + Total parameter count. + """ + params = nnx.state(model, nnx.Param) + return sum(np.prod(x.shape) for x in jax.tree_util.tree_leaves(params)) + +def _weighted_mean(terms: list[tuple[Array, Array]]) -> float: + """Combine per-batch (total_losses, normalizations) pairs into the loss over all records.""" + total_losses, normalizations = (jnp.stack(t) for t in zip(*terms)) + return float(jnp.sum(total_losses) / jnp.sum(normalizations)) + +def make_train_step(loss_fn: LossFn): + @nnx.jit + def train_step(model: nnx.Module, optimizer: nnx.Optimizer, x: Array, y: Array, w: Array): + def compute_total_loss(model: nnx.Module, x: Array, y: Array, w: Array): + loss_sum, normalization = loss_fn(model, x, y, w) + return loss_sum / normalization, (loss_sum, normalization) + + (_, (loss_sum, normalization)), grads = nnx.value_and_grad(compute_total_loss, has_aux=True)(model, x, y, w) + optimizer.update(model, grads) + return loss_sum, normalization + + return train_step + +def make_eval_step(loss_fn: LossFn): + @nnx.jit + def eval_step(model: nnx.Module, x: Array, y: Array, w: Array): + return loss_fn(model, x, y, w) + + return eval_step + + +def train_model( + model: nnx.Module, + cfg: DictConfig, + prepared_data: PreparedData, + loss_fn: LossFn, + name: str, + use_standardized_y: bool = False, + ) -> nnx.Module: + + log.info(f"Total parameters: {get_total_params(model) / 1e6:.2f}M") + + train_step = make_train_step(loss_fn) + eval_step = make_eval_step(loss_fn) + target_key = "y_std" if use_standardized_y else "y" + val_loader = make_loader( + data=prepared_data.val_data, + batch_size=cfg.batch_size, + seed=cfg.seed, + shuffle=False, + num_workers=cfg.num_workers, + drop_remainder=False, + ) + total_steps = cfg.num_epochs * len(prepared_data.train_data) // cfg.batch_size + optimizer = create_optimizer(model, cfg.lr, total_steps, weight_decay=cfg.weight_decay) + + # ---- training loop ---- + patience_n = 0 + best_val_loss = float("inf") + best_state = jax.tree.map(lambda x: x, nnx.state(model)) + epoch_pbar = tqdm(range(cfg.num_epochs), desc="Epoch") + for epoch in epoch_pbar: + # remake train loader each epoch to reshuffle with new seed + train_loader = make_loader( + data=prepared_data.train_data, + batch_size=cfg.batch_size, + seed=cfg.seed + epoch, + shuffle=True, + num_workers=cfg.num_workers, + drop_remainder=True, + ) + model.train() + train_terms = [train_step(model, optimizer, batch["x"], batch[target_key], batch["w"]) for batch in train_loader] + + model.eval() + val_terms = [eval_step(model, batch["x"], batch[target_key], batch["w"]) for batch in val_loader] + + avg_train_loss = _weighted_mean(train_terms) + avg_val_loss = _weighted_mean(val_terms) + epoch_pbar.set_postfix( + train=f"{avg_train_loss:.6f}", + val=f"{avg_val_loss:.6f}", + patience=f"{patience_n}/{cfg.patience}", + ) + if cfg.use_wandb: + wandb.log({"train/loss": avg_train_loss, "val/loss": avg_val_loss, "epoch": epoch}) + if avg_val_loss < best_val_loss: + best_val_loss = avg_val_loss + best_state = jax.tree.map(lambda x: x, nnx.state(model)) + patience_n = 0 + else: + patience_n += 1 + if epoch >= cfg.min_epochs and patience_n >= cfg.patience: + log.info(f"Early stopping at epoch {epoch} with best val loss {best_val_loss:.6f}") + break + + nnx.update(model, best_state) + save_checkpoint(cfg.checkpoint_dir, name, model) + + return model \ No newline at end of file diff --git a/tests/test_estimint.py b/tests/test_estimint.py index f8a38c1..03ec79f 100644 --- a/tests/test_estimint.py +++ b/tests/test_estimint.py @@ -8,94 +8,76 @@ class TestUtils: """Test utility functions.""" - + def test_r2(self): from estimint import r2 y = np.array([1, 2, 3, 4, 5]) yhat = np.array([1.1, 2.0, 2.9, 4.0, 5.1]) result = r2(y, yhat) assert 0.95 < result <= 1.0 - + def test_rmse(self): from estimint import rmse y = np.array([1, 2, 3]) yhat = np.array([1, 2, 3]) assert rmse(y, yhat) == 0.0 - + yhat2 = np.array([2, 3, 4]) assert rmse(y, yhat2) == 1.0 - + def test_mse(self): from estimint import mse y = np.array([1, 2, 3]) yhat = np.array([2, 3, 4]) assert mse(y, yhat) == 1.0 - + def test_mae(self): from estimint import mae y = np.array([1, 2, 3]) yhat = np.array([2, 3, 4]) assert mae(y, yhat) == 1.0 - + def test_median_ae(self): from estimint import median_ae y = np.array([1, 2, 3]) yhat = np.array([2, 3, 4]) assert median_ae(y, yhat) == 1.0 - + def test_safe_div(self): from estimint import safe_div result = safe_div(np.array([1, 2]), np.array([0, 2])) assert result[1] == 1.0 assert result[0] > 0 # Should not be inf - + def test_fit_predict_qmap(self): from estimint import fit_qmap_w, predict_qmap_w pred = np.array([1, 2, 3, 4, 5]) obs = np.array([1.5, 2.5, 3.5, 4.5, 5.5]) - + cal = fit_qmap_w(pred, obs) assert "xq" in cal assert "yq" in cal - + calibrated = predict_qmap_w(pred, cal) assert len(calibrated) == len(pred) - + def test_scale_pos(self): from estimint import scale_pos obs = np.array([1, 2, 3]) pred = np.array([0.5, 1, 1.5]) - + a = scale_pos(obs, pred) assert a > 0 class TestDataProcessing: """Test data processing functions.""" - + def test_make_value_weights(self): - from estimint import make_value_weights + from estimint.data_processing import make_value_weights eir = np.array([1.0, 1.0, 2.0, 2.0, 2.0, 3.0]) weights = make_value_weights(eir, digits=3) - + assert len(weights) == len(eir) assert np.isclose(np.mean(weights), 1.0) # Normalized to mean=1 - -class TestRun: - """Test model inference functions.""" - - def test_global_model_functions(self): - from estimint.run import set_global_model, get_global_model - - assert get_global_model() is None - - dummy_model = {"class": "estiMINT_model", "features": ["a", "b"]} - set_global_model(dummy_model) - - assert get_global_model() is not None - assert get_global_model()["class"] == "estiMINT_model" - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/test_flows.py b/tests/test_flows.py index 4d7a164..40119e5 100644 --- a/tests/test_flows.py +++ b/tests/test_flows.py @@ -1,57 +1,53 @@ """End-to-end flow tests: prevalence -> EIR, and the mosquito-delta HBR pipeline. -These exercise the bundled models in src/estimint/data, so they run offline. +The prevalence -> EIR flow exercises the bundled XGBoost models in src/estimint/data, +so it runs offline. The mosquito-delta pipeline uses the published RQS models, which +are downloaded (and then cached) from the estiMINT HuggingFace repo. """ -import pandas as pd import pytest from estimint import ( - load_xgb_model, - run_xgb_model, estimate_eir_with_mosquito_delta, ) +from estimint.eir_models import load_eir_models +from estimint.types import EirTarget, PreparedScenario INTERVENTIONS = dict( dn0_use=0.33, Q0=0.87, phi_bednets=0.82, seasonal=0.0, itn_use=0.6, irs_use=0.0, ) - -class TestPrevalenceToEir: - def test_predicts_positive_eir(self): - model = load_xgb_model("prevalence") - X = pd.DataFrame({"prev_y9": [0.30], **{k: [v] for k, v in INTERVENTIONS.items()}}) - eir = run_xgb_model(X, model) - assert len(eir) == 1 - assert eir[0] > 0 - - def test_higher_prevalence_gives_higher_eir(self): - model = load_xgb_model("prevalence") - rows = {"prev_y9": [0.10, 0.50], **{k: [v, v] for k, v in INTERVENTIONS.items()}} - eir = run_xgb_model(pd.DataFrame(rows), model) - assert eir[1] > eir[0] +def _prepared(delta: float, prevalence: float = 0.30) -> PreparedScenario: + """A prevalence-input scenario carrying a mosquito-density change.""" + return PreparedScenario( + eir_target=EirTarget(prevalence, "prevalence"), + mosquito_density_change=delta, + eir_model_features=dict(INTERVENTIONS), + summary_values={}, + emulator_covariates={}, + ) class TestMosquitoDelta: @pytest.fixture(scope="class") def models(self): - return {name: load_xgb_model(name) for name in ("prevalence", "hbr", "eir_to_hbr")} + return load_eir_models() def _run(self, models, delta): - inputs = pd.DataFrame([{"prevalence": 0.30, "mosquito_delta": delta, **INTERVENTIONS}]) - return estimate_eir_with_mosquito_delta(inputs, models=models).iloc[0] + return estimate_eir_with_mosquito_delta([_prepared(delta)], eir_models=models)[0] - def test_returns_expected_columns(self, models): + def test_returns_expected_keys(self, models): res = self._run(models, 0.25) - assert set(res.index) == { + assert set(res) == { "eir_baseline", "eir_new", "eir_multiplier", "hbr_baseline", "hbr_new", } + assert all(isinstance(value, float) for value in res.values()) def test_zero_delta_is_identity(self, models): res = self._run(models, 0.0) - assert res["eir_new"] == res["eir_baseline"] - assert res["eir_multiplier"] == 1.0 + assert res["eir_new"] == pytest.approx(res["eir_baseline"]) + assert res["eir_multiplier"] == pytest.approx(1.0) def test_more_mosquitoes_raises_eir(self, models): res = self._run(models, 0.25) @@ -67,7 +63,7 @@ def test_fewer_mosquitoes_lowers_eir(self, models): def test_batch_is_monotonic_in_delta(self, models): # a single batched call handles every row and preserves input order deltas = [-0.5, -0.25, 0.0, 0.25, 0.5, 1.0] - inputs = pd.DataFrame([{"prevalence": 0.30, "mosquito_delta": d, **INTERVENTIONS} for d in deltas]) - res = estimate_eir_with_mosquito_delta(inputs, models=models) - assert list(res.index) == list(range(len(deltas))) - assert list(res["eir_new"]) == sorted(res["eir_new"]) + results = estimate_eir_with_mosquito_delta([_prepared(d) for d in deltas], eir_models=models) + assert len(results) == len(deltas) + eir_new = [result["eir_new"] for result in results] + assert eir_new == sorted(eir_new) diff --git a/tests/test_scenarios.py b/tests/test_scenarios.py index e986d37..00222d1 100644 --- a/tests/test_scenarios.py +++ b/tests/test_scenarios.py @@ -10,12 +10,12 @@ import numpy as np import pytest +from estimint.eir_models import EirModels, load_eir_models from estimint.scenarios import ( - _PreparedScenario, + PreparedScenario, _apply_mosquito_delta_batch, _classify_prepared_scenario, _estimate_eir, - _load_eir_hbr_models, _prepare_scenario_inputs, run_scenarios, ) @@ -32,14 +32,14 @@ def mk(**kwargs: Any) -> Scenario: return Scenario(eir_target=EirTarget(input_value, input_mode), **defaults) -def _estimate_eir_single(scenario: Scenario, eir_models: dict[str, Any]) -> _PreparedScenario: +def _estimate_eir_single(scenario: Scenario, eir_models: EirModels) -> PreparedScenario: """Estimate EIR for a single scenario via the batch estimator.""" return _estimate_eir([scenario], eir_models)[0] @pytest.fixture(scope="module") def est(): - return _load_eir_hbr_models() + return load_eir_models() class TestEstimateEir: @@ -174,11 +174,11 @@ def test_batch_mixed_input_modes(self, est): class TestClassifyPreparedScenario: - def _make_prepared(self, input_mode: str, mosquito_density_change: float = 0.0) -> _PreparedScenario: + def _make_prepared(self, input_mode: str, mosquito_density_change: float = 0.0) -> PreparedScenario: from typing import cast from estimint.types import Input_Mode - return _PreparedScenario( + return PreparedScenario( eir_target=EirTarget(input_value=10.0, input_mode=cast(Input_Mode, input_mode)), mosquito_density_change=mosquito_density_change, eir_model_features={}, diff --git a/tests/v2/__init__.py b/tests/v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/v2/data/__init__.py b/tests/v2/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/v2/data/test_dataset.py b/tests/v2/data/test_dataset.py new file mode 100644 index 0000000..08987d0 --- /dev/null +++ b/tests/v2/data/test_dataset.py @@ -0,0 +1,73 @@ +""" +Tests for the Grain data source and loader. +""" + +import numpy as np +import pytest + +from estimint.v2.data.dataset import DataSource, make_loader + + +def make_records(n=10, n_features=3): + return [ + { + "x": np.full(n_features, i, dtype=np.float32), + "y_std": np.float32(i), + "w": np.float32(1.0), + "ps": np.array([i, 0], dtype=np.int32), + } + for i in range(n) + ] + + +@pytest.fixture +def records(): + return make_records() + + +class TestDataSource: + def test_length_matches_the_records(self, records): + assert len(DataSource(records)) == len(records) + + def test_getitem_returns_the_record(self, records): + assert DataSource(records)[3] is records[3] + + +class TestMakeLoader: + def test_batches_have_a_leading_batch_dimension(self, records): + batch = next(iter(make_loader(records, batch_size=5))) + assert batch["x"].shape == (5, 3) + assert batch["y_std"].shape == (5,) + assert batch["ps"].shape == (5, 2) + + def test_drop_remainder_discards_the_partial_batch(self, records): + batches = list(make_loader(records, batch_size=4, drop_remainder=True)) + assert [b["y_std"].shape[0] for b in batches] == [4, 4] + + def test_keeping_the_remainder_yields_every_record(self, records): + batches = list(make_loader(records, batch_size=4, drop_remainder=False)) + assert [b["y_std"].shape[0] for b in batches] == [4, 4, 2] + + def test_unshuffled_loader_preserves_record_order(self, records): + batches = list(make_loader(records, batch_size=5, shuffle=False)) + seen = np.concatenate([b["y_std"] for b in batches]) + np.testing.assert_array_equal(seen, np.arange(len(records), dtype=np.float32)) + + def test_shuffle_reorders_records_without_losing_any(self, records): + batches = list(make_loader(records, batch_size=5, shuffle=True, seed=0)) + seen = np.concatenate([b["y_std"] for b in batches]) + assert sorted(seen) == list(np.arange(len(records), dtype=np.float32)) + + def test_same_seed_gives_the_same_shuffle(self, records): + order = lambda seed: np.concatenate( + [b["y_std"] for b in make_loader(records, batch_size=5, shuffle=True, seed=seed)] + ) + np.testing.assert_array_equal(order(0), order(0)) + assert not np.array_equal(order(0), order(1)) + + def test_loader_covers_exactly_one_epoch(self, records): + assert sum(b["y_std"].shape[0] for b in make_loader(records, batch_size=2)) == len(records) + + def test_full_batch_loads_everything_at_once(self, records): + batches = list(make_loader(records, batch_size=len(records))) + assert len(batches) == 1 and batches[0]["x"].shape == (len(records), 3) diff --git a/tests/v2/data/test_features.py b/tests/v2/data/test_features.py new file mode 100644 index 0000000..4bed9b1 --- /dev/null +++ b/tests/v2/data/test_features.py @@ -0,0 +1,59 @@ +""" +Tests for the v2 feature naming and standardization helpers. +""" + +import numpy as np +import pytest + +from estimint.v2.data.features import FEATURES_BASE, StandardScaler, get_features + + +class TestGetFeatures: + @pytest.mark.parametrize("predictor", ["prev_y9", "eir", "hbr_y9"]) + def test_predictor_leads_the_feature_list(self, predictor): + features = get_features(predictor) + assert features[0] == predictor + assert features[1:] == FEATURES_BASE + + def test_feature_names_are_unique(self): + features = get_features("prev_y9") + assert len(set(features)) == len(features) + + +class TestStandardScaler: + @pytest.fixture + def X(self): + return np.array([[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]]) + + def test_unfitted_scaler_raises(self, X): + scaler = StandardScaler() + assert not scaler.is_fitted + with pytest.raises(ValueError, match="not fitted"): + scaler.transform(X) + with pytest.raises(ValueError, match="not fitted"): + scaler.inverse_transform(X) + + def test_fit_transform_standardizes_columns(self, X): + Z = StandardScaler().fit_transform(X) + np.testing.assert_allclose(Z.mean(axis=0), 0.0, atol=1e-12) + np.testing.assert_allclose(Z.std(axis=0), 1.0) + + def test_inverse_transform_round_trips(self, X): + scaler = StandardScaler().fit(X) + np.testing.assert_allclose(scaler.inverse_transform(scaler.transform(X)), X) + + def test_constant_column_does_not_divide_by_zero(self): + X = np.array([[5.0, 1.0], [5.0, 2.0], [5.0, 3.0]]) + scaler = StandardScaler().fit(X) + assert scaler.scale_[0] == 1.0 + assert np.all(np.isfinite(scaler.transform(X))) + + def test_fit_returns_self(self, X): + scaler = StandardScaler() + assert scaler.fit(X) is scaler + assert scaler.is_fitted + + def test_transform_uses_train_statistics(self, X): + scaler = StandardScaler().fit(X) + # a row equal to the training mean maps to zero + np.testing.assert_allclose(scaler.transform(X.mean(axis=0)[None, :]), 0.0, atol=1e-12) diff --git a/tests/v2/data/test_preprocess.py b/tests/v2/data/test_preprocess.py new file mode 100644 index 0000000..0c090fa --- /dev/null +++ b/tests/v2/data/test_preprocess.py @@ -0,0 +1,299 @@ +""" +Tests for the v2 data preparation pipeline: filtering, splitting, scaling, record building. +""" + +import pickle + +import numpy as np +import pandas as pd +import pytest +from omegaconf import OmegaConf + +from estimint.v2.data.features import FEATURES_BASE, StandardScaler, get_features +from estimint.v2.data.preprocess import ( + SplitParamSims, + _assign_param_group, + _assign_param_splits, + _build_data, + _create_split, + _filter_by_threshold, + _fit_features_scaler, + _fit_target_scaler, + _load_split, + _parameter_strata, + _save_split, + prepare_data, +) + +N_PARAMS = 20 +N_SIMS = 2 +ROWS_PER_SIM = 3 + + +def make_df(n_params=N_PARAMS, n_sims=N_SIMS, rows_per_sim=ROWS_PER_SIM, low_prev_params=()): + """Synthetic simulation frame: constant feature/target values within a (param, sim) pair.""" + rng = np.random.default_rng(0) + rows = [] + for p in range(n_params): + for s in range(n_sims): + static = {f: float(rng.uniform(0, 1)) for f in FEATURES_BASE} + static["prev_y9"] = 0.001 if p in low_prev_params else float(rng.uniform(0.05, 0.8)) + static["hbr_y9"] = float(rng.uniform(1, 100)) + static["eir"] = float(rng.uniform(0.1, 500)) + rows += [{"parameter_index": p, "simulation_index": s, **static}] * rows_per_sim + return pd.DataFrame(rows) + + +@pytest.fixture +def df(): + return make_df() + + +@pytest.fixture +def filtered_df(df): + return _filter_by_threshold(df) + + +@pytest.fixture +def cfg(tmp_path): + return OmegaConf.create( + { + "seed": 0, + "predictor": "prev_y9", + "target": "eir", + "stratify": False, + "use_existing_split": False, + "split_file": str(tmp_path / "split.csv"), + "output_dir": str(tmp_path / "out"), + } + ) + + +def param_indices(pairs): + return {p for p, _ in pairs} + + +class TestFilterByThreshold: + def test_drops_pairs_below_prevalence_threshold(self): + df = make_df(n_params=6, low_prev_params=(1, 4)) + out = _filter_by_threshold(df) + assert param_indices(out["_ps"]) == {0, 2, 3, 5} + + def test_adds_pair_column(self, filtered_df): + assert list(filtered_df["_ps"].iloc[0]) == [ + filtered_df["parameter_index"].iloc[0], + filtered_df["simulation_index"].iloc[0], + ] + + def test_keeps_everything_above_threshold(self, df, filtered_df): + assert len(filtered_df) == len(df) + + +class TestSplitting: + def test_splits_are_disjoint_and_cover_all_pairs(self, filtered_df): + split = _create_split(filtered_df, seed=0, calib_frac=0.1) + all_pairs = set(filtered_df["_ps"]) + parts = [split.train, split.val, split.calib, split.test] + assert set().union(*parts) == all_pairs + assert sum(len(p) for p in parts) == len(all_pairs) + + def test_a_parameter_never_straddles_two_splits(self, filtered_df): + split = _create_split(filtered_df, seed=0, calib_frac=0.1) + groups = [param_indices(p) for p in (split.train, split.val, split.calib, split.test)] + for i, a in enumerate(groups): + for b in groups[i + 1 :]: + assert not (a & b) + + def test_split_is_deterministic_given_a_seed(self, filtered_df): + assert _create_split(filtered_df, seed=0).train == _create_split(filtered_df, seed=0).train + + def test_different_seeds_give_different_splits(self, filtered_df): + assert _create_split(filtered_df, seed=0).train != _create_split(filtered_df, seed=1).train + + def test_train_gets_roughly_seventy_percent(self, filtered_df): + split = _create_split(filtered_df, seed=0, stratify=False) + assert len(param_indices(split.train)) == pytest.approx(0.7 * N_PARAMS, abs=1) + + def test_calibration_split_is_empty_by_default(self, filtered_df): + assert _create_split(filtered_df, seed=0).calib == set() + + def test_rejects_calibration_fraction_that_starves_validation(self, filtered_df): + with pytest.raises(ValueError, match="Validation fraction"): + _assign_param_splits(filtered_df, seed=0, calib_frac=0.4) + + def test_assign_param_group_respects_fractions(self): + assign = _assign_param_group( + np.arange(100), np.random.default_rng(0), train_frac=0.7, val_frac=0.15, calib_frac=0.05 + ) + counts = pd.Series(list(assign.values())).value_counts() + assert counts["train"] == 70 and counts["val"] == 15 and counts["calib"] == 5 and counts["test"] == 10 + + +class TestParameterStrata: + def test_unstratified_returns_one_group(self, filtered_df): + strata = _parameter_strata(filtered_df, "eir", n_bins=10, stratify=False) + assert len(strata) == 1 + assert len(strata[0]) == N_PARAMS + + def test_stratified_partitions_every_parameter_once(self, filtered_df): + strata = _parameter_strata(filtered_df, "eir", n_bins=5, stratify=True) + assert len(strata) <= 5 + assert sorted(np.concatenate(strata)) == sorted(filtered_df["parameter_index"].unique()) + + def test_strata_are_ordered_by_target_magnitude(self, filtered_df): + strata = _parameter_strata(filtered_df, "eir", n_bins=4, stratify=True) + means = filtered_df.groupby("parameter_index")["eir"].mean() + stratum_means = [means[s].mean() for s in strata] + assert stratum_means == sorted(stratum_means) + + +class TestSplitFileRoundTrip: + def test_save_then_load_preserves_splits(self, filtered_df, tmp_path): + split = _create_split(filtered_df, seed=0, calib_frac=0.1) + path = tmp_path / "split.csv" + _save_split(path, split) + loaded = _load_split(str(path), filtered_df) + assert loaded == split + + def test_load_ignores_pairs_absent_from_the_frame(self, filtered_df, tmp_path): + path = tmp_path / "split.csv" + split = _create_split(filtered_df, seed=0) + _save_split(path, split) + smaller = filtered_df[filtered_df["parameter_index"] < 5] + loaded = _load_split(str(path), smaller) + assert param_indices(loaded.train | loaded.val | loaded.test) <= {0, 1, 2, 3, 4} + + def test_load_maps_csv_split_names(self, filtered_df, tmp_path): + path = tmp_path / "split.csv" + _save_split(path, SplitParamSims(train={(0, 0)}, val={(1, 0)}, calib={(2, 0)}, test={(3, 0)})) + assert set(pd.read_csv(path)["split"]) == {"train", "validate", "calibrate", "test"} + loaded = _load_split(str(path), filtered_df) + assert loaded.val == {(1, 0)} and loaded.calib == {(2, 0)} + + +class TestScalers: + def test_feature_scaler_uses_train_pairs_only(self, filtered_df, tmp_path): + train_ps = {ps for ps in filtered_df["_ps"] if ps[0] < 5} + features = get_features("prev_y9") + scaler = _fit_features_scaler(filtered_df, train_ps, str(tmp_path), features) + + expected = ( + filtered_df[filtered_df["_ps"].isin(train_ps)] + .drop_duplicates(subset=["_ps"])[features] + .to_numpy(dtype=np.float32) + ) + np.testing.assert_allclose(scaler.mean_, expected.mean(axis=0), rtol=1e-5) + + def test_feature_scaler_is_pickled_to_the_output_dir(self, filtered_df, tmp_path): + train_ps = set(filtered_df["_ps"]) + scaler = _fit_features_scaler(filtered_df, train_ps, str(tmp_path), get_features("prev_y9")) + with open(tmp_path / "features_scaler.pkl", "rb") as f: + np.testing.assert_allclose(pickle.load(f).mean_, scaler.mean_) + + def test_target_scaler_is_fitted_in_log_space(self, filtered_df, tmp_path): + train_ps = set(filtered_df["_ps"]) + scaler = _fit_target_scaler(filtered_df, train_ps, str(tmp_path), target="eir") + + eir = filtered_df.drop_duplicates(subset=["_ps"])["eir"].to_numpy(dtype=np.float32) + np.testing.assert_allclose(scaler.mean_, np.log10(eir).mean(), rtol=1e-4) + assert (tmp_path / "target_scaler.pkl").exists() + + def test_wide_range_predictor_is_fitted_in_log_space(self, filtered_df, tmp_path): + # eir spans ~3 decades, so it is standardized in log space like the target. + # The bounded covariates in FEATURES_BASE are left alone. + features = get_features("eir") + scaler = _fit_features_scaler(filtered_df, set(filtered_df["_ps"]), str(tmp_path), features) + + static = filtered_df.drop_duplicates(subset=["_ps"]) + eir = static["eir"].to_numpy(dtype=np.float32) + np.testing.assert_allclose(scaler.mean_[0], np.log10(eir).mean(), rtol=1e-4) + np.testing.assert_allclose( + scaler.mean_[1:], static[FEATURES_BASE].to_numpy(dtype=np.float32).mean(axis=0), rtol=1e-5 + ) + + +class TestBuildData: + @pytest.fixture + def records(self, filtered_df, tmp_path): + filtered_df = filtered_df.assign(_weight=1.0) + features = get_features("prev_y9") + pairs = set(filtered_df["_ps"]) + feature_scaler = _fit_features_scaler(filtered_df, pairs, str(tmp_path), features) + target_scaler = _fit_target_scaler(filtered_df, pairs, str(tmp_path), "eir") + data = _build_data(filtered_df, pairs, feature_scaler, target_scaler, features, "eir") + return data, feature_scaler, target_scaler, features + + def test_one_record_per_pair_with_expected_keys(self, records, filtered_df): + data, *_ = records + assert len(data) == len(set(filtered_df["_ps"])) + assert set(data[0]) == {"x_raw", "x", "y_raw", "y", "y_std", "w", "ps"} + + def test_features_are_scaled_versions_of_the_raw_row(self, records): + data, feature_scaler, _, features = records + record = data[0] + assert record["x_raw"].shape == (len(features),) + np.testing.assert_allclose(record["x"], feature_scaler.transform(record["x_raw"]), rtol=1e-5) + + def test_target_is_log10_then_standardized(self, records): + data, _, target_scaler, _ = records + record = data[0] + assert record["y"] == pytest.approx(np.log10(record["y_raw"]), rel=1e-6) + assert record["y_std"] == pytest.approx( + target_scaler.transform(np.array([[record["y"]]], dtype=np.float32))[0, 0], rel=1e-5 + ) + + def test_unknown_pairs_are_skipped(self, filtered_df, tmp_path): + filtered_df = filtered_df.assign(_weight=1.0) + features = get_features("prev_y9") + scaler = _fit_features_scaler(filtered_df, set(filtered_df["_ps"]), str(tmp_path), features) + target_scaler = _fit_target_scaler(filtered_df, set(filtered_df["_ps"]), str(tmp_path), "eir") + data = _build_data(filtered_df, {(0, 0), (999, 999)}, scaler, target_scaler, features, "eir") + assert len(data) == 1 + + +class TestPrepareData: + def test_end_to_end(self, df, cfg): + prepared = prepare_data(df, cfg, calib_frac=0.1) + + assert prepared.input_size == len(get_features(cfg.predictor)) + assert all(len(d) > 0 for d in (prepared.train_data, prepared.val_data, prepared.calib_data, prepared.test_data)) + n_records = sum(len(d) for d in (prepared.train_data, prepared.val_data, prepared.calib_data, prepared.test_data)) + assert n_records == N_PARAMS * N_SIMS + assert isinstance(prepared.feature_scaler, StandardScaler) and prepared.feature_scaler.is_fitted + + def test_writes_split_and_scalers(self, df, cfg, tmp_path): + prepare_data(df, cfg, calib_frac=0.1) + assert (tmp_path / "split.csv").exists() + assert (tmp_path / "out" / "features_scaler.pkl").exists() + assert (tmp_path / "out" / "target_scaler.pkl").exists() + + def test_reuses_an_existing_split_file(self, df, cfg): + first = prepare_data(df, cfg, calib_frac=0.1) + cfg.use_existing_split = True + cfg.seed = 999 # would produce a different split if it were recreated + second = prepare_data(df, cfg, calib_frac=0.1) + assert second.train_param_sims == first.train_param_sims + + def test_low_prevalence_pairs_are_excluded(self, cfg): + prepared = prepare_data(make_df(low_prev_params=(0, 1)), cfg, calib_frac=0.1) + kept = set().union( + prepared.train_param_sims, prepared.val_param_sims, prepared.calib_param_sims, prepared.test_param_sims + ) + assert not param_indices(kept) & {0, 1} + + def test_scalers_ignore_non_train_pairs(self, df, cfg): + prepared = prepare_data(df, cfg, calib_frac=0.1) + train_x = np.stack([r["x_raw"] for r in prepared.train_data]) + np.testing.assert_allclose(prepared.feature_scaler.mean_, train_x.mean(axis=0), rtol=1e-4) + + def test_feature_scaler_reproduces_every_record_from_its_raw_row(self, df, cfg): + # Training consumes record["x"], while compute_metrics and RQSArtifact re-derive + # the context from record["x_raw"] via prepared.feature_scaler. Any feature + # transform has to live inside the scaler or those two paths drift apart. + cfg.predictor = "hbr_y9" # wide-range predictor, so a log transform is in play + prepared = prepare_data(df, cfg, calib_frac=0.1) + + for split in (prepared.train_data, prepared.val_data, prepared.calib_data, prepared.test_data): + x_raw = np.stack([r["x_raw"] for r in split]) + x = np.stack([r["x"] for r in split]) + np.testing.assert_allclose(x, prepared.feature_scaler.transform(x_raw), rtol=1e-5) diff --git a/tests/v2/eval/__init__.py b/tests/v2/eval/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/v2/eval/test_metrics.py b/tests/v2/eval/test_metrics.py new file mode 100644 index 0000000..6e78923 --- /dev/null +++ b/tests/v2/eval/test_metrics.py @@ -0,0 +1,72 @@ +""" +Tests for evaluation metric aggregation over a data loader. +""" + +import numpy as np +import pytest + +from estimint.v2.data.dataset import make_loader +from estimint.v2.eval.metrics import Metrics, compute_metrics, get_preds_targets + + +class ConstantOffset: + """Stand-in ModelArtifact: predicts the first raw feature plus a fixed offset.""" + + def __init__(self, offset=0.0): + self.offset = offset + + def predict(self, X_raw): + return X_raw[:, 0] + self.offset + + +@pytest.fixture +def records(): + rng = np.random.default_rng(0) + y = rng.uniform(1, 100, size=12).astype(np.float32) + return [{"x_raw": np.array([v, 0.0], dtype=np.float32), "y_raw": v} for v in y] + + +@pytest.fixture +def loader(records): + return make_loader(records, batch_size=4, shuffle=False) + + +class TestGetPredsTargets: + def test_returns_one_value_per_record(self, loader, records): + preds, targets = get_preds_targets(ConstantOffset(), loader) + assert preds.shape == targets.shape == (len(records),) + + def test_targets_follow_loader_order(self, loader, records): + _, targets = get_preds_targets(ConstantOffset(), loader) + np.testing.assert_allclose(targets, [r["y_raw"] for r in records]) + + def test_predictions_come_from_raw_features(self, loader, records): + preds, _ = get_preds_targets(ConstantOffset(offset=2.0), loader) + np.testing.assert_allclose(preds, [r["x_raw"][0] + 2.0 for r in records], rtol=1e-6) + + def test_dropped_remainder_shortens_the_result(self, records): + loader = make_loader(records, batch_size=5, drop_remainder=True) + preds, _ = get_preds_targets(ConstantOffset(), loader) + assert preds.shape == (10,) + + +class TestComputeMetrics: + def test_perfect_predictions_score_perfectly(self, loader): + metrics = compute_metrics(ConstantOffset(), loader) + assert isinstance(metrics, Metrics) + assert metrics.mse == pytest.approx(0.0, abs=1e-8) + assert metrics.rmse == pytest.approx(0.0, abs=1e-8) + assert metrics.mae == pytest.approx(0.0, abs=1e-8) + assert metrics.r2 == pytest.approx(1.0) + + def test_constant_offset_shows_up_as_bias(self, loader): + metrics = compute_metrics(ConstantOffset(offset=3.0), loader) + assert metrics.bias == pytest.approx(3.0, rel=1e-4) + assert metrics.mae == pytest.approx(3.0, rel=1e-4) + assert metrics.rmse == pytest.approx(3.0, rel=1e-4) + + def test_worse_predictions_lower_r2(self, records): + better = compute_metrics(ConstantOffset(offset=1.0), make_loader(records, batch_size=4)) + worse = compute_metrics(ConstantOffset(offset=50.0), make_loader(records, batch_size=4)) + assert worse.r2 < better.r2 + assert worse.mse > better.mse diff --git a/tests/v2/models/__init__.py b/tests/v2/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/v2/models/test_rqs.py b/tests/v2/models/test_rqs.py new file mode 100644 index 0000000..07d4a7a --- /dev/null +++ b/tests/v2/models/test_rqs.py @@ -0,0 +1,254 @@ +""" +Tests for the v2 conditional rational-quadratic spline (RQS) flow model. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax import nnx +from omegaconf import OmegaConf + +from estimint.v2.data.features import StandardScaler +from estimint.v2.models.rqs import ConditionalRQS, RQSArtifact, _rqs, rqs_loss + +BOUNDS = 6 +N_BINS = 8 +N_CONTEXT = 4 + + +@pytest.fixture +def raw_params(): + """Unconstrained (widths, heights, derivatives) for a batch of splines.""" + key = jax.random.key(0) + k1, k2, k3 = jax.random.split(key, 3) + n = 16 + return ( + jax.random.normal(k1, (n, N_BINS)), + jax.random.normal(k2, (n, N_BINS)), + jax.random.normal(k3, (n, N_BINS + 1)), + ) + + +@pytest.fixture +def model(): + return ConditionalRQS(N_CONTEXT, width=16, depth=2, n_bins=N_BINS, bounds=BOUNDS, rngs=nnx.Rngs(0)) + + +@pytest.fixture +def context(): + return jax.random.normal(jax.random.key(1), (16, N_CONTEXT)) + + +class TestSpline: + """The bare spline transform, independent of any network.""" + + def test_forward_inverse_round_trip(self, raw_params): + x = jnp.linspace(-BOUNDS + 0.1, BOUNDS - 0.1, 16) + z, _ = _rqs(x, *raw_params, BOUNDS, inverse=False) + x_back, _ = _rqs(z, *raw_params, BOUNDS, inverse=True) + np.testing.assert_allclose(x_back, x, atol=1e-4) + + def test_identity_outside_bounds(self, raw_params): + x = jnp.full((16,), BOUNDS + 2.0) + z, log_det = _rqs(x, *raw_params, BOUNDS, inverse=False) + np.testing.assert_allclose(z, x) + np.testing.assert_allclose(log_det, 0.0) + + def test_forward_is_monotonic(self, raw_params): + widths, heights, derivatives = raw_params + # Evaluate one spline (row 0) on an increasing grid of inputs. + grid = jnp.linspace(-BOUNDS + 0.1, BOUNDS - 0.1, 64) + row = lambda p: jnp.repeat(p[:1], grid.shape[0], axis=0) + z, _ = _rqs(grid, row(widths), row(heights), row(derivatives), BOUNDS, inverse=False) + assert jnp.all(jnp.diff(z) > 0) + + def test_stays_within_bounds(self, raw_params): + x = jnp.linspace(-BOUNDS + 0.1, BOUNDS - 0.1, 16) + z, _ = _rqs(x, *raw_params, BOUNDS, inverse=False) + assert jnp.all(jnp.abs(z) <= BOUNDS) + + def test_log_det_matches_numerical_derivative(self, raw_params): + x = jnp.linspace(-BOUNDS + 0.5, BOUNDS - 0.5, 16) + eps = 1e-3 + z_hi, _ = _rqs(x + eps, *raw_params, BOUNDS, inverse=False) + z_lo, _ = _rqs(x - eps, *raw_params, BOUNDS, inverse=False) + _, log_det = _rqs(x, *raw_params, BOUNDS, inverse=False) + np.testing.assert_allclose(jnp.exp(log_det), (z_hi - z_lo) / (2 * eps), rtol=1e-2) + + +class TestConditionalRQS: + def test_log_prob_shape_and_finite(self, model, context): + y0 = jax.random.normal(jax.random.key(2), (context.shape[0],)) + log_prob = model.log_prob(y0, context) + assert log_prob.shape == (context.shape[0],) + assert jnp.all(jnp.isfinite(log_prob)) + + def test_density_integrates_to_one(self, model, context): + """The flow is a normalized density in the standardized target space.""" + grid = jnp.linspace(-12, 12, 4001) + one_row = jnp.repeat(context[:1], grid.shape[0], axis=0) + density = jnp.exp(model.log_prob(grid, one_row)) + assert jnp.trapezoid(density, grid) == pytest.approx(1.0, abs=1e-3) + + def test_quantiles_increase_with_probability(self, model, context): + probs = jnp.array([0.05, 0.25, 0.5, 0.75, 0.95]) + y0 = model.quantiles(context, probs) # (Q, B) + assert y0.shape == (probs.shape[0], context.shape[0]) + assert jnp.all(jnp.diff(y0, axis=0) > 0) + + def test_quantiles_matches_single_quantile(self, model, context): + probs = jnp.array([0.1, 0.5, 0.9]) + batched = model.quantiles(context, probs) + for i, q in enumerate(probs): + np.testing.assert_allclose(batched[i], model.quantile(context, float(q)), atol=1e-5) + + def test_median_maps_back_to_base_zero(self, model, context): + """quantile(0.5) is the target value the flow maps to z = 0.""" + y0 = model.quantile(context, 0.5) + widths, heights, derivatives = model._params(context) + z, _ = _rqs(y0, widths, heights, derivatives, model.bounds, inverse=False) + np.testing.assert_allclose(z, 0.0, atol=1e-4) + + def test_from_cfg(self): + cfg = OmegaConf.create( + { + "seed": 0, + "width": 16, + "depth": 2, + "n_bins": N_BINS, + "rqs_bounds": BOUNDS, + "mlp_residual": False, + "dropout_rate": 0.0, + } + ) + model = ConditionalRQS.from_cfg(cfg, n_context=N_CONTEXT) + assert model.K == N_BINS + assert model.bounds == BOUNDS + # net emits K widths + K heights + (K + 1) derivatives + assert model.net(jnp.zeros((2, N_CONTEXT))).shape == (2, 3 * N_BINS + 1) + + def test_residual_variant_runs(self, context): + model = ConditionalRQS( + N_CONTEXT, width=16, depth=2, n_bins=N_BINS, bounds=BOUNDS, residual=True, rngs=nnx.Rngs(0) + ) + assert jnp.all(jnp.isfinite(model.log_prob(jnp.zeros(context.shape[0]), context))) + + +class TestLoss: + def test_loss_is_finite(self, model, context): + y0 = jax.random.normal(jax.random.key(3), (context.shape[0],)) + w = jnp.ones_like(y0) + total_loss, normalization = rqs_loss(model, context, y0, w) + assert jnp.isfinite(total_loss / normalization) + + def test_loss_is_gradable(self, model, context): + y0 = jax.random.normal(jax.random.key(3), (context.shape[0],)) + w = jnp.ones_like(y0) + def objective(model, X, y0, w): + total_loss, normalization = rqs_loss(model, X, y0, w) + return total_loss / normalization + + grads = nnx.grad(objective)(model, context, y0, w) + leaves = jax.tree_util.tree_leaves(grads) + assert leaves and all(jnp.all(jnp.isfinite(g)) for g in leaves) + + def test_zero_weight_rows_are_ignored(self, model, context): + y0 = jax.random.normal(jax.random.key(3), (context.shape[0],)) + w = jnp.ones_like(y0).at[8:].set(0.0) + masked_total, masked_normalization = rqs_loss(model, context, y0, w) + kept_total, kept_normalization = rqs_loss(model, context[:8], y0[:8], jnp.ones(8)) + assert masked_total / masked_normalization == pytest.approx(float(kept_total / kept_normalization), rel=1e-5) + + +def make_scaler(mean, scale): + scaler = StandardScaler() + scaler.mean_ = np.array(mean, dtype=np.float32) + scaler.scale_ = np.array(scale, dtype=np.float32) + return scaler + + +@pytest.fixture +def features(): + return ["prev_y9", "dn0_use", "Q0", "phi_bednets"] + + +@pytest.fixture +def artifact(model, features): + model.eval() + return RQSArtifact( + model=model, + feature_scaler=make_scaler(np.zeros(len(features)), np.ones(len(features))), + target_scaler=make_scaler([0.0], [1.0]), + features=features, + ) + + +class TestRQSArtifact: + def test_rejects_scaler_with_wrong_feature_count(self, model, features): + with pytest.raises(ValueError, match="features"): + RQSArtifact( + model=model, + feature_scaler=make_scaler(np.zeros(2), np.ones(2)), + target_scaler=make_scaler([0.0], [1.0]), + features=features, + ) + + def test_predict_shape_and_non_negative(self, artifact, features): + X = np.random.default_rng(0).normal(size=(5, len(features))).astype(np.float32) + preds = artifact.predict(X) + assert preds.shape == (5,) + assert np.all(preds >= 0) + + def test_predict_accepts_single_row(self, artifact, features): + X = np.zeros(len(features), dtype=np.float32) + assert artifact.predict(X).shape == (1,) + + def test_dict_input_matches_array_input(self, artifact, features): + values = [0.4, 0.1, 0.9, 0.5] + row = dict(zip(features, values)) + np.testing.assert_allclose(artifact.predict(row), artifact.predict(np.array(values, dtype=np.float32))) + + def test_dict_input_is_order_independent(self, artifact, features): + row = {f: v for f, v in zip(features, [0.4, 0.1, 0.9, 0.5])} + shuffled = dict(reversed(list(row.items()))) + np.testing.assert_allclose(artifact.predict(shuffled), artifact.predict(row)) + + def test_list_of_dicts_is_batched(self, artifact, features): + rows = [dict.fromkeys(features, 0.1), dict.fromkeys(features, 0.2)] + assert artifact.predict(rows).shape == (2,) + + @pytest.mark.parametrize( + "bad_row", + [ + {"prev_y9": 0.1}, # missing features + {"prev_y9": 0.1, "dn0_use": 0.1, "Q0": 0.1, "phi_bednets": 0.1, "nope": 0.1}, # unexpected feature + ], + ) + def test_rejects_malformed_dict_rows(self, artifact, bad_row): + with pytest.raises(KeyError): + artifact.predict(bad_row) + + def test_rejects_empty_input(self, artifact): + with pytest.raises(ValueError, match="empty"): + artifact.predict([]) + + def test_quantiles_are_ordered(self, artifact, features): + X = np.zeros((3, len(features)), dtype=np.float32) + low, mid, high = (artifact.quantile(X, q) for q in (0.1, 0.5, 0.9)) + assert np.all(low <= mid) and np.all(mid <= high) + + def test_interval_brackets_the_prediction(self, artifact, features): + X = np.zeros((3, len(features)), dtype=np.float32) + lower, upper = artifact.interval(X, alpha=0.10) + preds = artifact.predict(X) + assert np.all(lower >= 0) + assert np.all(lower <= preds) and np.all(preds <= upper) + + def test_conformal_offset_widens_the_interval(self, artifact, features): + X = np.zeros((3, len(features)), dtype=np.float32) + lower, upper = artifact.interval(X, alpha=0.10) + artifact.conformal[0.10] = 1.0 + wide_lower, wide_upper = artifact.interval(X, alpha=0.10) + assert np.all(wide_upper > upper) + assert np.all(wide_lower <= lower) diff --git a/tests/v2/training/__init__.py b/tests/v2/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/v2/training/test_calibrate.py b/tests/v2/training/test_calibrate.py new file mode 100644 index 0000000..1310ebb --- /dev/null +++ b/tests/v2/training/test_calibrate.py @@ -0,0 +1,51 @@ +""" +Tests for the split-conformal interval correction. +""" + +import numpy as np +import pytest + +from estimint.v2.training.calibrate import conformal_offset + + +@pytest.fixture +def y(): + return np.random.default_rng(0).normal(size=500) + + +def test_offset_shrinks_intervals_that_already_over_cover(): + y = np.linspace(0, 1, 100) + offset = conformal_offset(y - 1.0, y + 1.0, y, alpha=0.10) + # every score is negative here, so the correction narrows rather than widens + assert offset == pytest.approx(-1.0) + + +def test_offset_restores_target_coverage(y): + lower, upper = np.zeros_like(y), np.zeros_like(y) # degenerate point intervals + offset = conformal_offset(lower, upper, y, alpha=0.10) + coverage = np.mean((y >= lower - offset) & (y <= upper + offset)) + assert coverage >= 0.90 + + +def test_offset_is_the_score_quantile(y): + lower, upper = -np.ones_like(y), np.ones_like(y) + scores = np.maximum(lower - y, y - upper) + k = int(np.ceil((len(y) + 1) * 0.90)) + assert conformal_offset(lower, upper, y, alpha=0.10) == pytest.approx(np.sort(scores)[k - 1]) + + +def test_wider_intervals_need_a_smaller_offset(y): + tight = conformal_offset(-0.1 * np.ones_like(y), 0.1 * np.ones_like(y), y) + loose = conformal_offset(-1.0 * np.ones_like(y), 1.0 * np.ones_like(y), y) + assert loose < tight + + +def test_smaller_alpha_gives_a_larger_offset(y): + lower, upper = np.zeros_like(y), np.zeros_like(y) + assert conformal_offset(lower, upper, y, alpha=0.01) > conformal_offset(lower, upper, y, alpha=0.20) + +@pytest.mark.parametrize("alpha", [0.0, 1.0, -0.1, 1.1]) +def test_invalid_alpha_raises(y, alpha): + lower, upper = np.zeros_like(y), np.zeros_like(y) + with pytest.raises(ValueError): + conformal_offset(lower, upper, y, alpha=alpha) \ No newline at end of file diff --git a/tests/v2/training/test_checkpoint.py b/tests/v2/training/test_checkpoint.py new file mode 100644 index 0000000..7189626 --- /dev/null +++ b/tests/v2/training/test_checkpoint.py @@ -0,0 +1,74 @@ +""" +Tests for Orbax checkpoint saving and restoring. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax import nnx + +from estimint.v2.models.mlp import MLP +from estimint.v2.training.checkpoint import _resolve_checkpoint_dir, restore_model, save_checkpoint + +MODEL_NAME = "RQS" + + +def make_model(seed): + return MLP(3, 1, width=4, depth=1, dropout_rate=0.0, rngs=nnx.Rngs(seed)) + + +def weights(model): + return [np.asarray(w) for w in jax.tree_util.tree_leaves(nnx.state(model, nnx.Param))] + + +@pytest.fixture +def ckpt_dir(tmp_path): + return str(tmp_path / "ckpts") + + +def test_resolve_checkpoint_dir_appends_the_model_name(tmp_path): + assert _resolve_checkpoint_dir(str(tmp_path), MODEL_NAME).name == MODEL_NAME + + +def test_save_creates_the_checkpoint_directory(ckpt_dir): + save_checkpoint(ckpt_dir, MODEL_NAME, make_model(0)) + assert _resolve_checkpoint_dir(ckpt_dir, MODEL_NAME).exists() + + +def test_restore_recovers_the_saved_weights(ckpt_dir): + saved = make_model(0) + save_checkpoint(ckpt_dir, MODEL_NAME, saved) + + restored = restore_model(ckpt_dir, MODEL_NAME, make_model(1)) # different init + for got, want in zip(weights(restored), weights(saved)): + np.testing.assert_allclose(got, want) + + +def test_restored_model_reproduces_predictions(ckpt_dir): + saved = make_model(0) + save_checkpoint(ckpt_dir, MODEL_NAME, saved) + x = jnp.ones((2, 3)) + + restored = restore_model(ckpt_dir, MODEL_NAME, make_model(1)) + np.testing.assert_allclose(restored(x), saved(x), rtol=1e-6) + + +def test_saving_twice_keeps_the_latest_weights(ckpt_dir): + save_checkpoint(ckpt_dir, MODEL_NAME, make_model(0)) + latest = make_model(2) + save_checkpoint(ckpt_dir, MODEL_NAME, latest) + + restored = restore_model(ckpt_dir, MODEL_NAME, make_model(1)) + for got, want in zip(weights(restored), weights(latest)): + np.testing.assert_allclose(got, want) + + +def test_models_are_namespaced_by_name(ckpt_dir): + first, second = make_model(0), make_model(2) + save_checkpoint(ckpt_dir, "first", first) + save_checkpoint(ckpt_dir, "second", second) + + restored = restore_model(ckpt_dir, "first", make_model(1)) + for got, want in zip(weights(restored), weights(first)): + np.testing.assert_allclose(got, want) diff --git a/tests/v2/training/test_train_step.py b/tests/v2/training/test_train_step.py new file mode 100644 index 0000000..a5c1e31 --- /dev/null +++ b/tests/v2/training/test_train_step.py @@ -0,0 +1,224 @@ +""" +Tests for the optimizer, train/eval steps, and the training loop. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax import nnx +from omegaconf import OmegaConf + +from estimint.v2.data.features import StandardScaler +from estimint.v2.data.preprocess import PreparedData +from estimint.v2.models.mlp import MLP +from estimint.v2.training.checkpoint import _resolve_checkpoint_dir +from estimint.v2.training.train_step import ( + _weighted_mean, + create_optimizer, + get_total_params, + make_eval_step, + make_train_step, + train_model, +) + +N_FEATURES = 3 + + +def some_unweighted_loss(model, X, y0, w): + """Unweighted loss that ignores the weights, for testing eval_step.""" + per_record = -model.log_prob(y0, X) + return jnp.sum(per_record), jnp.asarray(per_record.shape[0], dtype=per_record.dtype) + +def mse_loss(model, x, y, w): + """Weighted MSE against a scalar target — stands in for rqs_loss.""" + preds = model(x)[:, 0] + return jnp.sum(w * (preds - y) ** 2), jnp.sum(w) + + +def scalar(terms): + """Collapse a (total, mass) pair to the batch loss.""" + total, mass = terms + return float(total / mass) + +def make_model(seed=0, width=8, depth=1): + return MLP(N_FEATURES, 1, width=width, depth=depth, dropout_rate=0.0, rngs=nnx.Rngs(seed)) + + +def make_records(n=32, seed=0): + """Records shaped like preprocess output: y_std is a linear function of x.""" + rng = np.random.default_rng(seed) + X = rng.normal(size=(n, N_FEATURES)).astype(np.float32) + y = X.sum(axis=1).astype(np.float32) + return [{"x": X[i], "y": y[i], "y_std": y[i], "w": np.float32(1.0)} for i in range(n)] + + +def weights(model): + return [np.asarray(w) for w in jax.tree_util.tree_leaves(nnx.state(model, nnx.Param))] + + +@pytest.fixture +def batch(): + records = make_records(16) + return ( + jnp.stack([r["x"] for r in records]), + jnp.array([r["y_std"] for r in records]), + jnp.ones(len(records)), + ) + + +class TestGetTotalParams: + def test_counts_every_weight_and_bias(self): + # inp: 3*8 + 8, hidden: 8*8 + 8, out: 8*1 + 1 + assert get_total_params(make_model(width=8, depth=1)) == 32 + 72 + 9 + + def test_deeper_models_have_more_parameters(self): + assert get_total_params(make_model(depth=3)) > get_total_params(make_model(depth=1)) + + +class TestCreateOptimizer: + def apply_updates(self, model, optimizer, batch, n): + for _ in range(n): + _, grads = nnx.value_and_grad(mse_loss, has_aux=True)(model, *batch) + optimizer.update(model, grads) + + def test_first_step_is_a_no_op_because_warmup_starts_at_zero(self, batch): + model = make_model() + optimizer = create_optimizer(model, learning_rate=1e-2, total_steps=100) + before = weights(model) + + self.apply_updates(model, optimizer, batch, 1) + + for a, b in zip(before, weights(model)): + np.testing.assert_allclose(a, b) + + def test_optimizer_updates_the_model_once_warmup_ramps_up(self, batch): + model = make_model() + optimizer = create_optimizer(model, learning_rate=1e-2, total_steps=100) + before = weights(model) + + self.apply_updates(model, optimizer, batch, 5) # warmup is 3% of total_steps + + assert any(not np.allclose(a, b) for a, b in zip(before, weights(model))) + + +class TestTrainStep: + def test_training_reduces_the_loss(self, batch): + model = make_model() + optimizer = create_optimizer(model, learning_rate=1e-2, total_steps=100) + train_step = make_train_step(mse_loss) + + first = scalar(train_step(model, optimizer, *batch)) + for _ in range(50): + last = scalar(train_step(model, optimizer, *batch)) + assert last < first + + def test_train_step_returns_the_pre_update_loss(self, batch): + model = make_model() + optimizer = create_optimizer(model, learning_rate=1e-2, total_steps=100) + eval_step = make_eval_step(mse_loss) + + expected = scalar(eval_step(model, *batch)) + assert scalar(make_train_step(mse_loss)(model, optimizer, *batch)) == pytest.approx(expected, rel=1e-5) + + def test_eval_step_leaves_the_model_unchanged(self, batch): + model = make_model() + before = weights(model) + make_eval_step(mse_loss)(model, *batch) + for a, b in zip(before, weights(model)): + np.testing.assert_array_equal(a, b) + + def test_aggregation_is_invariant_to_batch_partition(self): + """The whole point of (total, mass): unequal batches combine exactly.""" + model = make_model() + eval_step = make_eval_step(mse_loss) + records = make_records(10) + x = jnp.stack([r["x"] for r in records]) + y = jnp.array([r["y_std"] for r in records]) + w = jnp.arange(1.0, 11.0) # deliberately non-uniform + + whole = scalar(eval_step(model, x, y, w)) + split = _weighted_mean( + [eval_step(model, x[:7], y[:7], w[:7]), eval_step(model, x[7:], y[7:], w[7:])], + ) + assert split == pytest.approx(whole, rel=1e-5) + + + def test_validation_smaller_than_one_batch_still_trains(self, cfg, prepared_data): + cfg.batch_size = 64 # larger than the 32-record val split + train_model(make_model(), cfg, prepared_data, mse_loss, name="RQS") + + +@pytest.fixture +def cfg(tmp_path): + return OmegaConf.create( + { + "batch_size": 8, + "seed": 0, + "num_workers": 0, + "num_epochs": 8, + "min_epochs": 0, + "patience": 3, + "lr": 1e-2, + "weight_decay": 1e-4, + "use_wandb": False, + "checkpoint_dir": str(tmp_path / "ckpts"), + } + ) + + +@pytest.fixture +def prepared_data(): + scaler = StandardScaler().fit(np.zeros((2, N_FEATURES)) + np.arange(N_FEATURES)) + return PreparedData( + train_data=make_records(64, seed=0), + val_data=make_records(32, seed=1), + test_data=[], + input_size=N_FEATURES, + feature_scaler=scaler, + target_scaler=scaler, + ) + + +class TestTrainModel: + def test_training_improves_validation_loss(self, cfg, prepared_data): + model = make_model() + eval_step = make_eval_step(mse_loss) + val = ( + jnp.stack([r["x"] for r in prepared_data.val_data]), + jnp.array([r["y_std"] for r in prepared_data.val_data]), + jnp.ones(len(prepared_data.val_data)), + ) + before = scalar(eval_step(model, *val)) + + trained = train_model(model, cfg, prepared_data, mse_loss, name="RQS", use_standardized_y=True) + assert scalar(eval_step(trained, *val)) < before + + def test_training_writes_a_checkpoint(self, cfg, prepared_data): + train_model(make_model(), cfg, prepared_data, mse_loss, name="RQS") + assert _resolve_checkpoint_dir(cfg.checkpoint_dir, "RQS").exists() + + def test_returns_the_same_model_instance_updated_in_place(self, cfg, prepared_data): + model = make_model() + assert train_model(model, cfg, prepared_data, mse_loss, name="RQS") is model + + def test_early_stopping_ends_training_when_validation_stalls(self, cfg, prepared_data, caplog): + cfg.num_epochs = 50 + cfg.patience = 2 + constant_loss = lambda model, x, y, w: (jnp.sum(jnp.zeros_like(y)) + 1.0, jnp.asarray(1.0)) + + with caplog.at_level("INFO"): + train_model(make_model(), cfg, prepared_data, constant_loss, name="RQS") + + assert "Early stopping" in caplog.text + + def test_min_epochs_defers_early_stopping(self, cfg, prepared_data, caplog): + cfg.num_epochs = 4 + cfg.min_epochs = 4 # never eligible to stop or checkpoint a best model + cfg.patience = 1 + constant_loss = lambda model, x, y, w: (jnp.sum(jnp.zeros_like(y)) + 1.0, jnp.asarray(1.0)) + + with caplog.at_level("INFO"): + train_model(make_model(), cfg, prepared_data, constant_loss, name="RQS") + + assert "Early stopping" not in caplog.text diff --git a/uv.lock b/uv.lock index 657798c..8840f05 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -40,6 +40,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, ] +[[package]] +name = "annotated-types" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, +] + [[package]] name = "antlr4-python3-runtime" version = "4.9.3" @@ -68,6 +77,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3b/00/2344469e2084fb287c2e0b57b72910309874c3245463acd6cf5e3db69324/appdirs-1.4.4-py2.py3-none-any.whl", hash = "sha256:a841dacd6b99318a741b166adb07e19ee71a274450e68237b4650ca1055ab128", size = 9566, upload-time = "2020-05-11T07:59:49.499Z" }, ] +[[package]] +name = "array-record" +version = "0.8.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "sys_platform != 'win32'" }, + { name = "etils", extra = ["epath"], marker = "sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/ff/8238c28f9bfda365ba991fc2b8a566c71847fc5a2d0913158aeadafd4e63/array_record-0.8.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:dada305fa0dfa3fd6f5f263c43ed37546f815e4f33ce30b066175384dead752e", size = 3989102, upload-time = "2025-11-13T16:15:33.845Z" }, + { url = "https://files.pythonhosted.org/packages/fb/67/4b4cd9891a36aec236e35688b47c44648bb645946e98653a7ad54f890128/array_record-0.8.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:718403dc9a364519a5fc440ed9e2784077e965489d5a44c96970d3101101e1cd", size = 4838611, upload-time = "2025-11-13T16:15:35.016Z" }, + { url = "https://files.pythonhosted.org/packages/50/65/2f60b9ca59e5fce140be36f4136ebcc83385c9e56cd499e2942f9fb4b250/array_record-0.8.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:458f6658de86c9369b23ffe64dcb31393a919b91ae2f15147ee2beb2010b122a", size = 4995925, upload-time = "2025-11-13T16:15:37.091Z" }, + { url = "https://files.pythonhosted.org/packages/d6/63/fd734c5f2be8017d4f33cff0dd3d3e2f91de3ab988f93755cdc54d67b654/array_record-0.8.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:fa94fa053c1afaecc183a1c31463fd89d1b7b148cc526095bfc50ab58967e47c", size = 3989122, upload-time = "2025-11-13T16:15:38.375Z" }, + { url = "https://files.pythonhosted.org/packages/e9/b0/9f649781c8e630cbb8724fb56627430b4249b51e8536e6e0a31d6f85f147/array_record-0.8.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adfe6c92918363747539c0caa579f568f75aac079ad7bbd4ebcf7fdb02e5461b", size = 4838733, upload-time = "2025-11-13T16:15:39.767Z" }, + { url = "https://files.pythonhosted.org/packages/d7/6d/5575a1d64cffecb72de5743b35a1532b1801a0350b9b45f277090d02b3c2/array_record-0.8.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:13df1aec9b38afe98973bf5fe3cf523f83ca904b0423d85319930877145b8f28", size = 4995837, upload-time = "2025-11-13T16:15:40.973Z" }, + { url = "https://files.pythonhosted.org/packages/99/46/1449f2c36f75e563c5bdd59133146a10110bcf3a3cd53d7d74a535579603/array_record-0.8.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:14facb713fc55ea612cf853c0abbc602032ae619eb85037af3aa41dc4e27eed6", size = 3989205, upload-time = "2025-11-13T16:15:42.156Z" }, + { url = "https://files.pythonhosted.org/packages/35/ef/f85d329ac6cdcc8c3e8af832c30a9a0b719c6719a09191056ef7342a3627/array_record-0.8.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:06b17a0a90bc74a80c4ebad0b99f0e039adb01746b673db3d7a5632d8b138ddd", size = 4839320, upload-time = "2025-11-13T16:15:43.617Z" }, + { url = "https://files.pythonhosted.org/packages/f2/89/510594ff506d456d7dc0e8dfa9095f078e5f98da49d1b8f7be1fd9c0717e/array_record-0.8.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:89d2cf5f4709be3c6c2b30c0d366005223375633ffce66dcb13d927a0bdc5228", size = 4995667, upload-time = "2025-11-13T16:15:45.092Z" }, +] + [[package]] name = "black" version = "25.11.0" @@ -175,6 +204,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -361,36 +399,49 @@ wheels = [ [[package]] name = "estimint" -version = "1.5.4" +version = "2.1.0" source = { editable = "." } dependencies = [ + { name = "flax" }, + { name = "huggingface-hub" }, + { name = "jax" }, + { name = "jaxtyping" }, { name = "numpy" }, + { name = "omegaconf" }, + { name = "orbax-checkpoint" }, { name = "pandas" }, - { name = "scipy" }, - { name = "xgboost" }, + { name = "wandb" }, ] [package.optional-dependencies] all = [ { name = "appdirs" }, { name = "duckdb" }, + { name = "grain" }, + { name = "hydra-core" }, + { name = "jax", extra = ["cuda12"] }, { name = "matplotlib" }, { name = "mintstate" }, - { name = "pyarrow" }, + { name = "optax" }, { name = "requests" }, - { name = "scikit-learn" }, + { name = "tqdm" }, ] download = [ { name = "appdirs" }, { name = "requests" }, ] +gpu = [ + { name = "jax", extra = ["cuda12"] }, +] scenarios = [ { name = "mintstate" }, ] train = [ { name = "duckdb" }, - { name = "pyarrow" }, - { name = "scikit-learn" }, + { name = "grain" }, + { name = "hydra-core" }, + { name = "optax" }, + { name = "tqdm" }, ] viz = [ { name = "matplotlib" }, @@ -410,18 +461,26 @@ dev = [ requires-dist = [ { name = "appdirs", marker = "extra == 'download'", specifier = ">=1.4.0" }, { name = "duckdb", marker = "extra == 'train'", specifier = ">=0.8.0" }, - { name = "estimint", extras = ["train", "viz", "download", "scenarios"], marker = "extra == 'all'" }, + { name = "estimint", extras = ["train", "viz", "download", "scenarios", "gpu"], marker = "extra == 'all'" }, + { name = "flax", specifier = ">=0.12.7" }, + { name = "grain", marker = "extra == 'train'", specifier = ">=0.2.16" }, + { name = "huggingface-hub", specifier = ">=0.24.0" }, + { name = "hydra-core", marker = "extra == 'train'", specifier = ">=1.3.2" }, + { name = "jax", specifier = ">=0.10.1" }, + { name = "jax", extras = ["cuda12"], marker = "extra == 'gpu'", specifier = ">=0.10.1" }, + { name = "jaxtyping", specifier = ">=0.3.10" }, { name = "matplotlib", marker = "extra == 'viz'", specifier = ">=3.4.0" }, { name = "mintstate", marker = "extra == 'scenarios'", specifier = ">=0.3.0" }, { name = "numpy", specifier = ">=1.20.0" }, + { name = "omegaconf", specifier = ">=2.3" }, + { name = "optax", marker = "extra == 'train'", specifier = ">=0.2.8" }, + { name = "orbax-checkpoint", specifier = ">=0.12.0" }, { name = "pandas", specifier = ">=1.3.0" }, - { name = "pyarrow", marker = "extra == 'train'", specifier = ">=10.0.0" }, { name = "requests", marker = "extra == 'download'", specifier = ">=2.28.0" }, - { name = "scikit-learn", marker = "extra == 'train'", specifier = ">=1.0.0" }, - { name = "scipy", specifier = ">=1.7.0" }, - { name = "xgboost", specifier = ">=1.6.0" }, + { name = "tqdm", marker = "extra == 'train'", specifier = ">=4.67.3" }, + { name = "wandb", specifier = ">=0.28.0" }, ] -provides-extras = ["train", "viz", "download", "scenarios", "all"] +provides-extras = ["train", "gpu", "viz", "download", "scenarios", "all"] [package.metadata.requires-dev] dev = [ @@ -546,6 +605,58 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/22/4222d7ddf3da30f363edaa98e329c2bce6c65497c9cb2810931c8b2c0fbc/fsspec-2026.6.0-py3-none-any.whl", hash = "sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1", size = 203949, upload-time = "2026-06-16T01:57:26.358Z" }, ] +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.1.50" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/33/f6/354ae6491228b5eb40e10d89c4d13c651fe1cf7556e35ebdded50cff57ce/gitpython-3.1.50.tar.gz", hash = "sha256:80da2d12504d52e1f998772dc5baf6e553f8d2fcfe1fcc226c9d9a2ee3372dcc", size = 219798, upload-time = "2026-05-06T04:01:26.571Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/7a/1c6e3562dfd8950adbb11ffbc65d21e7c89d01a6e4f137fa981056de25c5/gitpython-3.1.50-py3-none-any.whl", hash = "sha256:d352abe2908d07355014abdd21ddf798c2a961469239afec4962e9da884858f9", size = 212507, upload-time = "2026-05-06T04:01:23.799Z" }, +] + +[[package]] +name = "grain" +version = "0.2.18" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "array-record", marker = "sys_platform != 'win32'" }, + { name = "cloudpickle" }, + { name = "etils", extra = ["epath", "epy"] }, + { name = "numpy" }, + { name = "portpicker" }, + { name = "protobuf" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/20/59fb32fd6aadc40e0dbde586a62903646ad778413c78280ea8d1e18e73de/grain-0.2.18-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1bcff4c0ec71cee0b8c26076e3b60d83897cd7586b8f182291969743fd79a7b9", size = 554329, upload-time = "2026-06-17T20:57:47.422Z" }, + { url = "https://files.pythonhosted.org/packages/c5/60/13bf95936145419a4673f33262d239477e1dc7da30d1a576945e2cba84af/grain-0.2.18-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1eb8f8424d3844944d1c37248b0430a5598b21f7ec5141119e0a70069edbd943", size = 610874, upload-time = "2026-06-17T20:57:48.907Z" }, + { url = "https://files.pythonhosted.org/packages/56/86/237d99cc246f87536ca3f3764fe94e2e364dababfe382ad227f3153079de/grain-0.2.18-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1ecd87427b4050b695f9b9b5a41d8058208586d3f90a2e56f14a64f920aa47f8", size = 610169, upload-time = "2026-06-17T20:57:50.073Z" }, + { url = "https://files.pythonhosted.org/packages/ee/4b/7ec021e2385094ab31f6e6dad3e6f5bb2251dff5c3543934c347ba7ada3e/grain-0.2.18-cp312-cp312-win_amd64.whl", hash = "sha256:c8a29696d98b2e84d8531df211006b2187475af7cebc45408023da6aa3ac1494", size = 539810, upload-time = "2026-06-17T20:57:51.414Z" }, + { url = "https://files.pythonhosted.org/packages/a2/3d/6ad2033ae65145bcef635741c0d351bfe18a6a241b803766490a432551ff/grain-0.2.18-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bec1c551574d3a1885ea0d1481586b4c2490a1332120bf626b459e2f9a6f1d06", size = 554349, upload-time = "2026-06-17T20:57:52.525Z" }, + { url = "https://files.pythonhosted.org/packages/60/d5/590925b5fdf33218dde9f72322b88b65f4043c6a6a06cd00520465ac6915/grain-0.2.18-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:229683c7714c3275cc4fee6cbc7356d47769eeaa286fbdc6af7a9542dcb468e5", size = 610930, upload-time = "2026-06-17T20:57:53.681Z" }, + { url = "https://files.pythonhosted.org/packages/9b/f3/0d3269dd6796c187f808708b381f6d85a62e5022f45f00c938a7f3a605f2/grain-0.2.18-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:61290a4af38fbddac41d6da1d19129df498dece53c3de4cc02e7ec557a0b0024", size = 610204, upload-time = "2026-06-17T20:57:55.104Z" }, + { url = "https://files.pythonhosted.org/packages/d4/ed/d491cd9e6406d4eca43443939c6b303f8cd66386f861988813ed92038e01/grain-0.2.18-cp313-cp313-win_amd64.whl", hash = "sha256:8d74e3f33403146fb9da0b101d7d18148e3979f5adc8d625237524cc144c0d24", size = 539915, upload-time = "2026-06-17T20:57:56.573Z" }, + { url = "https://files.pythonhosted.org/packages/3c/de/14532d47c14bcfaeb79878621ea4f28863df60ec1ce29f43a9781a1a2167/grain-0.2.18-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:5986e756a0803012a74af4d34a506ed2877f6a86b8d1164227e8a67e09af1688", size = 554398, upload-time = "2026-06-17T20:57:57.641Z" }, + { url = "https://files.pythonhosted.org/packages/0f/31/07c861e8b0a782d2b17b3e8a7b9191420d30d4a9ce86a9882502d4d883eb/grain-0.2.18-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:09bef4e8290d06e79cfddc71e777d0b5f60316e5a43af61f8a3d0e6a2d9feccc", size = 610928, upload-time = "2026-06-17T20:57:58.901Z" }, + { url = "https://files.pythonhosted.org/packages/b7/60/b257f1173a6c71c67a1de81798427a7e73ceb71cd6573ccccf7f528018f5/grain-0.2.18-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fc7204490aff21620fbdbd69e088991a4ae1ebc31586456191c11f9f61a8c0e6", size = 610179, upload-time = "2026-06-17T20:58:00.154Z" }, + { url = "https://files.pythonhosted.org/packages/a2/c8/d3f559b3691559054d57779ce9e9dbb832c7631b622e94c8c092ce26ca18/grain-0.2.18-cp314-cp314-win_amd64.whl", hash = "sha256:fe83461d55e81c9c9cdba3d9d9ed0185db7b41cddb8e28223aa495ca3516877c", size = 547828, upload-time = "2026-06-17T20:58:01.419Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -645,6 +756,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/10/dd/220f0e91743136725352497e98540772a01fc7c3ab96ff16c3c74424e984/hydra_core-1.3.4.tar.gz", hash = "sha256:ad0f7b05a0242255a8984d5a4ed2f6847f7b783ed727368a2c0155ec52d6c34c", size = 3263348, upload-time = "2026-07-04T16:25:38.891Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/cd/a568610bafe991fdd3f628fb606316b3b2be52ded019284e895d9beb3a1e/hydra_core-1.3.4-py3-none-any.whl", hash = "sha256:e58683692904a09f1fdfffa1a9b86bfd94e215b59f1ee17e7cd7d92738090d33", size = 155478, upload-time = "2026-07-04T16:25:37.291Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -688,6 +813,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/89/82/5ab5211079a151b6f661529369c0c8e98ec64cabf5c0cf22a0a05af124d8/jax-0.10.2-py3-none-any.whl", hash = "sha256:724d73c4678d8b06f6a6ab4db1b8a2fea8cd4f1e2c2564f99601634ec7b8d1c6", size = 3219515, upload-time = "2026-06-17T23:42:41.259Z" }, ] +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", extra = ["with-cuda"] }, + { name = "jaxlib" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.10.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bf/b0/e1d392ee24ecd53caa202194746cb12058befde5882d1f3307922829783f/jax_cuda12_pjrt-0.10.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b6f2d66548ee1ee910a836b5fe3d0ebbe1da04a62d0f468ba4822189036a32b3", size = 169847949, upload-time = "2026-06-17T23:42:44.729Z" }, + { url = "https://files.pythonhosted.org/packages/77/00/326facd192b7dc0731e43b30f5ada5ef235e2a3325a151d94ed3db041536/jax_cuda12_pjrt-0.10.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:806d1fd29038b6acf5a2b289dd62192abea977ee26aef60ea295b1d28a23acf8", size = 174364763, upload-time = "2026-06-17T23:42:49.931Z" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.10.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax-cuda12-pjrt" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/4a/382ec14e57d774148b079c3d251fc9b9e1ed7f92dd2327823f35e0a968a7/jax_cuda12_plugin-0.10.2-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:767a1482dbf652688403c4e22ff01470e1add13c48f9d817169fd8f7440afce5", size = 8093630, upload-time = "2026-06-17T23:42:57.019Z" }, + { url = "https://files.pythonhosted.org/packages/ca/25/19ee8936b00ca74a12bfaebd4703695cee2407f954a6e4e67cdd7f82b7bf/jax_cuda12_plugin-0.10.2-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:4eb6e8e0992fd9897db2468da33f276ae414ec84dc436804124404430502eeac", size = 8141521, upload-time = "2026-06-17T23:42:58.483Z" }, + { url = "https://files.pythonhosted.org/packages/94/33/a9c4ea148d31843764383b26fd66b38cfedad9a6dbf08db634a23d9312c3/jax_cuda12_plugin-0.10.2-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:69776ac849112f4fc2d6fa616f8772106daff38fcdb825c4e39f47e878f27610", size = 8093605, upload-time = "2026-06-17T23:43:00.071Z" }, + { url = "https://files.pythonhosted.org/packages/92/07/3098ae537e68d76bbc98ed61d53e9856f63f4f1bd64f8aeb4d432dad1290/jax_cuda12_plugin-0.10.2-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:948a988927cb10843b501a215181ccb1daed3fa70531b2f1ae0842fe1c08b05e", size = 8141236, upload-time = "2026-06-17T23:43:01.55Z" }, + { url = "https://files.pythonhosted.org/packages/72/0c/ff1e1975ef108cb2e45deacf5dea812822efaf3210af075f464abd40399c/jax_cuda12_plugin-0.10.2-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:febd9d3f488d182090758936042ad647a60470923715b48a385bc312cb943c63", size = 8108408, upload-time = "2026-06-17T23:43:03.398Z" }, + { url = "https://files.pythonhosted.org/packages/09/c2/04c70eb73beb4df23f49b72d159ca74ee31cccb1be95cdab7491a5ebd259/jax_cuda12_plugin-0.10.2-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:1f1e20fba0aaa3f28ab0a3273754c8c098fa40b74bf4dac5b76c9f2a70425487", size = 8151184, upload-time = "2026-06-17T23:43:05.106Z" }, + { url = "https://files.pythonhosted.org/packages/95/34/5d5b2993f6b7bbfe781920671de61036faffa5a8bf3740b4ca23d790b6f9/jax_cuda12_plugin-0.10.2-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:4946ec57ddb795e2da8288ceb1ceddd58c2a204a0f7765af89eae486a1d3ace7", size = 8094107, upload-time = "2026-06-17T23:43:06.546Z" }, + { url = "https://files.pythonhosted.org/packages/02/ac/85bc07d7a0766c7999f9bdfcf7f7cb2487fd0eb5d6bbe2bd41a3a54cd379/jax_cuda12_plugin-0.10.2-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:0c7a0204155585cc1c4fcb30c66b7f881b38e57e1f0d8e97d48e1654bf0d27f8", size = 8141889, upload-time = "2026-06-17T23:43:08.018Z" }, + { url = "https://files.pythonhosted.org/packages/f3/7d/9fa25e24b96ab06217e1e3dd311a80f3534076bb1b19d46cbc8cf9aa55b9/jax_cuda12_plugin-0.10.2-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:c26a7dd5e71c9edf203b95480f02ec4ad60c9c7e11e38b7310b1fd1dd03ba116", size = 8108639, upload-time = "2026-06-17T23:43:09.651Z" }, + { url = "https://files.pythonhosted.org/packages/20/c2/98fa02ba204f5dc88680128129fbc7cd833d5b5ddae2734af93ef85db853/jax_cuda12_plugin-0.10.2-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:c539b052056e8181082fa2635f17e098b3f625d87da53ba35470a2c13480801d", size = 8151505, upload-time = "2026-06-17T23:43:11.023Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" }, +] + [[package]] name = "jaxlib" version = "0.10.2" @@ -730,15 +906,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/38/c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0/jaxtyping-0.3.11-py3-none-any.whl", hash = "sha256:8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3", size = 56593, upload-time = "2026-06-13T18:35:22.01Z" }, ] -[[package]] -name = "joblib" -version = "1.5.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e8/5d/447af5ea094b9e4c4054f82e223ada074c552335b9b4b2d14bd9b35a67c4/joblib-1.5.2.tar.gz", hash = "sha256:3faa5c39054b2f03ca547da9b2f52fde67c06240c31853f306aea97f13647b55", size = 331077, upload-time = "2025-08-27T12:15:46.575Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/e8/685f47e0d754320684db4425a0967f7d3fa70126bffd76110b7009a0090f/joblib-1.5.2-py3-none-any.whl", hash = "sha256:4e1f0bdbb987e6d843c70cf43714cb276623def372df3c22fe5266b2670bc241", size = 308396, upload-time = "2025-08-27T12:15:45.188Z" }, -] - [[package]] name = "kiwisolver" version = "1.4.9" @@ -1173,6 +1340,113 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/6a/d3a169aaf8536cf228d56a09e04bcb713a2fe4410d4e2105b9419b5a9c89/numpy-2.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:016623417bb330d719d579daf2d6b9a01ddc52e41a9ed61a47f39fde46dcd865", size = 10686451, upload-time = "2026-06-21T20:57:49.313Z" }, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.9.2.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/a2/c96163a0fff1839c0c9548bbdeae7b853b867009e33b9b9264adc238b1cf/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:5572131a59c3eebeeb1c4c8144f772d49372c20124916e072a0e3fc30df421d5", size = 575012079, upload-time = "2026-04-08T18:51:47.303Z" }, + { url = "https://files.pythonhosted.org/packages/cb/c0/0a517bfe63ccd3b92eb254d264e28fca3c7cab75d07daea315250fb1bf73/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:e4f53a8ca8c5d6e8c492d0d0a3d565ecb59a751b19cfdaa4f6da0ab2104c1702", size = 581240110, upload-time = "2026-04-08T18:52:31.532Z" }, +] + +[[package]] +name = "nvidia-cuda-cccl-cu12" +version = "12.9.27" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/7e/82e49956b046bdc506c789235c587d9b3ef58b8bc1782258c1e247229647/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7898b38aa68beaa234d48f0868273702342a196d6e2e9d0ef058dca2390ebea", size = 3152245, upload-time = "2025-05-01T19:32:04.802Z" }, + { url = "https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:37869e17ce2e1ecec6eddf1927cca0f8c34e64fd848d40453df559091e2d7117", size = 3152243, upload-time = "2025-05-01T19:32:13.955Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/78/351b5c8cdbd9a6b4fb0d6ee73fb176dcdc1b6b6ad47c2ffff5ae8ca4a1f7/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:791853b030602c6a11d08b5578edfb957cadea06e9d3b26adbf8d036135a4afe", size = 10077166, upload-time = "2025-06-05T20:01:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:096bcf334f13e1984ba36685ad4c1d6347db214de03dbb6eebb237b41d9d934f", size = 10814997, upload-time = "2025-06-05T20:01:10.168Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/48/b54a06168a2190572a312bfe4ce443687773eb61367ced31e064953dd2f7/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5d6a0d32fdc7ea39917c20065614ae93add6f577d840233237ff08e9a38f58f0", size = 40546229, upload-time = "2025-06-05T20:01:53.357Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/85/e4af82cc9202023862090bfca4ea827d533329e925c758f0cde964cb54b7/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:210cf05005a447e29214e9ce50851e83fc5f4358df8b453155d5e1918094dcb4", size = 89568129, upload-time = "2025-06-05T20:02:41.973Z" }, + { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e0/0279bd94539fda525e0c8538db29b72a5a8495b0c12173113471d28bce78/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83469a846206f2a733db0c42e223589ab62fd2fabac4432d2f8802de4bded0a4", size = 3515012, upload-time = "2025-06-05T20:00:35.519Z" }, + { url = "https://files.pythonhosted.org/packages/bc/46/a92db19b8309581092a3add7e6fceb4c301a3fd233969856a8cbf042cd3c/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:25bba2dfb01d48a9b59ca474a1ac43c6ebf7011f1b0b8cc44f54eb6ac48a96c3", size = 3493179, upload-time = "2025-06-05T20:00:53.735Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.24.0.43" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/f1/cd42563325fa827f54ff30da05686c747652bdbd4cb5654cea54d7d0ad4f/nvidia_cudnn_cu12-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a42996943f0cd78ddfd61c8bf59361672a19b63e0491aa22a53d6fe63a3f854a", size = 856490582, upload-time = "2026-07-02T16:21:30.924Z" }, + { url = "https://files.pythonhosted.org/packages/10/13/b8887c869cf2471339a24b60d3c28e761facbb534935f572b61423371abb/nvidia_cudnn_cu12-9.24.0.43-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:f424192dd85e7d29f44be18df2dae4c80d32c67a29c0d42f5c283c40cfdf871c", size = 799083985, upload-time = "2026-07-02T16:25:37.467Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.4.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/2b/76445b0af890da61b501fde30650a1a4bd910607261b209cccb5235d3daa/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1a28c9b12260a1aa7a8fd12f5ebd82d027963d635ba82ff39a1acfa7c4c0fbcf", size = 200822453, upload-time = "2025-06-05T20:05:27.889Z" }, + { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.5.82" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" }, + { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.10.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" }, + { url = "https://files.pythonhosted.org/packages/12/46/b0fd4b04f86577921feb97d8e2cf028afe04f614d17fb5013de9282c9216/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:73060ce019ac064a057267c585bf1fd5a353734151f87472ff02b2c5c9984e78", size = 366465088, upload-time = "2025-06-05T20:08:20.413Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.28.9" @@ -1182,6 +1456,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4a/4e/44dbb46b3d1b0ec61afda8e84837870f2f9ace33c564317d59b70bc19d3e/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:485776daa8447da5da39681af455aa3b2c2586ddcf4af8772495e7c532c7e5ab", size = 296782137, upload-time = "2025-11-18T05:49:34.248Z" }, ] +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, + { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-cccl-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/9a/a60b98b629e3b7c4ef6385b08c0e169f2048382c75476f377da54b83a8ca/nvidia_nvshmem_cu12-3.7.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7de83f2aa2b58570b4058543fd3c3e992c19f96571380d46ec9b00e0062c6015", size = 229879072, upload-time = "2026-06-30T19:25:39.671Z" }, + { url = "https://files.pythonhosted.org/packages/61/a2/d1d065914f1860782171c6bd2e5eabcbe4488c66352b468569a6b69d4bf6/nvidia_nvshmem_cu12-3.7.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:faca8717c321f088d3c17490fbfc682bac214986e80b97e232c5ae94f900c45c", size = 230087166, upload-time = "2026-06-30T19:26:24.652Z" }, +] + [[package]] name = "omegaconf" version = "2.3.1" @@ -1403,6 +1698,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "portpicker" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "psutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/d0/cda2fc582f09510c84cd6b7d7b9e22a02d4e45dbad2b2ef1c6edd7847e00/portpicker-1.6.0.tar.gz", hash = "sha256:bd507fd6f96f65ee02781f2e674e9dc6c99bbfa6e3c39992e3916204c9d431fa", size = 25676, upload-time = "2023-08-15T04:37:08.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/2d/440e4d7041fff89f28f483733eb617127aa866135c2dc719e05893f089e1/portpicker-1.6.0-py3-none-any.whl", hash = "sha256:b2787a41404cf7edbe29b07b9e0ed863b09f2665dcc01c1eb0c2261c1e7d0755", size = 16613, upload-time = "2023-08-15T04:37:07.327Z" }, +] + [[package]] name = "prometheus-client" version = "0.25.0" @@ -1455,49 +1762,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, ] -[[package]] -name = "pyarrow" -version = "24.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/91/13/13e1069b351bdc3881266e11147ffccf687505dbb0ea74036237f5d454a5/pyarrow-24.0.0.tar.gz", hash = "sha256:85fe721a14dd823aca09127acbb06c3ca723efbd436c004f16bca601b04dcc83", size = 1180261, upload-time = "2026-04-21T10:51:25.837Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b4/a9/9686d9f07837f91f775e8932659192e02c74f9d8920524b480b85212cc68/pyarrow-24.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:6233c9ed9ab9d1db47de57d9753256d9dcffbf42db341576099f0fd9f6bf4810", size = 34981559, upload-time = "2026-04-21T10:47:22.17Z" }, - { url = "https://files.pythonhosted.org/packages/80/b6/0ddf0e9b6ead3474ab087ae598c76b031fc45532bf6a63f3a553440fb258/pyarrow-24.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:f7616236ec1bc2b15bfdec22a71ab38851c86f8f05ff64f379e1278cf20c634a", size = 36663654, upload-time = "2026-04-21T10:47:28.315Z" }, - { url = "https://files.pythonhosted.org/packages/7c/3b/926382efe8ce27ba729071d3566ade6dfb86bdf112f366000196b2f5780a/pyarrow-24.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:1617043b99bd33e5318ae18eb2919af09c71322ef1ca46566cdafc6e6712fb66", size = 45679394, upload-time = "2026-04-21T10:47:34.821Z" }, - { url = "https://files.pythonhosted.org/packages/b3/7a/829f7d9dfd37c207206081d6dad474d81dde29952401f07f2ba507814818/pyarrow-24.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:6165461f55ef6314f026de6638d661188e3455d3ec49834556a0ebbdbace18bb", size = 48863122, upload-time = "2026-04-21T10:47:42.056Z" }, - { url = "https://files.pythonhosted.org/packages/5f/e8/f88ce625fe8babaae64e8db2d417c7653adb3019b08aae85c5ed787dc816/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3b13dedfe76a0ad2d1d859b0811b53827a4e9d93a0bcb05cf59333ab4980cc7e", size = 49376032, upload-time = "2026-04-21T10:47:48.967Z" }, - { url = "https://files.pythonhosted.org/packages/36/7a/82c363caa145fff88fb475da50d3bf52bb024f61917be5424c3392eaf878/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:25ea65d868eb04015cd18e6df2fbe98f07e5bda2abefabcb88fce39a947716f6", size = 51929490, upload-time = "2026-04-21T10:47:55.981Z" }, - { url = "https://files.pythonhosted.org/packages/66/1c/e3e72c8014ad2743ca64a701652c733cc5cbcee15c0463a32a8c55518d9e/pyarrow-24.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:295f0a7f2e242dabd513737cf076007dc5b2d59237e3eca37b05c0c6446f3826", size = 27355660, upload-time = "2026-04-21T10:48:01.718Z" }, - { url = "https://files.pythonhosted.org/packages/6f/d3/a1abf004482026ddc17f4503db227787fa3cfe41ec5091ff20e4fea55e57/pyarrow-24.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:02b001b3ed4723caa44f6cd1af2d5c86aa2cf9971dacc2ffa55b21237713dfba", size = 34976759, upload-time = "2026-04-21T10:48:07.258Z" }, - { url = "https://files.pythonhosted.org/packages/4f/4a/34f0a36d28a2dd32225301b79daad44e243dc1a2bb77d43b60749be255c4/pyarrow-24.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:04920d6a71aabd08a0417709efce97d45ea8e6fb733d9ca9ecffb13c67839f68", size = 36658471, upload-time = "2026-04-21T10:48:13.347Z" }, - { url = "https://files.pythonhosted.org/packages/1f/78/543b94712ae8bb1a6023bcc1acf1a740fbff8286747c289cd9468fced2a5/pyarrow-24.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a964266397740257f16f7bb2e4f08a0c81454004beab8ff59dd531b73610e9f2", size = 45675981, upload-time = "2026-04-21T10:48:20.201Z" }, - { url = "https://files.pythonhosted.org/packages/84/9f/8fb7c222b100d314137fa40ec050de56cd8c6d957d1cfff685ce72f15b17/pyarrow-24.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6f066b179d68c413374294bc1735f68475457c933258df594443bb9d88ddc2a0", size = 48859172, upload-time = "2026-04-21T10:48:27.541Z" }, - { url = "https://files.pythonhosted.org/packages/a7/d3/1ea72538e6c8b3b475ed78d1049a2c518e655761ea50fe1171fc855fcab7/pyarrow-24.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1183baeb14c5f587b1ec52831e665718ce632caab84b7cd6b85fd44f96114495", size = 49385733, upload-time = "2026-04-21T10:48:34.7Z" }, - { url = "https://files.pythonhosted.org/packages/c3/be/c3d8b06a1ba35f2260f8e1f771abbee7d5e345c0937aab90675706b1690a/pyarrow-24.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:806f24b4085453c197a5078218d1ee08783ebbba271badd153d1ae22a3ee804f", size = 51934335, upload-time = "2026-04-21T10:48:42.099Z" }, - { url = "https://files.pythonhosted.org/packages/9c/62/89e07a1e7329d2cde3e3c6994ba0839a24977a2beda8be6005ea3d860b99/pyarrow-24.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:e4505fc6583f7b05ab854934896bcac8253b04ac1171a77dfb73efef92076d91", size = 27271748, upload-time = "2026-04-21T10:49:42.532Z" }, - { url = "https://files.pythonhosted.org/packages/17/1a/cff3a59f80b5b1658549d46611b67163f65e0664431c076ad728bf9d5af4/pyarrow-24.0.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:1a4e45017efbf115032e4475ee876d525e0e36c742214fbe405332480ecd6275", size = 35238554, upload-time = "2026-04-21T10:48:48.526Z" }, - { url = "https://files.pythonhosted.org/packages/a8/99/cce0f42a327bfef2c420fb6078a3eb834826e5d6697bf3009fe11d2ad051/pyarrow-24.0.0-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:7986f1fa71cee060ad00758bcc79d3a93bab8559bf978fab9e53472a2e25a17b", size = 36782301, upload-time = "2026-04-21T10:48:55.181Z" }, - { url = "https://files.pythonhosted.org/packages/2a/66/8e560d5ff6793ca29aca213c53eec0dd482dd46cb93b2819e5aab52e4252/pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:d3e0b61e8efb24ed38898e5cdc5fffa9124be480008d401a1f8071500494ae42", size = 45721929, upload-time = "2026-04-21T10:49:03.676Z" }, - { url = "https://files.pythonhosted.org/packages/27/0c/a26e25505d030716e078d9f16eb74973cbf0b33b672884e9f9da1c83b871/pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:55a3bc1e3df3b5567b7d27ef551b2283f0c68a5e86f1cd56abc569da4f31335b", size = 48825365, upload-time = "2026-04-21T10:49:11.714Z" }, - { url = "https://files.pythonhosted.org/packages/5f/eb/771f9ecb0c65e73fe9dccdd1717901b9594f08c4515d000c7c62df573811/pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:641f795b361874ac9da5294f8f443dfdbee355cf2bd9e3b8d97aaac2306b9b37", size = 49451819, upload-time = "2026-04-21T10:49:21.474Z" }, - { url = "https://files.pythonhosted.org/packages/48/da/61ae89a88732f5a785646f3ec6125dbb640fa98a540eb2b9889caa561403/pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8adc8e6ce5fccf5dc707046ae4914fd537def529709cc0d285d37a7f9cd442ca", size = 51909252, upload-time = "2026-04-21T10:49:31.164Z" }, - { url = "https://files.pythonhosted.org/packages/cb/1a/8dd5cafab7b66573fa91c03d06d213356ad4edd71813aa75e08ce2b3a844/pyarrow-24.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:9b18371ad2f44044b81a8d23bc2d8a9b6a6226dca775e8e16cfee640473d6c5d", size = 27388127, upload-time = "2026-04-21T10:49:37.334Z" }, - { url = "https://files.pythonhosted.org/packages/ad/80/d022a34ff05d2cbedd8ccf841fc1f532ecfa9eb5ed1711b56d0e0ea71fc9/pyarrow-24.0.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:1cc9057f0319e26333b357e17f3c2c022f1a83739b48a88b25bfd5fa2dc18838", size = 35007997, upload-time = "2026-04-21T10:49:48.796Z" }, - { url = "https://files.pythonhosted.org/packages/1a/ff/f01485fda6f4e5d441afb8dd5e7681e4db18826c1e271852f5d3957d6a80/pyarrow-24.0.0-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:e6f1278ee4785b6db21229374a1c9e54ec7c549de5d1efc9630b6207de7e170b", size = 36678720, upload-time = "2026-04-21T10:49:55.858Z" }, - { url = "https://files.pythonhosted.org/packages/9e/c2/2d2d5fea814237923f71b36495211f20b43a1576f9a4d6da7e751a64ec6f/pyarrow-24.0.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:adbbedc55506cbdabb830890444fb856bfb0060c46c6f8026c6c2f2cf86ae795", size = 45741852, upload-time = "2026-04-21T10:50:04.624Z" }, - { url = "https://files.pythonhosted.org/packages/8e/3a/28ba9c1c1ebdbb5f1b94dfebb46f207e52e6a554b7fe4132540fde29a3a0/pyarrow-24.0.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:ae8a1145af31d903fa9bb166824d7abe9b4681a000b0159c9fb99c11bc11ad26", size = 48889852, upload-time = "2026-04-21T10:50:12.293Z" }, - { url = "https://files.pythonhosted.org/packages/df/51/4a389acfd31dca009f8fb82d7f510bb4130f2b3a8e18cf00194d0687d8ac/pyarrow-24.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:d7027eba1df3b2069e2e8d80f644fa0918b68c46432af3d088ddd390d063ecde", size = 49445207, upload-time = "2026-04-21T10:50:20.677Z" }, - { url = "https://files.pythonhosted.org/packages/19/4b/0bab2b23d2ae901b1b9a03c0efd4b2d070256f8ce3fc43f6e58c167b2081/pyarrow-24.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e56a1ffe9bf7b727432b89104cc0849c21582949dd7bdcb34f17b2001a351a76", size = 51954117, upload-time = "2026-04-21T10:50:29.14Z" }, - { url = "https://files.pythonhosted.org/packages/29/88/f4e9145da0417b3d2c12035a8492b35ff4a3dbc653e614fcfb51d9dedb38/pyarrow-24.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:38be1808cdd068605b787e6ca9119b27eb275a0234e50212c3492331680c3b1e", size = 28001155, upload-time = "2026-04-21T10:51:22.337Z" }, - { url = "https://files.pythonhosted.org/packages/79/4f/46a49a63f43526da895b1a45bbb51d5baf8e4d77159f8528fc3e5490007f/pyarrow-24.0.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:418e48ce50a45a6a6c73c454677203a9c75c966cb1e92ca3370959185f197a05", size = 35250387, upload-time = "2026-04-21T10:50:35.552Z" }, - { url = "https://files.pythonhosted.org/packages/a0/da/d5e0cd5ef00796922404806d5f00325cdadc3441ce2c13fe7115f2df9a64/pyarrow-24.0.0-cp314-cp314t-macosx_12_0_x86_64.whl", hash = "sha256:2f16197705a230a78270cdd4ea8a1d57e86b2fdcbc34a1f6aebc72e65c986f9a", size = 36797102, upload-time = "2026-04-21T10:50:42.417Z" }, - { url = "https://files.pythonhosted.org/packages/34/c7/5904145b0a593a05236c882933d439b5720f0a145381179063722fbfc123/pyarrow-24.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:fb24ac194bfc5e86839d7dcd52092ee31e5fe6733fe11f5e3b06ef0812b20072", size = 45745118, upload-time = "2026-04-21T10:50:49.324Z" }, - { url = "https://files.pythonhosted.org/packages/13/d3/cca42fe166d1c6e4d5b80e530b7949104d10e17508a90ae202dac205ce2a/pyarrow-24.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:9700ebd9a51f5895ce75ff4ac4b3c47a7d4b42bc618be8e713e5d56bacf5f931", size = 48844765, upload-time = "2026-04-21T10:50:55.579Z" }, - { url = "https://files.pythonhosted.org/packages/b0/49/942c3b79878ba928324d1e17c274ed84581db8c0a749b24bcf4cbdf15bd3/pyarrow-24.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:d8ddd2768da81d3ee08cfea9b597f4abb4e8e1dc8ae7e204b608d23a0d3ab699", size = 49471890, upload-time = "2026-04-21T10:51:02.439Z" }, - { url = "https://files.pythonhosted.org/packages/76/97/ff71431000a75d84135a1ace5ca4ba11726a231a8007bbb320a4c54075d5/pyarrow-24.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:61a3d7eaa97a14768b542f3d284dc6400dd2470d9f080708b13cd46b6ae18136", size = 51932250, upload-time = "2026-04-21T10:51:10.576Z" }, - { url = "https://files.pythonhosted.org/packages/51/be/6f79d55816d5c22557cf27533543d5d70dfe692adfbee4b99f2760674f38/pyarrow-24.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c91d00057f23b8d353039520dc3a6c09d8608164c692e9f59a175a42b2ae0c19", size = 28131282, upload-time = "2026-04-21T10:51:16.815Z" }, -] - [[package]] name = "pycodestyle" version = "2.14.0" @@ -1507,6 +1771,96 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d7/27/a58ddaf8c588a3ef080db9d0b7e0b97215cee3a45df74f3a94dbbf5c893a/pycodestyle-2.14.0-py2.py3-none-any.whl", hash = "sha256:dd6bf7cb4ee77f8e016f9c8e74a35ddd9f67e1d5fd4184d86c3b98e07099f42d", size = 31594, upload-time = "2025-06-20T18:49:47.491Z" }, ] +[[package]] +name = "pydantic" +version = "2.13.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/8c/af022f0af448d7747c5154288d46b5f2bc5f17366eaa0e23e9aa04d59f3b/pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2", size = 2106158, upload-time = "2026-05-06T13:38:57.215Z" }, + { url = "https://files.pythonhosted.org/packages/19/95/6195171e385007300f0f5574592e467c568becce2d937a0b6804f218bc49/pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f", size = 1951724, upload-time = "2026-05-06T13:37:02.697Z" }, + { url = "https://files.pythonhosted.org/packages/8e/bc/f47d1ff9cbb1620e1b5b697eef06010035735f07820180e74178226b27b3/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7", size = 1975742, upload-time = "2026-05-06T13:37:09.448Z" }, + { url = "https://files.pythonhosted.org/packages/5b/11/9b9a5b0306345664a2da6410877af6e8082481b5884b3ddd78d47c6013ce/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7", size = 2052418, upload-time = "2026-05-06T13:37:38.234Z" }, + { url = "https://files.pythonhosted.org/packages/f1/b7/a65fec226f5d78fc39f4a13c4cc0c768c22b113438f60c14adc9d2865038/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712", size = 2232274, upload-time = "2026-05-06T13:38:27.753Z" }, + { url = "https://files.pythonhosted.org/packages/68/f0/92039db98b907ef49269a8271f67db9cb78ae2fc68062ef7e4e77adb5f61/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4", size = 2309940, upload-time = "2026-05-06T13:38:05.353Z" }, + { url = "https://files.pythonhosted.org/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" }, + { url = "https://files.pythonhosted.org/packages/22/37/a8aca44d40d737dde2bc05b3c6c07dff0de07ce6f82e9f3167aeaf4d5dea/pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987", size = 2136854, upload-time = "2026-05-06T13:40:22.59Z" }, + { url = "https://files.pythonhosted.org/packages/24/99/fcef1b79238c06a8cbec70819ac722ba76e02bc8ada9b0fd66eba40da01b/pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b", size = 2180306, upload-time = "2026-05-06T13:40:10.666Z" }, + { url = "https://files.pythonhosted.org/packages/ae/6c/fc44000918855b42779d007ae63b0532794739027b2f417321cddbc44f6a/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458", size = 2190044, upload-time = "2026-05-06T13:40:43.231Z" }, + { url = "https://files.pythonhosted.org/packages/6b/65/d9cadc9f1920d7a127ad2edba16c1db7916e59719285cd6c94600b0080ba/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b", size = 2329133, upload-time = "2026-05-06T13:39:57.365Z" }, + { url = "https://files.pythonhosted.org/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" }, + { url = "https://files.pythonhosted.org/packages/47/bd/6f2fc8188f31bf10590f1e98e7b306336161fac930a8c514cd7bd828c7dc/pydantic_core-2.46.4-cp312-cp312-win32.whl", hash = "sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894", size = 1974823, upload-time = "2026-05-06T13:40:47.985Z" }, + { url = "https://files.pythonhosted.org/packages/40/8c/985c1d41ea1107c2534abd9870e4ed5c8e7669b5c308297835c001e7a1c4/pydantic_core-2.46.4-cp312-cp312-win_amd64.whl", hash = "sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89", size = 2072919, upload-time = "2026-05-06T13:39:21.153Z" }, + { url = "https://files.pythonhosted.org/packages/c4/ba/f463d006e0c47373ca7ec5e1a261c59dc01ef4d62b2657af925fb0deee3a/pydantic_core-2.46.4-cp312-cp312-win_arm64.whl", hash = "sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a", size = 2027604, upload-time = "2026-05-06T13:39:03.753Z" }, + { url = "https://files.pythonhosted.org/packages/51/a2/5d30b469c5267a17b39dec53208222f76a8d351dfac4af661888c5aee77d/pydantic_core-2.46.4-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008", size = 2106306, upload-time = "2026-05-06T13:37:48.029Z" }, + { url = "https://files.pythonhosted.org/packages/c1/81/4fa520eaffa8bd7d1525e644cd6d39e7d60b1592bc5b516693c7340b50f1/pydantic_core-2.46.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4", size = 1951906, upload-time = "2026-05-06T13:37:17.012Z" }, + { url = "https://files.pythonhosted.org/packages/03/d5/fd02da45b659668b05923b17ba3a0100a0a3d5541e3bd8fcc4ecb711309e/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76", size = 1976802, upload-time = "2026-05-06T13:37:35.113Z" }, + { url = "https://files.pythonhosted.org/packages/21/f2/95727e1368be3d3ed485eaab7adbd7dda408f33f7a36e8b48e0144002b91/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3", size = 2052446, upload-time = "2026-05-06T13:37:12.313Z" }, + { url = "https://files.pythonhosted.org/packages/9c/86/5d99feea3f77c7234b8718075b23db11532773c1a0dbd9b9490215dc2eeb/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76", size = 2232757, upload-time = "2026-05-06T13:39:01.149Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3a/508ac615935ef7588cf6d9e9b91309fdc2da751af865e02a9098de88258c/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4", size = 2309275, upload-time = "2026-05-06T13:37:41.406Z" }, + { url = "https://files.pythonhosted.org/packages/07/f8/41db9de19d7987d6b04715a02b3b40aea467000275d9d758ffaa31af7d50/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a", size = 2094467, upload-time = "2026-05-06T13:39:18.847Z" }, + { url = "https://files.pythonhosted.org/packages/2c/e2/f35033184cb11d0052daf4416e8e10a502ea2ac006fc4f459aee872727d1/pydantic_core-2.46.4-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262", size = 2134417, upload-time = "2026-05-06T13:40:17.944Z" }, + { url = "https://files.pythonhosted.org/packages/7e/7b/6ceeb1cc90e193862f444ebe373d8fdf613f0a82572dde03fb10734c6c71/pydantic_core-2.46.4-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e", size = 2179782, upload-time = "2026-05-06T13:40:32.618Z" }, + { url = "https://files.pythonhosted.org/packages/5a/f2/c8d7773ede6af08036423a00ae0ceffce266c3c52a096c435d68c896083f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd", size = 2188782, upload-time = "2026-05-06T13:36:51.018Z" }, + { url = "https://files.pythonhosted.org/packages/59/31/0c864784e31f09f05cdd87606f08923b9c9e7f6e51dd27f20f62f975ce9f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be", size = 2328334, upload-time = "2026-05-06T13:40:37.764Z" }, + { url = "https://files.pythonhosted.org/packages/c2/eb/4f6c8a41efa30baa755590f4141abf3a8c370fab610915733e74134a7270/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d", size = 2372986, upload-time = "2026-05-06T13:39:34.152Z" }, + { url = "https://files.pythonhosted.org/packages/5b/24/b375a480d53113860c299764bfe9f349a3dc9108b3adc0d7f0d786492ebf/pydantic_core-2.46.4-cp313-cp313-win32.whl", hash = "sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb", size = 1973693, upload-time = "2026-05-06T13:37:55.072Z" }, + { url = "https://files.pythonhosted.org/packages/7e/e8/cff247591966f2d22ec8c003cd7587e27b7ba7b81ab2fb888e3ab75dc285/pydantic_core-2.46.4-cp313-cp313-win_amd64.whl", hash = "sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292", size = 2071819, upload-time = "2026-05-06T13:38:49.139Z" }, + { url = "https://files.pythonhosted.org/packages/c6/1a/f4aee670d5670e9e148e0c82c7db98d780be566c6e6a97ee8035528ca0b3/pydantic_core-2.46.4-cp313-cp313-win_arm64.whl", hash = "sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d", size = 2027411, upload-time = "2026-05-06T13:40:45.796Z" }, + { url = "https://files.pythonhosted.org/packages/8d/74/228a26ddad29c6672b805d9fd78e8d251cd04004fa7eed0e622096cd0250/pydantic_core-2.46.4-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb", size = 2102079, upload-time = "2026-05-06T13:38:41.019Z" }, + { url = "https://files.pythonhosted.org/packages/ad/1f/8970b150a4b4365623ae00fc88603491f763c627311ae8031e3111356d6e/pydantic_core-2.46.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462", size = 1952179, upload-time = "2026-05-06T13:36:59.812Z" }, + { url = "https://files.pythonhosted.org/packages/95/30/5211a831ae054928054b2f79731661087a2bc5c01e825c672b3a4a8f1b3e/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9", size = 1978926, upload-time = "2026-05-06T13:37:39.933Z" }, + { url = "https://files.pythonhosted.org/packages/57/e9/689668733b1eb67adeef047db3c2e8788fcf65a7fd9c9e2b46b7744fe245/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4", size = 2046785, upload-time = "2026-05-06T13:38:01.995Z" }, + { url = "https://files.pythonhosted.org/packages/60/d9/6715260422ff50a2109878fd24d948a6c3446bb2664f34ee78cd972b3acd/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914", size = 2228733, upload-time = "2026-05-06T13:40:50.371Z" }, + { url = "https://files.pythonhosted.org/packages/18/ae/fdb2f64316afca925640f8e70bb1a564b0ec2721c1389e25b8eb4bf9a299/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28", size = 2307534, upload-time = "2026-05-06T13:37:21.531Z" }, + { url = "https://files.pythonhosted.org/packages/89/1d/8eff589b45bb8190a9d12c49cfad0f176a5cbd1534908a6b5125e2886239/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b", size = 2099732, upload-time = "2026-05-06T13:39:31.942Z" }, + { url = "https://files.pythonhosted.org/packages/06/d5/ee5a3366637fee41dee51a1fc91562dcf12ddbc68fda34e6b253da2324bb/pydantic_core-2.46.4-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c", size = 2129627, upload-time = "2026-05-06T13:37:25.033Z" }, + { url = "https://files.pythonhosted.org/packages/94/33/2414be571d2c6a6c4d08be21f9292b6d3fdb08949a97b6dfe985017821db/pydantic_core-2.46.4-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb", size = 2179141, upload-time = "2026-05-06T13:37:14.046Z" }, + { url = "https://files.pythonhosted.org/packages/7b/79/7daa95be995be0eecc4cf75064cb33f9bbbfe3fe0158caf2f0d4a996a5c7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898", size = 2184325, upload-time = "2026-05-06T13:36:53.615Z" }, + { url = "https://files.pythonhosted.org/packages/9f/cb/d0a382f5c0de8a222dc61c65348e0ce831b1f68e0a018450d31c2cace3a5/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e", size = 2323990, upload-time = "2026-05-06T13:40:29.971Z" }, + { url = "https://files.pythonhosted.org/packages/05/db/d9ba624cc4a5aced1598e88c04fdbd8310c8a69b9d38b9a3d39ce3a61ed7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519", size = 2369978, upload-time = "2026-05-06T13:37:23.027Z" }, + { url = "https://files.pythonhosted.org/packages/f2/20/d15df15ba918c423461905802bfd2981c3af0bfa0e40d05e13edbfa48bc3/pydantic_core-2.46.4-cp314-cp314-win32.whl", hash = "sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4", size = 1966354, upload-time = "2026-05-06T13:38:03.499Z" }, + { url = "https://files.pythonhosted.org/packages/fc/b6/6b8de4c0a7d7ab3004c439c80c5c1e0a3e8d78bbae19379b01960383d9e5/pydantic_core-2.46.4-cp314-cp314-win_amd64.whl", hash = "sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac", size = 2072238, upload-time = "2026-05-06T13:39:40.807Z" }, + { url = "https://files.pythonhosted.org/packages/32/36/51eb763beec1f4cf59b1db243a7dcc39cbb41230f050a09b9d69faaf0a48/pydantic_core-2.46.4-cp314-cp314-win_arm64.whl", hash = "sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a", size = 2018251, upload-time = "2026-05-06T13:37:26.72Z" }, + { url = "https://files.pythonhosted.org/packages/e8/91/855af51d625b23aa987116a19e231d2aaef9c4a415273ddc189b79a45fee/pydantic_core-2.46.4-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0", size = 2099593, upload-time = "2026-05-06T13:39:47.682Z" }, + { url = "https://files.pythonhosted.org/packages/fb/1b/8784a54c65edb5f49f0a14d6977cf1b209bba85a4c77445b255c2de58ab3/pydantic_core-2.46.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d", size = 1935226, upload-time = "2026-05-06T13:40:40.428Z" }, + { url = "https://files.pythonhosted.org/packages/e8/e7/1955d28d1afc56dd4b3ad7cc0cf39df1b9852964cf16e5d13912756d6d6b/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b", size = 1974605, upload-time = "2026-05-06T13:37:32.029Z" }, + { url = "https://files.pythonhosted.org/packages/93/e2/3fedbf0ba7a22850e6e9fd78117f1c0f10f950182344d8a6c535d468fdd8/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000", size = 2030777, upload-time = "2026-05-06T13:38:55.239Z" }, + { url = "https://files.pythonhosted.org/packages/f8/61/46be275fcaaba0b4f5b9669dd852267ce1ff616592dccf7a7845588df091/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e", size = 2236641, upload-time = "2026-05-06T13:37:08.096Z" }, + { url = "https://files.pythonhosted.org/packages/60/db/12e93e46a8bac9988be3c016860f83293daea8c716c029c9ace279036f2f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd", size = 2286404, upload-time = "2026-05-06T13:40:20.221Z" }, + { url = "https://files.pythonhosted.org/packages/e2/4a/4d8b19008f38d31c53b8219cfedc2e3d5de5fe99d90076b7e767de29274f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3", size = 2109219, upload-time = "2026-05-06T13:38:12.153Z" }, + { url = "https://files.pythonhosted.org/packages/88/70/3cbc40978fefb7bb09c6708d40d4ad1a5d70fd7213c3d17f971de868ec1f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7", size = 2110594, upload-time = "2026-05-06T13:40:02.971Z" }, + { url = "https://files.pythonhosted.org/packages/9d/20/b8d36736216e29491125531685b2f9e61aa5b4b2599893f8268551da3338/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff", size = 2159542, upload-time = "2026-05-06T13:39:27.506Z" }, + { url = "https://files.pythonhosted.org/packages/1d/a2/367df868eb584dacf6bf82a389272406d7178e301c4ac82545ab98bc2dd9/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424", size = 2168146, upload-time = "2026-05-06T13:38:31.93Z" }, + { url = "https://files.pythonhosted.org/packages/c1/b8/4460f77f7e201893f649a29ab355dddd3beee8a97bcb1a320db414f9a06e/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6", size = 2306309, upload-time = "2026-05-06T13:37:44.717Z" }, + { url = "https://files.pythonhosted.org/packages/64/c4/be2639293acd87dc8ddbcec41a73cee9b2ebf996fe6d892a1a74e88ad3f7/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565", size = 2369736, upload-time = "2026-05-06T13:37:05.645Z" }, + { url = "https://files.pythonhosted.org/packages/30/a6/9f9f380dbb301f67023bf8f707aaa75daadf84f7152d95c410fd7e81d994/pydantic_core-2.46.4-cp314-cp314t-win32.whl", hash = "sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02", size = 1955575, upload-time = "2026-05-06T13:38:51.116Z" }, + { url = "https://files.pythonhosted.org/packages/40/1f/f1eb9eb350e795d1af8586289746f5c5677d16043040d63710e22abc43c9/pydantic_core-2.46.4-cp314-cp314t-win_amd64.whl", hash = "sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5", size = 2051624, upload-time = "2026-05-06T13:38:21.672Z" }, + { url = "https://files.pythonhosted.org/packages/f6/d2/42dd53d0a85c27606f316d3aa5d2869c4e8470a5ed6dec30e4a1abe19192/pydantic_core-2.46.4-cp314-cp314t-win_arm64.whl", hash = "sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596", size = 2017325, upload-time = "2026-05-06T13:40:52.723Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1d/8987ad40f65ae1432753072f214fb5c74fe47ffbd0698bb9cbbb585664f8/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7", size = 2095527, upload-time = "2026-05-06T13:39:52.283Z" }, + { url = "https://files.pythonhosted.org/packages/64/d3/84c282a7eee1d3ac4c0377546ef5a1ea436ce26840d9ac3b7ed54a377507/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df", size = 1936024, upload-time = "2026-05-06T13:40:15.671Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ca/eac61596cdeb4d7e174d3dc0bd8a6238f14f75f97a24e7b7db4c7e7340a0/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526", size = 1990696, upload-time = "2026-05-06T13:38:34.717Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, +] + [[package]] name = "pyflakes" version = "3.4.0" @@ -1659,40 +2013,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, ] -[[package]] -name = "scikit-learn" -version = "1.7.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "joblib" }, - { name = "numpy" }, - { name = "scipy" }, - { name = "threadpoolctl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/aa/3996e2196075689afb9fce0410ebdb4a09099d7964d061d7213700204409/scikit_learn-1.7.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8d91a97fa2b706943822398ab943cde71858a50245e31bc71dba62aab1d60a96", size = 9259818, upload-time = "2025-09-09T08:20:43.19Z" }, - { url = "https://files.pythonhosted.org/packages/43/5d/779320063e88af9c4a7c2cf463ff11c21ac9c8bd730c4a294b0000b666c9/scikit_learn-1.7.2-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:acbc0f5fd2edd3432a22c69bed78e837c70cf896cd7993d71d51ba6708507476", size = 8636997, upload-time = "2025-09-09T08:20:45.468Z" }, - { url = "https://files.pythonhosted.org/packages/5c/d0/0c577d9325b05594fdd33aa970bf53fb673f051a45496842caee13cfd7fe/scikit_learn-1.7.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e5bf3d930aee75a65478df91ac1225ff89cd28e9ac7bd1196853a9229b6adb0b", size = 9478381, upload-time = "2025-09-09T08:20:47.982Z" }, - { url = "https://files.pythonhosted.org/packages/82/70/8bf44b933837ba8494ca0fc9a9ab60f1c13b062ad0197f60a56e2fc4c43e/scikit_learn-1.7.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b4d6e9deed1a47aca9fe2f267ab8e8fe82ee20b4526b2c0cd9e135cea10feb44", size = 9300296, upload-time = "2025-09-09T08:20:50.366Z" }, - { url = "https://files.pythonhosted.org/packages/c6/99/ed35197a158f1fdc2fe7c3680e9c70d0128f662e1fee4ed495f4b5e13db0/scikit_learn-1.7.2-cp312-cp312-win_amd64.whl", hash = "sha256:6088aa475f0785e01bcf8529f55280a3d7d298679f50c0bb70a2364a82d0b290", size = 8731256, upload-time = "2025-09-09T08:20:52.627Z" }, - { url = "https://files.pythonhosted.org/packages/ae/93/a3038cb0293037fd335f77f31fe053b89c72f17b1c8908c576c29d953e84/scikit_learn-1.7.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0b7dacaa05e5d76759fb071558a8b5130f4845166d88654a0f9bdf3eb57851b7", size = 9212382, upload-time = "2025-09-09T08:20:54.731Z" }, - { url = "https://files.pythonhosted.org/packages/40/dd/9a88879b0c1104259136146e4742026b52df8540c39fec21a6383f8292c7/scikit_learn-1.7.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:abebbd61ad9e1deed54cca45caea8ad5f79e1b93173dece40bb8e0c658dbe6fe", size = 8592042, upload-time = "2025-09-09T08:20:57.313Z" }, - { url = "https://files.pythonhosted.org/packages/46/af/c5e286471b7d10871b811b72ae794ac5fe2989c0a2df07f0ec723030f5f5/scikit_learn-1.7.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:502c18e39849c0ea1a5d681af1dbcf15f6cce601aebb657aabbfe84133c1907f", size = 9434180, upload-time = "2025-09-09T08:20:59.671Z" }, - { url = "https://files.pythonhosted.org/packages/f1/fd/df59faa53312d585023b2da27e866524ffb8faf87a68516c23896c718320/scikit_learn-1.7.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7a4c328a71785382fe3fe676a9ecf2c86189249beff90bf85e22bdb7efaf9ae0", size = 9283660, upload-time = "2025-09-09T08:21:01.71Z" }, - { url = "https://files.pythonhosted.org/packages/a7/c7/03000262759d7b6f38c836ff9d512f438a70d8a8ddae68ee80de72dcfb63/scikit_learn-1.7.2-cp313-cp313-win_amd64.whl", hash = "sha256:63a9afd6f7b229aad94618c01c252ce9e6fa97918c5ca19c9a17a087d819440c", size = 8702057, upload-time = "2025-09-09T08:21:04.234Z" }, - { url = "https://files.pythonhosted.org/packages/55/87/ef5eb1f267084532c8e4aef98a28b6ffe7425acbfd64b5e2f2e066bc29b3/scikit_learn-1.7.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:9acb6c5e867447b4e1390930e3944a005e2cb115922e693c08a323421a6966e8", size = 9558731, upload-time = "2025-09-09T08:21:06.381Z" }, - { url = "https://files.pythonhosted.org/packages/93/f8/6c1e3fc14b10118068d7938878a9f3f4e6d7b74a8ddb1e5bed65159ccda8/scikit_learn-1.7.2-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:2a41e2a0ef45063e654152ec9d8bcfc39f7afce35b08902bfe290c2498a67a6a", size = 9038852, upload-time = "2025-09-09T08:21:08.628Z" }, - { url = "https://files.pythonhosted.org/packages/83/87/066cafc896ee540c34becf95d30375fe5cbe93c3b75a0ee9aa852cd60021/scikit_learn-1.7.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:98335fb98509b73385b3ab2bd0639b1f610541d3988ee675c670371d6a87aa7c", size = 9527094, upload-time = "2025-09-09T08:21:11.486Z" }, - { url = "https://files.pythonhosted.org/packages/9c/2b/4903e1ccafa1f6453b1ab78413938c8800633988c838aa0be386cbb33072/scikit_learn-1.7.2-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:191e5550980d45449126e23ed1d5e9e24b2c68329ee1f691a3987476e115e09c", size = 9367436, upload-time = "2025-09-09T08:21:13.602Z" }, - { url = "https://files.pythonhosted.org/packages/b5/aa/8444be3cfb10451617ff9d177b3c190288f4563e6c50ff02728be67ad094/scikit_learn-1.7.2-cp313-cp313t-win_amd64.whl", hash = "sha256:57dc4deb1d3762c75d685507fbd0bc17160144b2f2ba4ccea5dc285ab0d0e973", size = 9275749, upload-time = "2025-09-09T08:21:15.96Z" }, - { url = "https://files.pythonhosted.org/packages/d9/82/dee5acf66837852e8e68df6d8d3a6cb22d3df997b733b032f513d95205b7/scikit_learn-1.7.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fa8f63940e29c82d1e67a45d5297bdebbcb585f5a5a50c4914cc2e852ab77f33", size = 9208906, upload-time = "2025-09-09T08:21:18.557Z" }, - { url = "https://files.pythonhosted.org/packages/3c/30/9029e54e17b87cb7d50d51a5926429c683d5b4c1732f0507a6c3bed9bf65/scikit_learn-1.7.2-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:f95dc55b7902b91331fa4e5845dd5bde0580c9cd9612b1b2791b7e80c3d32615", size = 8627836, upload-time = "2025-09-09T08:21:20.695Z" }, - { url = "https://files.pythonhosted.org/packages/60/18/4a52c635c71b536879f4b971c2cedf32c35ee78f48367885ed8025d1f7ee/scikit_learn-1.7.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9656e4a53e54578ad10a434dc1f993330568cfee176dff07112b8785fb413106", size = 9426236, upload-time = "2025-09-09T08:21:22.645Z" }, - { url = "https://files.pythonhosted.org/packages/99/7e/290362f6ab582128c53445458a5befd471ed1ea37953d5bcf80604619250/scikit_learn-1.7.2-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96dc05a854add0e50d3f47a1ef21a10a595016da5b007c7d9cd9d0bffd1fcc61", size = 9312593, upload-time = "2025-09-09T08:21:24.65Z" }, - { url = "https://files.pythonhosted.org/packages/8e/87/24f541b6d62b1794939ae6422f8023703bbf6900378b2b34e0b4384dfefd/scikit_learn-1.7.2-cp314-cp314-win_amd64.whl", hash = "sha256:bb24510ed3f9f61476181e4db51ce801e2ba37541def12dc9333b946fc7a9cf8", size = 8820007, upload-time = "2025-09-09T08:21:26.713Z" }, -] - [[package]] name = "scipy" version = "1.16.3" @@ -1754,6 +2074,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/47/a494741db7280eae6dc033510c319e34d42dd41b7ac0c7ead39354d1a2b5/scipy-1.16.3-cp314-cp314t-win_arm64.whl", hash = "sha256:21d9d6b197227a12dcbf9633320a4e34c6b0e51c57268df255a0942983bac562", size = 26464127, upload-time = "2025-10-28T17:38:11.34Z" }, ] +[[package]] +name = "sentry-sdk" +version = "2.64.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/60/31/b7341f156a5f6f36f0b4845d6f1c28a2ae4799171dba7007f3a1e9b234b4/sentry_sdk-2.64.0.tar.gz", hash = "sha256:68be2c29e14ae310f8a39e1a79916b6d85c6cb41dcce789d14ff05fe293e4c55", size = 921020, upload-time = "2026-06-30T08:13:47.682Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/a8/3fb9a4319efa3b26f5be0e90e6d8918df43fa7c7e977d26390f589501d82/sentry_sdk-2.64.0-py3-none-any.whl", hash = "sha256:715ea91ca860a819e8d8a50a7bde3a80d0df3b4ed7b6660a20fb9a2d084188f1", size = 498901, upload-time = "2026-06-30T08:13:45.566Z" }, +] + [[package]] name = "shellingham" version = "1.5.4" @@ -1825,6 +2158,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + [[package]] name = "tensorstore" version = "0.1.84" @@ -1856,15 +2198,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/74/35a1d41343f86f6e2ef135e81f6b8107b9f16c777a3e8be9e3fbce541d18/tensorstore-0.1.84-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bd20de85c1b83dd3ca94db24e7bd449bdb055590a1b162b05691f6b81fff00f", size = 20986107, upload-time = "2026-05-16T06:17:56.571Z" }, ] -[[package]] -name = "threadpoolctl" -version = "3.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274, upload-time = "2025-03-13T13:49:23.031Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" }, -] - [[package]] name = "tqdm" version = "4.68.3" @@ -1913,6 +2246,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + [[package]] name = "tzdata" version = "2025.2" @@ -1973,21 +2318,32 @@ wheels = [ ] [[package]] -name = "xgboost" -version = "3.1.2" +name = "wandb" +version = "0.28.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, - { name = "nvidia-nccl-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, - { name = "scipy" }, + { name = "click" }, + { name = "gitpython" }, + { name = "packaging" }, + { name = "platformdirs" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sentry-sdk" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f0/64/42310363ecd814de5930981672d20da3d35271721ad2d2b4970b4092825b/xgboost-3.1.2.tar.gz", hash = "sha256:0f94496db277f5c227755e1f3ec775c59bafae38f58c94aa97c5198027c93df5", size = 1237438, upload-time = "2025-11-20T18:33:29.614Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/35/1e/efdd603db8cb37422b01d925f9cce1baaac46508661c73f6aafd5b9d7c51/xgboost-3.1.2-py3-none-macosx_10_15_x86_64.whl", hash = "sha256:b44f6ee43a28b998e289ab05285bd65a65d7999c78cf60b215e523d23dc94c5d", size = 2377854, upload-time = "2025-11-20T18:06:21.217Z" }, - { url = "https://files.pythonhosted.org/packages/f4/c6/ed928cb106f56ab73b3f4edb5287c1352251eb9225b5932d3dd5e2803f60/xgboost-3.1.2-py3-none-macosx_12_0_arm64.whl", hash = "sha256:09690f7430504fcd3b3e62bf826bb1282bb49873b68b07120d2696ab5638df41", size = 2211078, upload-time = "2025-11-20T18:06:47.063Z" }, - { url = "https://files.pythonhosted.org/packages/70/2f/5418f4b1deaf0886caf81c5e056299228ac2fc09b965a2dfe5e4496331c8/xgboost-3.1.2-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:f9b83f39340e5852bbf3e918318e7feb7a2a700ac7e8603f9bc3a06787f0d86b", size = 4953319, upload-time = "2025-11-20T18:28:29.851Z" }, - { url = "https://files.pythonhosted.org/packages/5e/ab/c60fcc137fa685533bb31e721de3ecc88959d393830d59d0204c5cbd2c85/xgboost-3.1.2-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:24879ac75c0ee21acae0101f791bc43303f072a86d70fdfc89dae10a0008767f", size = 115885060, upload-time = "2025-11-20T18:32:00.773Z" }, - { url = "https://files.pythonhosted.org/packages/30/7d/41847e45ff075f3636c95d1000e0b75189aed4f1ae18c36812575bb42b4b/xgboost-3.1.2-py3-none-win_amd64.whl", hash = "sha256:e627c50003269b4562aa611ed348dff8cb770e11a9f784b3888a43139a0f5073", size = 71979118, upload-time = "2025-11-20T18:27:55.23Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/5f/a7/683bfbd6cbade3012bc90d3e9c4cfc72dd62566195bf4c30321946d64b77/wandb-0.28.0.tar.gz", hash = "sha256:b20e5af0fe80e2e2a466b0466a1d60cedcc578dce0f036eca04f4a0adcad95b6", size = 40558332, upload-time = "2026-06-23T00:38:50.115Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c0/47/1723605f76c5d6446b6d0db65b83eda1599721bc8c1e65bd76cc1682b1a7/wandb-0.28.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:c3dab1205a5aca4abbad1eca08902cdba86add0edfa83d8d61b4429d0e79fa87", size = 24335272, upload-time = "2026-06-23T00:38:26.002Z" }, + { url = "https://files.pythonhosted.org/packages/81/ff/42b539bc75bc48fc86981dccde89327ba9b71504b805b9ba42cba7c26de9/wandb-0.28.0-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:ae255da18726ee8e731ef82cbc85035b901a28ae14cf91604c361b44b8d44ce0", size = 25557959, upload-time = "2026-06-23T00:38:28.993Z" }, + { url = "https://files.pythonhosted.org/packages/15/55/c3db03d04aeab3726066a418b2ef6a1f8119774ee510f4fbe992f52b7472/wandb-0.28.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:6dbcba12ab168aa37561f2f32dcdef8713495fc25fa7d30fdc9bfb37989694dd", size = 24878557, upload-time = "2026-06-23T00:38:31.417Z" }, + { url = "https://files.pythonhosted.org/packages/d8/5d/1385ce3c219cb5bd30d4027687e3f8d25969c7dfd09adad1cbd5080e1a72/wandb-0.28.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:325b2d0bd88be6eda5db10542499bad3710927f2569c81a84dc5eeaffc76825c", size = 26764727, upload-time = "2026-06-23T00:38:33.775Z" }, + { url = "https://files.pythonhosted.org/packages/00/58/23b6c17a6d3d5422b007707961c4496b2f6f892624d2910c9f7742fcc202/wandb-0.28.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:8954bc1c62ae43914dce2bebfd1d9957f72350f8fbb78e5cdfe2ca9b6be8a7b8", size = 25051656, upload-time = "2026-06-23T00:38:36.281Z" }, + { url = "https://files.pythonhosted.org/packages/89/67/9be00fb2db2281063af24a148636d2dd363d337317642ab5d8e93572c794/wandb-0.28.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:9fec6c908554c2dad33110c1312bc3028cc2e430f0679f16b84f82c8ea801e3b", size = 27074113, upload-time = "2026-06-23T00:38:38.737Z" }, + { url = "https://files.pythonhosted.org/packages/59/b1/f7a96c09cab0c5131b1e6466659b093b401e1653cbe6bb77b462fc1c361d/wandb-0.28.0-py3-none-win32.whl", hash = "sha256:8834ef3a7c8c43b701654162783caa7ad37af48a0ff06fc35d0d65a411f76ccd", size = 24525206, upload-time = "2026-06-23T00:38:42.041Z" }, + { url = "https://files.pythonhosted.org/packages/c6/c4/c7bed5e981679c74e9fbb22c03ff31c42e95f266199d03d8d325f4d0e6df/wandb-0.28.0-py3-none-win_amd64.whl", hash = "sha256:ac1f82292e2da4f98297b78c3a46726b3a6c5734ecb75fc39b8db2c8a4989159", size = 24525214, upload-time = "2026-06-23T00:38:44.549Z" }, + { url = "https://files.pythonhosted.org/packages/f0/77/b5ce9696c8cb955521a7941fbc443e78b2f504894c6ae1a2d0b1de6e12ae/wandb-0.28.0-py3-none-win_arm64.whl", hash = "sha256:c5b0faf1b84cf79ebabed77538c1940a4c6053e815f767a4004e877a1354bed1", size = 22378208, upload-time = "2026-06-23T00:38:47.148Z" }, ] [[package]]