Feat PyTorch calibrator and two-phase training pipeline by JemmaLDaniel · Pull Request #190 · instadeepai/winnow

JemmaLDaniel · 2026-04-14T08:27:29Z

Summary

Replaces the scikit-learn MLPClassifier calibrator with a PyTorch CalibratorNetwork, saved as model.safetensors + config.json. Adds a scalable two-phase train path (compute features to Parquet, then fit_from_features) and revises compute-features / predict / diagnose-calibration CLI orchestration to match the revisions branch pipeline.

What changed

Calibrator

PyTorch MLP with configurable hidden_dims, dropout, Adam training, GPU when available
Early stopping: n_iter_no_change + tol; optional validation subsampling (val_early_stopping_max_psms)
TrainingHistory (epoch metrics, JSON export, plot)
fit() / fit_from_features(); predict() expects pre-computed features (no internal compute_features)
Save/load via safetensors; feature configs in JSON (get_config()); Koina CE/frag runtime keys stripped on save
Predict-time merge of koina.input_constants / input_columns onto loaded models

Data & training

FeatureDataset + from_parquet() (file or directory, optional feature_columns)
winnow train: single-phase (features on full set → split → fit) or two-phase (features_path / val_features_path)
resolve_data_path() for local paths and HuggingFace Hub models/datasets

CLI & config

Batched compute-features / predict (per-experiment directories, shared filtering)
configs/koina.yaml extracted; slim predict.yaml / diagnose_calibration.yaml (no full train calibrator tree included where unnecessary)
irt_calibration predict overrides; hybrid iRT (skip/warn on bad pools instead of raising to keep training runs going, keep Fix per experiment irt regression #188 loaded-regressor pinning)
Sample Makefile targets use n_iter_no_change

Test plan

make tests / pre-commit
make train-sample and make predict-sample
Calibrator: save/load roundtrip, Koina override merge, early stopping, fit_from_features, empty-dataset guards
FeatureDataset: Parquet dir + feature_columns
iRT: per-experiment fit, loaded-regressor skip, low-diversity warning
test_koina_intensity_config, test_paths (HF path resolution)

…librator

…h models

…ing training history recording fix: convert ProbabilityCalibrator OmegaConf objects into plain Python types

…s-compatibility' into feat-pytorch-calibrator

…formatting from TrainingHistory.plot()

github-actions · 2026-04-14T08:29:31Z

Coverage Report

File	Stmts	Miss	Cover	Missing
__init__.py	0	0	100%
data_types.py	4	0	100%
calibration
__init__.py	0	0	100%
calibration_features.py	9	0	100%
calibrator.py	410	21	94%	56, 59, 74, 79–80, 198, 291–292, 294, 324–325, 520, 606–607, 610, 877–878, 1014, 1018, 1025, 1044
diagnostics.py	168	50	70%	70, 96, 101, 111, 115, 137, 146, 203–218, 261–262, 266, 307, 309–324, 335–341
calibration/features
__init__.py	10	0	100%
base.py	29	3	89%	83, 89–90
beam.py	47	0	100%
chimeric.py	82	1	98%	213
constants.py	9	0	100%
fragment_match.py	78	1	98%	203
mass_error.py	67	2	97%	16, 20
retention_time.py	206	17	91%	111, 114–115, 218, 229, 283–284, 338, 341, 398–400, 407, 410–411, 420, 487
sequence.py	19	0	100%
token_score.py	37	1	97%	82
utils.py	261	3	98%	96, 368, 594
compat
__init__.py	0	0	100%
instanovo.py	10	6	40%	12, 14–15, 17, 24–25
datasets
__init__.py	0	0	100%
calibration_dataset.py	109	17	84%	155, 169, 171, 173, 183, 196, 249, 251–252, 258–261, 263–266
feature_dataset.py	40	0	100%
interfaces.py	3	0	100%
psm_dataset.py	25	0	100%
datasets/data_loaders
__init__.py	5	0	100%
instanovo.py	119	19	84%	90, 93, 119, 142, 168–169, 172–174, 176–177, 179, 182–183, 185, 343–345, 356
mztab.py	215	55	74%	103, 106, 157, 161, 210–211, 223, 236–240, 287, 290, 302–303, 315–317, 319–320, 322, 324, 330, 334–336, 338–339, 343–346, 350, 514–515, 518, 521, 528, 542–546, 550–555, 561, 570–571, 573, 599
pointnovo.py	7	0	100%
utils.py	59	1	98%	11
winnow.py	39	4	89%	54–55, 91–92
fdr
__init__.py	0	0	100%
base.py	58	15	74%	81, 85–86, 91, 98–99, 105, 126, 129–130, 135, 137–138, 144, 186
database_grounded.py	28	1	96%	52
nonparametric.py	25	4	84%	62, 68–69, 72
scripts
__init__.py	0	0	100%
main.py	427	353	17%	23–26, 59, 61–62, 73, 78–79, 81–82, 84, 95, 100, 102–103, 105–106, 120, 124, 137–138, 147–148, 150, 167–168, 198, 200–202, 204, 206–211, 216–217, 224–226, 248, 251, 258–262, 265–266, 295–297, 299, 301, 306, 308–311, 313, 315–316, 318, 320–321, 323–324, 327, 332, 334, 341–343, 348–351, 353–354, 356, 358, 373–374, 376, 378–380, 382–387, 389, 403–404, 406–408, 410–411, 413–415, 417–418, 425–428, 442, 444, 446–447, 449, 457–459, 461, 465, 470, 484, 486, 488–489, 491, 499–501, 503–504, 506–510, 512, 535–536, 538–539, 541–543, 546, 550, 561–563, 565–568, 570–573, 588–592, 596, 607–608, 612–615, 619–623, 627–629, 632, 651, 672, 679, 686, 704–708, 710–713, 739–741, 743, 745, 750, 752–755, 757, 759, 761–762, 764–767, 769, 777–778, 780, 782, 784, 787–789, 791–793, 800, 817–819, 822, 825, 830, 832–836, 838, 840–841, 843–845, 847–848, 851–852, 857, 860, 865–867, 869–872, 875–876, 878, 880, 888–890, 892, 895–896, 899–900, 903–904, 907–908, 916–918, 922, 925, 929, 938–940, 942–943, 950–951, 953, 955, 960, 962–966, 968–969, 971–972, 975, 977–978, 980–983, 985–986, 988–989, 991–993, 999–1000, 1004, 1007, 1010–1014, 1017, 1024, 1029–1030, 1035–1037, 1040, 1045, 1055, 1062, 1064, 1068, 1070–1071, 1075–1076, 1115–1116, 1152–1153, 1194–1195, 1225–1226, 1256–1257, 1284–1285, 1315–1316, 1340–1341
utils
__init__.py	4	0	100%
config_formatter.py	53	40	24%	29, 37–38, 40–42, 44, 55, 58–60, 62–63, 66–69, 72–74, 77–78, 80, 91, 102, 113, 127–128, 130–132, 145–147, 150, 153–154, 157–158, 160
config_path.py	76	5	93%	24–26, 117–118
hydra_overrides.py	12	2	83%	12, 15
irt_calibration_config.py	72	72	0%	3, 5, 7, 9, 11–12, 15–19, 22, 27–34, 38–39, 42, 47–50, 52–55, 57, 59–60, 72–73, 76–77, 79–82, 85, 89–96, 99, 107–110, 117–118, 126, 133–136, 139, 141–142, 146, 148–149, 156, 158–159
koina_intensity_config.py	90	33	63%	42, 68, 76, 78, 108, 111, 131, 160–172, 177–178, 182–183, 189–192, 201–202, 204–205, 209
paths.py	20	0	100%
peptide.py	16	0	100%
rich_console.py	7	0	100%
TOTAL	2955	726	75%

Tests	Skipped	Failures	Errors	Time
480	0 💤	0 ❌	0 🔥	44.276s ⏱️

… selection

…ipelines

BioGeek · 2026-06-25T16:38:40Z

    "matchms>=0.31.0",
    "scikit-learn>=1.3.0",
    "matplotlib>=3.7.0",
+    "torch>=2.10.0",


I would recommend keeping the same torch version ranges as InstaNovo.

Now Winnow is at "torch>=2.10.0" and InstaNovo at <2.5/<2.!/<2.9 (depending on the paltform and CUDA version), which makes it impossible to install Winnow and InstaNovo in the same environment or Google Colab.

BioGeek · 2026-06-25T16:54:13Z

+```
+
+Koina model names, collision-energy / fragmentation inputs, and validity constraints
+live in [`configs/koina.yaml`](configs/koina.yaml) (composed via `defaults: - koina` in


Running make docs retuns:

uv run mkdocs build --strict │ ⚠ Warning from the Material for MkDocs team │ │ MkDocs 2.0, the underlying framework of Material for MkDocs, │ will introduce backward-incompatible changes, including: │ │ × All plugins will stop working – the plugin system has been removed │ × All theme overrides will break – the theming system has been rewritten │ × No migration path exists – existing projects cannot be upgraded │ × Closed contribution model – community members can't report bugs │ × Currently unlicensed – unsuitable for production use │ │ Our full analysis: │ │ https://squidfunk.github.io/mkdocs-material/blog/2026/02/18/mkdocs-2.0/ INFO - Cleaning site directory INFO - Building documentation to directory: /home/j-vangoey/code/winnow/docs_public WARNING - Doc file 'configuration.md' contains a link 'configs/koina.yaml', but the target is not found among documentation files. Aborted with 1 warnings in strict mode! make: *** [Makefile:122: docs] Error 1

MkDocs resolves relative links only within the documentation source tree, not arbitrary repo files. Since configs/koina.yaml is not present as a docs page/file under the configured docs directory, MkDocs reports a broken link. Because --strict is enabled, that warning aborts the build.

The Material/MkDocs 2.0 message is just an upstream warning banner, not the build-breaking error. But something we propably will have to look into.

Likely fix is to use a full GitHub URL.

Fixed with GitHub URL

BioGeek · 2026-06-25T16:58:34Z

+                    "dropped" if not self.learn_from_missing else "imputed as zero"
+                )
+                warnings.warn(
+                    f"Skipping experiment '{exp_name}':\n{e}\n"


When the input has no experiment_name and the global iRT regressor is skipped (for example because there are fewer than min_train_points or fewer than two unique peptides), _collect_training_data records "__global__" in _skipped_experiments, but this branch only applies skips when an experiment_name column exists. compute() then keeps those rows as valid and _apply_regressors() indexes self.irt_predictors["__global__"], raising KeyError instead of imputing or dropping the spectra as the warning says.

To reproduce: test_compute_imputes_skipped_global_regressor

This is fixed. When the global RT -> iRT regressor is skipped and there is no experiment_name column, _mark_missing_spectra now marks all spectra as missing if __global__ is in _skipped_experiments, so compute() imputes them (or drops them when learn_from_missing=False) instead of reaching _apply_regressors() and raising a KeyError. The test now passes.

BioGeek · 2026-06-25T17:02:13Z

+    if spectrum_path.is_dir():
+        return _compute_features_directory(


With data_loader=winnow, the valid input is a saved Winnow dataset directory containing metadata.csv and optionally predictions.pkl . This branch intercepts every directory and sends it to _discover_experiment_files, which only accepts .parquet , .ipc, or .mgf files, so predict/compute-features now fail before WinnowDatasetLoader.load() can read a previously saved Winnow dataset directory.

To reproduce, see: test_compute_features_loads_saved_winnow_dataset_directory

Fixed, _compute_features_batched_metadata now detects saved Winnow dataset directories (those containing metadata.csv) and passes them straight to _compute_features_single_file

BioGeek · 2026-06-25T17:04:00Z

+
+
+def notebook_safe_rich_handler(**kwargs: object) -> RichHandler:
+    """Return a Rich log handler that uses ANSI colours, not OSC-8 hyperlinks.


We might borrow this for InstaNovo as well!

Thanks! Something I did for the PAASTA practical

BioGeek · 2026-06-25T17:41:08Z

+    """
    irt_regressor_output_path = cfg.get("irt_regressor_output_path")
    if irt_regressor_output_path:
        from winnow.calibration.calibration_features import RetentionTimeFeature


A single calibrator/RetentionTimeFeature instance is reused per-file in directory mode, so _skipped_experiments accumulates. If file A skips an experiment and a later file B reuses that experiment_name, B's valid iRT features are silently marked is_missing_irt_error=True and imputed as zero (or dropped) → wrong calibrated scores, silently.

Should be fixed with a _skipped_experiments reset in prepare() now, and I've added a test to catch it in future: test_skipped_experiments_reset_between_prepare_calls

BioGeek · 2026-06-25T17:44:05Z

+            if should_stop:
+                break
+
+        history.epochs_trained = epoch + 1  # noqa: F821 — loop always runs


history.epochs_trained = epoch + 1 references the loop variable after the for loop. With max_epochs=0 (a valid, unvalidated int) the loop never runs, epoch is unbound and the result is an unhandled `NameError instead of a clear config error.

Added a catch and a test test_max_epochs_must_be_positive

BioGeek · 2026-06-25T17:45:41Z

+        )
+
+    @torch.no_grad()
+    def predict(self, dataset: CalibrationDataset) -> None:


Training uses batched DataLoaders but prediction materializes the entire feature tensor + all activations at once. On millions of PSMs this OOMs VRAM (CUDA OOM) or exhausts host RAM.

BioGeek · 2026-06-25T17:49:46Z

+    ) -> ProbabilityCalibrator:
+        """Load a calibrator from a local directory or HuggingFace Hub repo.
+
+        Expects the directory to contain ``model.safetensors`` and


load() no longer detects legacy pickle (calibrator.pkl) checkpoints. Pointing at a pre-PyTorch checkpoint dir now yields a bare FileNotFoundError. Raise a proper error that the old format is no longer supported.

Added an error to handle this and a test test_load_legacy_pickle_checkpoint_raises

BioGeek · 2026-06-25T17:52:58Z



 @app.command(
    name="diagnose-calibration",


diagnose_calibration loads via a single data_loader.load(**dataset_params) while predict streams directories one-at-a-time via _compute_features_batched_metadata. The same directory dataset that works under predict either errors or loads everything into RAM under diagnose-calibration.

BioGeek · 2026-06-25T17:56:05Z

+    n_val = max(1, int(n * validation_fraction))
+    n_train = n - n_val
+
+    generator = torch.Generator().manual_seed(cfg.get("calibrator", {}).get("seed", 42))


_maybe_split_validation uses torch.Generator().manual_seed(seed) while _maybe_split_calibration_dataset uses np.random.default_rng(seed). Same seed will result in different orderings, so single-phase vs two-phase training produce different train/val partitions and can't be reproduced across modes. Factor into one shared index-generating helper.

Refactored, the two modes should now agree for the same seed and dataset size

BioGeek · 2026-06-25T17:58:13Z

+    import numpy as np
+    from winnow.datasets.feature_dataset import FeatureDataset
+
+    features, labels = calibrator._extract_feature_matrix(train_data, labelled=True)


_fit_from_calibration_datasets duplicates the feature-extraction + FeatureDataset wrapping from fit(), reaching into the private _extract_feature_matrix. If that return contract changes, fit() picks it up but the CLI copy silently doesn't.

… fallback data

…iles

JemmaLDaniel added 19 commits April 11, 2026 16:04

Merge branch 'fix-per-experiment-irt-regression' into feat-pytorch-ca…

01e803e

…librator

feat: experiment-aware data loading for compute-features

3bd0e8e

feat: add option to compute feature matrix in compute-features

f92d76f

feat: add FeatureDataset, resolve_data_path and dependencies for torc…

26b3ebd

…h models

feat: rewrite ProbabilityCalibrator with PyTorch and safetensors, add…

9a8f364

…ing training history recording fix: convert ProbabilityCalibrator OmegaConf objects into plain Python types

feat: two-phase training pipeline with optional validation splitting

bd3d737

chore: update configs for PyTorch calibrator and two-phase training

79b2fc1

test: update and extend tests

ad071d3

chore: ignore saved safetensors

b0501c5

docs: update ProbabilityCalibrator and RetentionTimeFeature docs

551625a

chore: add docs and package build Make commands

b1083bd

Merge branch 'feat-prediction-column-remapping-for-instanovo-backward…

0576f6b

…s-compatibility' into feat-pytorch-calibrator

fix: make InstaNovoDatasetLoader residue_remapping optional

de6f79f

chore: use progress bar to log calibrator training progress

b704181

chore: replace mentions of 'Prosit' with 'Koina'

079f072

chore: standardise iRT metadata columns to lowercase snake case

967c825

feat: modify fit() to accept a CalibrationDataset directly and strip …

f3c4ca1

…formatting from TrainingHistory.plot()

docs: update calibrator training documentation

868342f

test: update calibrator training tests

bb6bf62

JemmaLDaniel self-assigned this Apr 14, 2026

JemmaLDaniel added the enhancement New feature or request label Apr 14, 2026

JemmaLDaniel mentioned this pull request Apr 14, 2026

feat: record training history #181

Closed

JemmaLDaniel changed the base branch from main to fix-per-experiment-irt-regression April 14, 2026 10:48

JemmaLDaniel changed the base branch from fix-per-experiment-irt-regression to main April 14, 2026 10:48

JemmaLDaniel added 5 commits June 22, 2026 17:04

Merge branch 'main' into feat-pytorch-calibrator

0a088d2

chore: update sample run commands to new calibrator arguments

405741d

feat: add Koina and iRT predict-time config helpers

2f145cd

chore: port PyTorch calibrator training API and FeatureDataset column…

fe3ee8b

… selection

chore: wire single-phase train and batched compute-features/predict p…

f99a903

…ipelines

JemmaLDaniel added 3 commits June 22, 2026 18:29

docs: document calibrator and config changes

13215ea

chore: fix sample train Makefile overrides

bbb0882

chore: add pickle and tex files to .gitignore

c6f2905

JemmaLDaniel changed the title ~~Feat pytorch calibrator~~ Feat PyTorch calibrator and two-phase training pipeline Jun 22, 2026

JemmaLDaniel marked this pull request as ready for review June 22, 2026 17:40

JemmaLDaniel requested a review from BioGeek June 22, 2026 17:40

BioGeek added 2 commits June 25, 2026 18:41

test: cover skipped global irt imputation

681679c

test: cover saved winnow directory loading

8b8834d

BioGeek requested changes Jun 25, 2026

View reviewed changes

BioGeek reviewed Jun 25, 2026

View reviewed changes

BioGeek requested changes Jun 25, 2026

View reviewed changes

JemmaLDaniel mentioned this pull request Jun 26, 2026

Feat use mgf for sample spectra #173

Merged

JemmaLDaniel added 9 commits June 27, 2026 17:20

Merge branch 'main' into feat-pytorch-calibrator

d01638d

fix: improve iRT regressor error handling for insufficient and global…

be21a68

… fallback data

Merge branch 'main' into feat-pytorch-calibrator

f70cb04

build: pin torch to InstaNovo-compatible versions

0cf750b

docs: use GitHub URL hyperlink

d32808f

fix: reset skipped experiments during iRT regressor fitting between f…

a0fba55

…iles

fix: catch max_epochs < 1

1540c38

fix: provide clear error for loading legacy pickle models

85800f2

fix: use same seed for single-phase and two-phase training

404177b

JemmaLDaniel mentioned this pull request Jul 1, 2026

Feat feature customisation #222

Open

JemmaLDaniel changed the base branch from main to feat-hf-feature-customisation July 1, 2026 18:39

BioGeek mentioned this pull request Jul 6, 2026

fix: The model in hugging face is "model.safetensors" #224

Open

		if spectrum_path.is_dir():
		return _compute_features_directory(



		def notebook_safe_rich_handler(**kwargs: object) -> RichHandler:
		"""Return a Rich log handler that uses ANSI colours, not OSC-8 hyperlinks.

Uh oh!

Conversation

JemmaLDaniel commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Calibrator

Data & training

CLI & config

Test plan

Uh oh!

github-actions Bot commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JemmaLDaniel Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JemmaLDaniel Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

JemmaLDaniel commented Apr 14, 2026 •

edited

Loading

github-actions Bot commented Apr 14, 2026 •

edited

Loading

JemmaLDaniel Jun 27, 2026 •

edited

Loading

JemmaLDaniel Jun 27, 2026 •

edited

Loading