diff --git a/.gitignore b/.gitignore index 6a577691..22b9d6cd 100644 --- a/.gitignore +++ b/.gitignore @@ -15,13 +15,14 @@ docs_public *.mztab *.fasta *.mgf -*.pkl +*.safetensors *.json *.yaml *.pdf *.png - *.ipynb +*.pkl +*.tex examples/winnow-general-model examples/winnow-ms-datasets diff --git a/Makefile b/Makefile index 41959a63..2fb59b38 100644 --- a/Makefile +++ b/Makefile @@ -86,7 +86,7 @@ install-all: ## Development commands # ################################################################################# -.PHONY: tests clean-coverage test-docker bash build-package clean-build clean-workspace test-build clean-all-build test-cli-isolated test-cli-config set-gcp-credentials +.PHONY: tests clean-coverage test-docker bash build-package clean-build test-build check-build docs docs-serve clean-docs set-gcp-credentials ## Run all tests tests: @@ -112,8 +112,22 @@ build-package: clean-build: rm -rf dist/ build/ *.egg-info/ winnow_fdr.egg-info/ -## Build the package and then clean up (safe test build) -test-build: build-package clean-build +## Build the package cleanly from scratch (fails on errors, cleans up on success) +check-build: clean-build build-package + @ls dist/*.whl dist/*.tar.gz >/dev/null 2>&1 && echo "Package build OK ✓" || (echo "Build produced no artifacts" && exit 1) + $(MAKE) clean-build + +## Build mkdocs site +docs: + uv run mkdocs build --strict + +## Serve mkdocs locally with live-reload +docs-serve: + uv run mkdocs serve + +## Remove mkdocs build output +clean-docs: + rm -rf docs_public/ ## Set the GCP credentials set-gcp-credentials: @@ -135,9 +149,9 @@ train-sample: dataset_output_path=results/calibrated_dataset.csv \ calibrator.features.retention_time_feature.train_fraction=0.3 \ calibrator.features.retention_time_feature.min_train_points=2 \ - calibrator.hidden_layer_sizes="[32, 16]" \ - calibrator.early_stopping=false \ - calibrator.max_iter=100 + calibrator.hidden_dims="[32, 16]" \ + calibrator.n_iter_no_change=100 \ + calibrator.max_epochs=100 ## Run winnow predict with sample data (uses locally trained model from models/new_model) predict-sample: diff --git a/docs/api/calibration.md b/docs/api/calibration.md index 5b0ffd47..47944b03 100644 --- a/docs/api/calibration.md +++ b/docs/api/calibration.md @@ -6,7 +6,7 @@ The `winnow.calibration` module implements confidence calibration for peptide-sp ### ProbabilityCalibrator -The main calibration model that transforms raw confidence scores into calibrated probabilities using a multi-layer perceptron classifier with various peptide and spectral features. +The main calibration model that transforms raw confidence scores into calibrated probabilities using a PyTorch neural network (`CalibratorNetwork`) with various peptide and spectral features. ```python from winnow.calibration import ProbabilityCalibrator @@ -31,20 +31,20 @@ residue_masses = { } # Create and configure calibrator -calibrator = ProbabilityCalibrator(seed=42) +calibrator = ProbabilityCalibrator(seed=42, hidden_dims=[128, 64]) # Add features for calibration calibrator.add_feature(MassErrorDaFeature(residue_masses=residue_masses)) calibrator.add_feature(FragmentMatchFeatures(mz_tolerance=0.02, mz_tolerance_unit="da")) calibrator.add_feature(BeamFeatures()) -# Train the calibrator -calibrator.fit(training_dataset) +# Train directly from a labelled CalibrationDataset +calibrator.fit(train_dataset) -# Make predictions +# Make predictions on new data calibrator.predict(test_dataset) -# Save/load trained models +# Save/load trained models (safetensors + config.json) ProbabilityCalibrator.save(calibrator, Path("calibrator_checkpoint")) # Load models - supports multiple sources @@ -60,18 +60,21 @@ loaded_calibrator = ProbabilityCalibrator.load("calibrator_checkpoint") **Key Features:** -- **Neural Network Classifier**: Uses MLPClassifier with standardised feature scaling +- **PyTorch Neural Network**: Uses a custom `CalibratorNetwork` (`nn.Module`) with feature normalisation - **Feature Management**: Add, remove and track multiple calibration features - **Dependency Handling**: Automatic computation of feature dependencies -- **Model Persistence**: Save and load trained calibrators -- **Feature Extraction**: Computes features and handles both labelled and unlabelled data +- **Model Persistence**: Save/load using `safetensors` (weights) and `config.json` (architecture, normalisation stats, feature definitions) +- **Two-phase Training**: Supports training from pre-computed Parquet feature matrices via `FeatureDataset.from_parquet()` and `fit_from_features()` +- **GPU Support**: Automatic GPU detection with CPU fallback **Main Methods:** - `add_feature(feature)`: Add a calibration feature -- `fit(dataset)`: Train the calibrator on a labelled dataset +- `compute_features(dataset)`: Run feature computation on a `CalibrationDataset`, mutating its metadata in place +- `fit(dataset, val_dataset)`: Compute features and train the calibrator from a `CalibrationDataset` +- `fit_from_features(dataset, val_dataset)`: Train from a pre-computed `FeatureDataset` (two-phase workflow) - `predict(dataset)`: Generate calibrated confidence scores -- `save(calibrator, path)`: Save trained model to disk +- `save(calibrator, path)`: Save trained model to disk (`model.safetensors` + `config.json`) - `load(pretrained_model_name_or_path, cache_dir)`: Load trained model from Hugging Face Hub or local directory - Default: Loads `"InstaDeepAI/winnow-general-model"` from Hugging Face @@ -94,9 +97,16 @@ The calibrator uses a feature-based approach where multiple feature extractors c 1. **Create Calibrator**: Initialise `ProbabilityCalibrator` 2. **Add Features**: Use `add_feature()` to include desired calibration features -3. **Fit Model**: Call `fit()` with labelled `CalibrationDataset` +3. **Fit Model**: Call `fit()` with a labelled `CalibrationDataset` — feature computation and training happen in one step 4. **Save Model**: Use `save()` to persist trained calibrator +For the two-phase workflow (compute features once, save to Parquet, train later): + +1. Call `compute_features(dataset)` to populate metadata columns +2. Export to Parquet via `dataset.to_parquet()` +3. Reload with `FeatureDataset.from_parquet()` +4. Train with `fit_from_features(train_ds, val_dataset=val_ds)` + ### Prediction workflow 1. **Load Calibrator**: Use `load()` to restore trained model from a Hugging Face repository or a local directory diff --git a/docs/cli.md b/docs/cli.md index 92b64539..f2996af4 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -78,14 +78,21 @@ winnow train dataset.spectrum_path_or_directory=data/spectra.parquet dataset.pre - `dataset.predictions_path`: Path to predictions file (set to `null` for Winnow format) - `model_output_dir`: Directory to save trained calibrator - `dataset_output_path`: Path to save training results CSV +- `features_path`: Optional path to pre-computed feature Parquet file or directory (enables two-phase training) +- `val_features_path`: Optional path to validation feature Parquet file or directory +- `validation_fraction`: Automatic validation split fraction (default: 0.1, used when `val_features_path` is null) +- `training_history_path`: Optional path to save epoch-level training history as JSON **Advanced calibrator configuration:** You can customise the calibrator architecture and features using nested parameters: ```bash -# Change MLP architecture -winnow train calibrator.hidden_layer_sizes=[100,50,25] +# Change network architecture +winnow train calibrator.hidden_dims=[128,64,32] + +# Adjust training hyperparameters +winnow train calibrator.learning_rate=0.01 calibrator.max_epochs=200 calibrator.n_iter_no_change=15 winnow train calibrator.features.fragment_match_features.mz_tolerance=0.01 \ calibrator.features.fragment_match_features.mz_tolerance_unit=da @@ -98,14 +105,18 @@ For comprehensive calibrator configuration options, see: ### `winnow compute-features` -Compute calibration features and write one enriched metadata CSV. Uses the same `data_loader`, `dataset` paths and `calibrator.features` stack as training, but does **not** fit the MLP or save a model. +Compute calibration features and write an enriched metadata CSV and optionally a lean numeric Parquet for model training. Uses the same `data_loader`, `dataset` paths and `calibrator.features` stack as training, but does **not** fit the calibrator or save a model. ```bash # Defaults (configs/compute_features.yaml) winnow compute-features -# Paths and output file -winnow compute-features dataset.spectrum_path_or_directory=data/spectra.parquet dataset.predictions_path=data/preds.csv dataset_output_path=results/features.csv +# Paths and output files +winnow compute-features \ + dataset.spectrum_path_or_directory=data/spectra.mgf \ + dataset.predictions_path=data/preds.csv \ + metadata_output_path=results/metadata.csv \ + training_matrix_output_path=results/features.parquet # De novo spectra (no ground truth): labelled=false; remove retention_time_feature if present winnow compute-features labelled=false '~calibrator.features.retention_time_feature' @@ -114,8 +125,8 @@ winnow compute-features labelled=false '~calibrator.features.retention_time_feat **Common parameters:** - `data_loader`, `dataset.*`: Same as `winnow train` -- `dataset_output_path`: Output CSV path -- `filter_empty_predictions`: Drop empty or invalid predictions (default: true) +- `metadata_output_path`: Full metadata CSV for EDA (all columns) +- `training_matrix_output_path`: Optional lean numeric Parquet containing only features and labels for model training (used with two-phase `features_path` workflow) - `labelled`: If true (default), spectrum data must include a `sequence` column; runs each feature's `prepare()` (e.g. iRT calibrator). Feature selection matches training: override `calibrator.features` or use `~calibrator.features.` to drop entries. See [Configuration guide](configuration.md#compute-features-configuration). @@ -273,12 +284,16 @@ winnow predict fdr_method=database_grounded fdr_control.fdr_threshold=0.05 Training produces: -1. **Model checkpoints** (`model_output_dir`): - - `calibrator.pkl`: Complete trained calibrator with all features and parameters +1. **Model directory** (`model_output_dir`): + - `model.safetensors`: Trained network weights in safetensors format + - `config.json`: Model architecture, feature definitions, normalisation statistics and training history 2. **Training results** (`dataset_output_path`): - CSV file with calibrated scores and evaluation metrics +3. **Optional** (`training_history_path`): + - JSON file with epoch-level training and validation metrics + ### Prediction output Prediction produces two CSV files in the `output-folder` directory: @@ -306,7 +321,8 @@ This separation allows users to work with metadata and features separately from ### Compute-features output -A single CSV at `dataset_output_path`: full metadata after feature computation. No calibrated scores, FDR columns or FDR row filtering (unlike `winnow predict`). +- **Metadata CSV** at `metadata_output_path`: Full metadata after feature computation, suitable for EDA. No calibrated scores, FDR columns or FDR row filtering (unlike `winnow predict`). +- **Training matrix Parquet** at `training_matrix_output_path` (optional): Lean numeric matrix containing only feature columns and labels, suitable for two-phase training via `features_path`. ## Example workflows @@ -360,6 +376,26 @@ winnow predict \ fdr_control.fdr_threshold=0.05 ``` +### Two-phase training (large-scale) + +For large datasets (e.g. millions of spectra across many projects), pre-compute features per project then train from the Parquet files: + +```bash +# Phase 1: Compute features per project +winnow compute-features \ + dataset.spectrum_path_or_directory=data/project_01/spectra.parquet \ + dataset.predictions_path=data/project_01/predictions.csv \ + training_matrix_output_path=features/project_01.parquet + +winnow compute-features \ + dataset.spectrum_path_or_directory=data/project_02/spectra.parquet \ + dataset.predictions_path=data/project_02/predictions.csv \ + training_matrix_output_path=features/project_02.parquet + +# Phase 2: Train from pre-computed features (directory of Parquets) +winnow train features_path=features/train/ val_features_path=features/val/ +``` + ### Advanced configuration ```bash @@ -367,10 +403,11 @@ winnow predict \ winnow train \ dataset.spectrum_path_or_directory=data/spectra.parquet \ dataset.predictions_path=data/predictions.csv \ - calibrator.hidden_layer_sizes=[100,50,25] \ - calibrator.learning_rate_init=0.01 \ - calibrator.max_iter=500 \ - calibrator.features.fragment_match_features.mz_tolerance=10 + calibrator.hidden_dims=[128,64,32] \ + calibrator.learning_rate=0.01 \ + calibrator.max_epochs=200 \ + calibrator.features.fragment_match_features.mz_tolerance=0.01 \ + calibrator.features.fragment_match_features.mz_tolerance_unit="da" # Predict with database-grounded FDR winnow predict \ @@ -385,7 +422,7 @@ winnow predict \ Winnow comes with sensible default settings for all parameters: -- **Calibrator**: 2-layer MLP with 50 hidden units per layer +- **Calibrator**: PyTorch feed-forward neural network - **Features**: Mass error, fragment match features, retention time deviation, chimeric features, beam features - **FDR**: Non-parametric method with 5% threshold - **Model**: Pretrained general model from Hugging Face diff --git a/docs/configuration.md b/docs/configuration.md index 6e0c0fa1..ab3ed442 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -40,8 +40,9 @@ configs/ │ ├── nonparametric.yaml │ └── database_grounded.yaml ├── train.yaml # Main training config -├── compute_features.yaml # Feature-only export (no MLP fit) +├── compute_features.yaml # Feature-only export (no calibrator fit) ├── calibrator.yaml # Model architecture and features +├── koina.yaml # Koina model names, inputs, and constraints ├── predict.yaml # Main prediction config └── diagnose_calibration.yaml # Tail calibration diagnostic (sTECE / TECE) ``` @@ -83,8 +84,8 @@ Access nested configuration values using dot notation: # Change calibrator seed winnow train calibrator.seed=123 -# Change MLP hidden layer sizes -winnow train calibrator.hidden_layer_sizes=[100,50,25] +# Change network architecture +winnow train calibrator.hidden_dims=[128,64,32] winnow train calibrator.features.fragment_match_features.mz_tolerance=0.01 \ calibrator.features.fragment_match_features.mz_tolerance_unit=da @@ -115,6 +116,13 @@ defaults: - calibrator - data_loader: instanovo # Options: instanovo, mztab, pointnovo, winnow +# Two-phase training: set features_path to skip raw data loading and train +# directly from pre-computed feature Parquets (produced by compute-features). +features_path: null +val_features_path: null +validation_fraction: 0.1 + +# Single-phase dataset config (ignored when features_path is set): dataset: # Path to the spectrum data file or to folder containing saved internal Winnow dataset spectrum_path_or_directory: data/spectra.mgf @@ -125,6 +133,8 @@ dataset: # Output paths model_output_dir: models/new_model dataset_output_path: results/calibrated_dataset.csv +irt_regressor_output_path: null +training_history_path: null ``` **Key parameters:** @@ -132,14 +142,19 @@ dataset_output_path: results/calibrated_dataset.csv - `data_loader`: Format of input data loader to use (via defaults: `instanovo`, `mztab`, `pointnovo`, `winnow`) - `dataset.spectrum_path_or_directory`: Path to spectrum/metadata file (InstaNovo: `.parquet`, `.ipc`, or `.mgf`; or directory for Winnow format) - `dataset.predictions_path`: Path to predictions file (set to null for Winnow format) -- `model_output_dir`: Where to save trained model +- `features_path`: Path to pre-computed feature Parquet(s) for two-phase training (file or directory) +- `val_features_path`: Explicit validation Parquet(s); overrides `validation_fraction` +- `validation_fraction`: Automatic random validation split fraction (default: 0.1) +- `model_output_dir`: Where to save trained model (`model.safetensors` + `config.json`) - `dataset_output_path`: Where to save calibrated training results +- `irt_regressor_output_path`: Optional path to save per-experiment iRT regressors +- `training_history_path`: Optional path to save epoch-level training history as JSON ## Compute-features configuration ### Main config (`configs/compute_features.yaml`) -Loads data like `train.yaml` (same `defaults`: `residues`, `calibrator`, `data_loader`), runs `ProbabilityCalibrator.compute_features` only, and writes one CSV. No `model_output_dir` and no MLP training. +Loads data like `train.yaml` (same `defaults`: `residues`, `calibrator`, `data_loader`), runs `ProbabilityCalibrator.compute_features` only, and writes outputs. No `model_output_dir` and no calibrator training. ```yaml defaults: @@ -152,16 +167,16 @@ dataset: spectrum_path_or_directory: data/spectra.mgf predictions_path: data/predictions.csv -dataset_output_path: results/metadata.csv -filter_empty_predictions: true +metadata_output_path: results/metadata.csv +# training_matrix_output_path: results/training_matrix.parquet labelled: true ``` **Key parameters:** - `dataset.*`, `data_loader`: Same meaning as in training config -- `dataset_output_path`: CSV path for metadata after feature computation -- `filter_empty_predictions`: If true, apply the same empty-prediction filter as train/predict +- `metadata_output_path`: Full metadata CSV for EDA +- `training_matrix_output_path`: Optional lean numeric Parquet for model training (used with two-phase `features_path` workflow) - `labelled`: If true, spectrum data must include `sequence` (ground truth). The feature set is the `calibrator.features` block from `calibrator.yaml` (shared with training). Override or drop features with Hydra the same way as for `winnow train`. @@ -174,13 +189,20 @@ Controls model architecture and calibration features: calibrator: _target_: winnow.calibration.calibrator.ProbabilityCalibrator - seed: 42 - hidden_layer_sizes: [50, 50] # The number of neurons in each hidden layer of the MLP classifier. - learning_rate_init: 0.001 # The initial learning rate for the MLP classifier. - alpha: 0.0001 # L2 regularisation parameter for the MLP classifier. - max_iter: 1000 # Maximum number of training iterations for the MLP classifier. - early_stopping: true # Whether to use early stopping to terminate training. - validation_fraction: 0.1 # Proportion of training data to use for early stopping validation. + # Network architecture + hidden_dims: [128, 64] # The number of neurons in each hidden layer of the network. + dropout: 0.1 # Dropout probability between hidden layers. + + # Training hyperparameters + learning_rate: 0.001 # Learning rate for the Adam optimiser. + weight_decay: 0.0001 # L2 regularisation (weight decay) parameter. + max_epochs: 100 # Maximum number of training epochs. + batch_size: 1024 # Mini-batch size for DataLoader. + n_iter_no_change: 10 # Early stopping: epochs without validation improvement by at least tol. + tol: 0.0001 # Minimum validation loss improvement to reset the early-stopping counter. + val_early_stopping_max_psms: null # Subsample large validation sets during early stopping only. + val_subsample_seed: null # RNG seed for validation subsampling (defaults to seed). + seed: 42 # Random seed for reproducibility. features: mass_error: @@ -202,7 +224,7 @@ calibrator: retention_time_feature: _target_: winnow.calibration.calibration_features.RetentionTimeFeature train_fraction: 0.1 # Top fraction of spectra (by confidence, descending) used to train the per-experiment RT->iRT regressor. - min_train_points: 10 # Minimum high-confidence spectra needed per experiment. Raises an error if fewer are available. + min_train_points: 10 # Minimum high-confidence spectra needed per experiment. Experiments below this are skipped with a warning. learn_from_missing: false # If True, impute missing features and add an indicator column. If False, filter invalid entries with a warning. seed: 42 # Random seed for reproducibility. irt_model_name: ${koina.irt_model} # The name of the Koina iRT model to use. @@ -214,7 +236,7 @@ calibrator: mz_tolerance: 20 mz_tolerance_unit: ppm learn_from_missing: false # If True, impute missing features and add an indicator column. If False, filter invalid entries with a warning. - prosit_intensity_model_name: ${koina.intensity_model} # The name of the Koina intensity model to use. + intensity_model_name: ${koina.intensity_model} # The name of the Koina intensity model to use. max_precursor_charge: ${koina.constraints.max_precursor_charge} # Maximum precursor charge accepted by the Koina intensity model. Applied to the runner-up sequence. max_peptide_length: ${koina.constraints.max_peptide_length} # Maximum peptide length accepted by the Koina intensity model. Applied to the runner-up (second-best) sequence. unsupported_residues: ${koina.constraints.unsupported_residues} # Residues unsupported by the configured Koina intensity model. @@ -223,51 +245,52 @@ calibrator: beam_features: _target_: winnow.calibration.calibration_features.BeamFeatures +``` + +Koina model names, collision-energy / fragmentation inputs, and validity constraints +live in [`winnow/configs/koina.yaml`](https://github.com/instadeepai/winnow/blob/main/winnow/configs/koina.yaml) (composed via `defaults: - koina` in +`calibrator.yaml` and inference configs). Feature blocks reference `${koina.*}` as shown above. + +**Key parameters:** + +- `seed`: Random seed for reproducibility +- `hidden_dims`: Architecture of the neural network (list of hidden layer sizes) +- `dropout`: Dropout probability between hidden layers +- `learning_rate`: Learning rate for the Adam optimiser +- `weight_decay`: L2 regularisation (weight decay) parameter +- `max_epochs`: Maximum number of training epochs +- `batch_size`: Mini-batch size for DataLoader +- `n_iter_no_change`: Early stopping — stop after this many epochs without validation loss improving by at least `tol` +- `tol`: Minimum validation loss improvement (absolute) to count as progress +- `val_early_stopping_max_psms`: When set, subsample validation PSMs for per-epoch early stopping (full validation metrics recorded after training) +- `val_subsample_seed`: RNG seed for validation subsampling (defaults to `seed`) +- `features.*`: Individual calibration feature configurations + +### Koina config (`configs/koina.yaml`) +Shared Koina settings for intensity- and iRT-based features. Composed automatically when +training (`calibrator.yaml` includes `defaults: - koina`) and explicitly for inference +(`predict.yaml`, `diagnose_calibration.yaml`). -# Koina model configuration — shared settings for all intensity- and iRT-based features. +```yaml koina: - # Model names intensity_model: Prosit_2025_intensity_22PTM irt_model: Prosit_2025_irt_22PTM - - # Model inputs — applied to FragmentMatchFeatures and ChimericFeatures. - # To use a constant value tiled across all rows, specify it under input_constants. - # To use per-row values from a metadata column, add the column mapping under input_columns. - # Each input key must appear in at most one of these two dicts. input_constants: collision_energies: 27 fragmentation_types: HCD input_columns: {} - - # Model constraints — adjust to match the capabilities of your Koina models. - # See docs/configuration.md for guidance on choosing these values. constraints: - max_precursor_charge: 6 # Maximum precursor charge accepted by the intensity models. - max_peptide_length: 30 # Maximum peptide length (residue token count) accepted by the intensity/iRT models. - # Residues unsupported by the configured Koina models. - # These residues must be specified using UNIMOD PTM IDs. - # Check that all PTMs unsupported by your selected Koina models but supported by Winnow are included. - unsupported_residues: - # Residue modifications (amino acid + modification) - - "N[UNIMOD:7]" # Deamidated asparagine - - "Q[UNIMOD:7]" # Deamidated glutamine - # ... - # N-terminal modifications (standalone tokens) - - "[UNIMOD:1]" # N-terminal acetylation - # ... + max_precursor_charge: 6 + max_peptide_length: 30 + unsupported_residues: [...] ``` **Key parameters:** -- `seed`: Random seed for reproducibility -- `hidden_layer_sizes`: Architecture of MLP classifier -- `learning_rate_init`: Initial learning rate -- `alpha`: L2 regularisation parameter -- `max_iter`: Maximum training iterations -- `early_stopping`: Whether to use early stopping -- `validation_fraction`: Proportion of data for validation -- `features.*`: Individual calibration feature configurations +- `intensity_model` / `irt_model`: Koina model identifiers (used when instantiating features at train time; saved in the checkpoint) +- `input_constants` / `input_columns`: Collision energy and fragmentation type — **required at predict time** (not persisted in the checkpoint). Override with e.g. `koina.input_constants.collision_energies=30` +- `constraints.*`: Validity filters interpolated into feature configs at train time ### Koina model input validation @@ -358,6 +381,7 @@ Controls dataset loading, FDR estimation and output: defaults: - _self_ - residues + - koina - data_loader: instanovo # Options: instanovo, mztab, pointnovo, winnow - fdr_method: nonparametric # Options: nonparametric, database_grounded @@ -369,12 +393,13 @@ dataset: predictions_path: data/predictions.csv calibrator: - # Path to the local calibrator directory or the Hugging Face model identifier - # If the path is a local directory path, it will be used directly - # If it is a Hugging Face repository identifier, it will be downloaded from Hugging Face pretrained_model_name_or_path: InstaDeepAI/winnow-general-model - # Directory to cache the Hugging Face model - cache_dir: null # can be set to null if using local model or for the default cache directory + cache_dir: null + irt_regressor_path: null + # Optional RT->iRT regressor fitting overrides (when irt_regressor_path is null): + # irt_calibration: + # train_fraction: 0.3 + # min_train_points: 10 fdr_control: # Target FDR threshold (e.g. 0.01 for 1%, 0.05 for 5% etc.) @@ -393,6 +418,9 @@ output_folder: results/predictions - `dataset.predictions_path`: Path to predictions file - `calibrator.pretrained_model_name_or_path`: Hugging Face model identifier or local model directory path - `calibrator.cache_dir`: Directory to cache Hugging Face models (null for default) +- `calibrator.irt_regressor_path`: Optional pre-fitted iRT regressors from training +- `calibrator.irt_calibration.*`: Optional RT→iRT regressor fitting overrides at predict time +- `koina.input_constants.*` / `koina.input_columns.*`: Koina collision energy and fragmentation (predict-time; match training overrides) - `fdr_method`: FDR estimation method (via defaults: `nonparametric` or `database_grounded`) - `fdr_control.fdr_threshold`: Target FDR threshold (e.g. 0.01 for 1%, 0.05 for 5%) - `fdr_control.confidence_column`: Column name with confidence scores @@ -437,6 +465,7 @@ Runs tail calibration diagnostics on a labelled holdout set: loads data like `wi defaults: - _self_ - residues + - koina - data_loader: instanovo dataset: @@ -465,6 +494,8 @@ diagnostics: **Key parameters:** - `data_loader`, `dataset.*`: Same meaning as in `predict.yaml` (holdout spectra + predictions) +- `koina.*`: Koina collision energy / fragmentation inputs (same as predict) +- `calibrator.pretrained_model_name_or_path`, `calibrator.cache_dir`: Model loading - `calibrator.pretrained_model_name_or_path`, `calibrator.cache_dir`: Calibrator checkpoint to score the holdout (same as predict) - `fdr_control.fdr_threshold`: Nominal FDR target $\alpha$ used to set $\tau$ via `NonParametricFDRControl.get_confidence_cutoff` - `fdr_control.confidence_column`: Column used for scores $S$ after calibration (default `calibrated_confidence`) @@ -726,6 +757,7 @@ Your custom config directory should mirror the structure of the package configs: ```text my_configs/ ├── residues.yaml # Override residue masses/modifications +├── koina.yaml # Override Koina models, inputs, constraints ├── calibrator.yaml # Override calibrator features ├── train.yaml # Override training config (if needed) ├── compute_features.yaml # Override compute-features config (if needed) @@ -771,7 +803,7 @@ When you use `--config-dir`, Winnow will: **What this means:** - ✅ **Partial configs at file level**: You only need to include the files you want to override (e.g., just `residues.yaml` and `calibrator.yaml`). Files not in your custom directory use package defaults. -- ❌ **Partial configs at key level don't work**: If you provide `calibrator.yaml` with only `seed: 999`, the other settings (`hidden_layer_sizes`, `features`, etc.) will be **missing**, not using package defaults. This will cause errors. +- ❌ **Partial configs at key level don't work**: If you provide `calibrator.yaml` with only `seed: 999`, the other settings (`hidden_dims`, `features`, etc.) will be **missing**, not using package defaults. This will cause errors. **Example - What happens with minimal config:** @@ -782,7 +814,7 @@ calibrator: seed: 99999 ``` -**Result**: Only `_target_` and `seed` are present. All other keys (`hidden_layer_sizes`, `learning_rate_init`, `features`, etc.) are **missing** from the final config. This will cause errors when running the pipeline in most cases. +**Result**: Only `_target_` and `seed` are present. All other keys (`hidden_dims`, `learning_rate`, `features`, etc.) are **missing** from the final config. This will cause errors when running the pipeline in most cases. **Example - What you need (complete structure):** @@ -791,12 +823,14 @@ calibrator: calibrator: _target_: winnow.calibration.calibrator.ProbabilityCalibrator seed: 99999 # Your custom value - hidden_layer_sizes: [50, 50] # Must include all settings - learning_rate_init: 0.001 - alpha: 0.0001 - max_iter: 1000 - early_stopping: true - validation_fraction: 0.1 + hidden_dims: [128, 64] # Must include all settings + dropout: 0.1 + learning_rate: 0.001 + weight_decay: 0.0001 + max_epochs: 100 + batch_size: 1024 + n_iter_no_change: 10 + tol: 0.0001 features: mass_error: _target_: winnow.calibration.calibration_features.MassErrorDaFeature diff --git a/pyproject.toml b/pyproject.toml index cff2556c..d5f22713 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,10 @@ dependencies = [ "matchms>=0.31.0", "scikit-learn>=1.3.0", "matplotlib>=3.7.0", + "torch>=2.6,<2.9; sys_platform == 'linux' or sys_platform == 'win32'", + "torch>=2.6,<2.8; sys_platform == 'darwin'", + "safetensors>=0.7.0", + "polars>=1.39.3", ] [project.urls] diff --git a/tests/calibration/features/test_retention_time.py b/tests/calibration/features/test_retention_time.py index 232457ff..c4564709 100644 --- a/tests/calibration/features/test_retention_time.py +++ b/tests/calibration/features/test_retention_time.py @@ -228,8 +228,8 @@ def test_prepare_refits_regressors_if_not_loaded_from_checkpoint(self, mock_koin assert reg_after_b.coef_[0] != coef_a @patch("winnow.calibration.features.retention_time.koinapy.Koina") - def test_prepare_raises_on_insufficient_data(self, mock_koina): - """Test that prepare raises ValueError when min_train_points is not met.""" + def test_prepare_skips_on_insufficient_data(self, mock_koina): + """Test that prepare skips experiments when min_train_points is not met.""" mock_model_instance = Mock() mock_koina.return_value = mock_model_instance mock_model_instance.model_inputs = ["peptide_sequences"] @@ -247,9 +247,145 @@ def test_prepare_raises_on_insufficient_data(self, mock_koina): feature = RetentionTimeFeature(train_fraction=0.1, min_train_points=10) - with pytest.raises(ValueError, match="insufficient data for iRT"): + with pytest.warns( + UserWarning, match="Skipping RT->iRT regressor fit for experiment 'exp_a'" + ): feature.prepare(dataset) + assert "exp_a" not in feature.irt_predictors + + @patch("winnow.calibration.features.retention_time.koinapy.Koina") + def test_skipped_experiments_reset_between_prepare_calls(self, mock_koina): + """Stale skipped experiments must not affect a later file with the same name.""" + mock_model_instance = Mock() + mock_koina.return_value = mock_model_instance + mock_model_instance.model_inputs = ["peptide_sequences"] + mock_model_instance.predict.side_effect = lambda inputs: pd.DataFrame( + {"irt": range(len(inputs))} + ) + + feature = RetentionTimeFeature( + train_fraction=1.0, + min_train_points=10, + learn_from_missing=True, + ) + + insufficient_metadata = pd.DataFrame( + { + "confidence": [0.95, 0.90], + "prediction": [["A", "G"], ["G", "A"]], + "retention_time": [10.5, 15.2], + "spectrum_id": [0, 1], + "experiment_name": ["exp_a", "exp_a"], + } + ) + insufficient_dataset = CalibrationDataset( + metadata=insufficient_metadata, predictions=None + ) + + with pytest.warns( + UserWarning, match="Skipping RT->iRT regressor fit for experiment 'exp_a'" + ): + feature.prepare(insufficient_dataset) + assert feature._skipped_experiments == ["exp_a"] + assert "exp_a" not in feature.irt_predictors + + feature.compute(insufficient_dataset) + assert insufficient_dataset.metadata["is_missing_irt_error"].tolist() == [ + True, + True, + ] + + repeated_peptides = [ + ["P", "E", "P", "T", "I", "D", "E"], + ["G", "L", "Y", "G", "A", "T"], + ["K", "V", "L", "V", "A", "P"], + ["A", "I", "V", "E", "G"], + ["S", "T", "D", "K"], + ] + sufficient_metadata = pd.DataFrame( + { + "confidence": [ + 0.95, + 0.93, + 0.91, + 0.89, + 0.87, + 0.85, + 0.83, + 0.81, + 0.79, + 0.77, + ], + "prediction": repeated_peptides + repeated_peptides, + "retention_time": [ + 12.5, + 15.2, + 18.1, + 21.0, + 24.3, + 26.8, + 29.1, + 31.5, + 34.0, + 36.2, + ], + "spectrum_id": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109], + "experiment_name": ["exp_a"] * 10, + } + ) + sufficient_dataset = CalibrationDataset( + metadata=sufficient_metadata, predictions=None + ) + + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + feature.prepare(sufficient_dataset) + assert feature._skipped_experiments == [] + assert "exp_a" in feature.irt_predictors + + feature.compute(sufficient_dataset) + assert not sufficient_dataset.metadata["is_missing_irt_error"].any() + assert sufficient_dataset.metadata["irt_error"].notna().all() + assert sufficient_dataset.metadata["predicted iRT"].notna().all() + + @patch("winnow.calibration.features.retention_time.koinapy.Koina") + def test_compute_imputes_skipped_global_regressor(self, mock_koina): + """Test global skipped iRT fits are imputed when experiment_name is absent.""" + mock_model_instance = Mock() + mock_koina.return_value = mock_model_instance + mock_model_instance.model_inputs = ["peptide_sequences"] + mock_model_instance.predict.return_value = pd.DataFrame({"irt": [10.0, 20.0]}) + + metadata = pd.DataFrame( + { + "confidence": [0.95, 0.90], + "prediction": [["A", "G"], ["G", "A"]], + "retention_time": [10.5, 15.2], + "spectrum_id": [0, 1], + } + ) + dataset = CalibrationDataset(metadata=metadata, predictions=None) + + feature = RetentionTimeFeature( + train_fraction=1.0, + min_train_points=10, + learn_from_missing=True, + ) + + with pytest.warns(UserWarning, match="Skipping global RT->iRT regressor fit"): + feature.prepare(dataset) + + assert "__global__" not in feature.irt_predictors + assert feature._skipped_experiments == ["__global__"] + + feature.compute(dataset) + + assert dataset.metadata["is_missing_irt_error"].tolist() == [True, True] + assert dataset.metadata["irt_error"].tolist() == [0.0, 0.0] + assert dataset.metadata["iRT"].isna().all() + assert dataset.metadata["predicted iRT"].isna().all() + @patch("winnow.calibration.features.retention_time.koinapy.Koina") def test_compute_with_mock( self, mock_koina, retention_time_feature, sample_dataset_with_rt @@ -439,10 +575,43 @@ def test_legacy_pickle_state_backfills_min_train_points(self, mock_koina): metadata = pd.DataFrame( { - "confidence": [0.95 - (idx * 0.01) for idx in range(10)], - "prediction": [["A", "G"] for _ in range(10)], - "retention_time": [float(idx) for idx in range(10)], - "spectrum_id": list(range(10)), + "confidence": [ + 0.95, + 0.93, + 0.91, + 0.89, + 0.87, + 0.85, + 0.83, + 0.81, + 0.79, + 0.77, + ], + "prediction": [ + ["P", "E", "P", "T", "I", "D", "E"], + ["G", "L", "Y", "G", "A", "T"], + ["K", "V", "L", "V", "A", "P"], + ["A", "I", "V", "E", "G"], + ["S", "T", "D", "K"], + ["L", "L", "G", "E"], + ["H", "G", "K", "T"], + ["Q", "F", "S", "R"], + ["M", "D", "P", "S"], + ["N", "F", "Y", "R"], + ], + "retention_time": [ + 12.5, + 15.2, + 18.1, + 21.0, + 24.3, + 26.8, + 29.1, + 31.5, + 34.0, + 36.2, + ], + "spectrum_id": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109], "experiment_name": ["exp_a"] * 10, } ) @@ -452,6 +621,69 @@ def test_legacy_pickle_state_backfills_min_train_points(self, mock_koina): assert "exp_a" in feature.irt_predictors + @patch("winnow.calibration.features.retention_time.koinapy.Koina") + def test_prepare_warns_on_low_peptide_diversity(self, mock_koina): + """Test warning when unique peptides are below min_train_points but PSM count is not.""" + mock_model_instance = Mock() + mock_koina.return_value = mock_model_instance + mock_model_instance.model_inputs = ["peptide_sequences"] + mock_model_instance.predict.return_value = pd.DataFrame({"irt": range(10)}) + + # Five distinct peptides, each observed twice (10 PSMs total). + repeated_peptides = [ + ["P", "E", "P", "T", "I", "D", "E"], + ["G", "L", "Y", "G", "A", "T"], + ["K", "V", "L", "V", "A", "P"], + ["A", "I", "V", "E", "G"], + ["S", "T", "D", "K"], + ] + + metadata = pd.DataFrame( + { + "confidence": [ + 0.95, + 0.93, + 0.91, + 0.89, + 0.87, + 0.85, + 0.83, + 0.81, + 0.79, + 0.77, + ], + "prediction": repeated_peptides + repeated_peptides, + "retention_time": [ + 12.5, + 15.2, + 18.1, + 21.0, + 24.3, + 26.8, + 29.1, + 31.5, + 34.0, + 36.2, + ], + "spectrum_id": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109], + "experiment_name": ["exp_a"] * 10, + } + ) + dataset = CalibrationDataset(metadata=metadata, predictions=None) + + feature = RetentionTimeFeature(train_fraction=1.0, min_train_points=10) + + with pytest.warns( + UserWarning, + match=( + r"Experiment 'exp_a': iRT calibration pool \(top 100%, 10 PSMs\):\n" + r" Only 5 unique peptide\(s\), below min_train_points=10\." + ), + ): + feature.prepare(dataset) + + assert "exp_a" in feature.irt_predictors + def test_save_and_load_regressors(self, tmp_path): """Test save_regressors and load_regressors round-trip.""" from sklearn.linear_model import LinearRegression @@ -463,7 +695,7 @@ def test_save_and_load_regressors(self, tmp_path): reg_b.fit([[3], [4]], [30, 40]) feature.irt_predictors = {"exp_a": reg_a, "exp_b": reg_b} - path = tmp_path / "regressors.pkl" + path = tmp_path / "regressors.safetensors" feature.save_regressors(path) assert path.exists() diff --git a/tests/calibration/test_calibrator.py b/tests/calibration/test_calibrator.py index 85cf9cc6..3c7172c6 100644 --- a/tests/calibration/test_calibrator.py +++ b/tests/calibration/test_calibrator.py @@ -1,15 +1,25 @@ """Unit tests for winnow ProbabilityCalibrator.""" -import pickle +import json + import numpy as np import pandas as pd import pytest -from winnow.calibration.calibrator import ProbabilityCalibrator +import torch + +from winnow.calibration.calibrator import ( + CalibratorNetwork, + ProbabilityCalibrator, + TrainingHistory, +) from winnow.calibration.calibration_features import ( CalibrationFeatures, FeatureDependency, ) from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.datasets.feature_dataset import FeatureDataset +from winnow.calibration.features.fragment_match import FragmentMatchFeatures +from winnow.utils.koina_intensity_config import KOINA_RUNTIME_CONFIG_KEYS class MockFeatureDependency(FeatureDependency): @@ -23,7 +33,6 @@ def name(self) -> str: return self._name def compute(self, dataset: CalibrationDataset): - # Add mock data to dataset - ensure it matches dataset length dataset.metadata[f"{self.name}_data"] = list(range(len(dataset.metadata))) @@ -51,11 +60,107 @@ def prepare(self, dataset): pass def compute(self, dataset): - # Add mock feature data for col in self.columns: dataset.metadata[col] = np.random.random(len(dataset.metadata)) +class MockKoinaFeature(MockCalibrationFeature): + """Feature with Koina-style model inputs for override tests.""" + + def __init__(self): + super().__init__("mock_koina", ["kcol"]) + self.model_input_constants = {"collision_energies": 20} + self.model_input_columns = {"fragmentation_types": "frag_col"} + + +class TestCalibratorNetwork: + """Test the CalibratorNetwork nn.Module.""" + + def test_forward_shape(self): + """Test that output has correct shape.""" + net = CalibratorNetwork(input_dim=5, hidden_dims=(16, 8)) + x = torch.randn(10, 5) + out = net(x) + assert out.shape == (10,) + + def test_single_sample(self): + """Test with a single sample.""" + net = CalibratorNetwork(input_dim=3, hidden_dims=(4,)) + x = torch.randn(1, 3) + out = net(x) + assert out.shape == (1,) + + def test_custom_dropout(self): + """Test that dropout parameter is accepted.""" + net = CalibratorNetwork(input_dim=3, hidden_dims=(8,), dropout=0.5) + x = torch.randn(5, 3) + net.eval() + out = net(x) + assert out.shape == (5,) + + +class TestTrainingHistory: + """Test TrainingHistory dataclass.""" + + def test_save_load_roundtrip(self, tmp_path): + """Test that save/load produces identical history.""" + history = TrainingHistory( + train_losses=[0.5, 0.4, 0.3], + val_losses=[0.6, 0.5, 0.4], + val_accuracies=[0.7, 0.8, 0.85], + best_epoch=2, + epochs_trained=3, + ) + path = tmp_path / "history.json" + history.save(path) + loaded = TrainingHistory.load(path) + + assert loaded.train_losses == history.train_losses + assert loaded.val_losses == history.val_losses + assert loaded.val_accuracies == history.val_accuracies + assert loaded.best_epoch == history.best_epoch + assert loaded.epochs_trained == history.epochs_trained + + def test_save_without_validation(self, tmp_path): + """Test save/load when no validation data was used.""" + history = TrainingHistory( + train_losses=[0.5, 0.4], + best_epoch=1, + epochs_trained=2, + ) + path = tmp_path / "history.json" + history.save(path) + loaded = TrainingHistory.load(path) + + assert loaded.val_losses is None + assert loaded.val_accuracies is None + + def test_json_format(self, tmp_path): + """Test that saved file is valid JSON with expected keys.""" + history = TrainingHistory(train_losses=[0.5], epochs_trained=1) + path = tmp_path / "history.json" + history.save(path) + + with open(path) as f: + data = json.load(f) + + assert "train_losses" in data + assert "epochs_trained" in data + + def test_plot_saves_file(self, tmp_path): + """Test that plot saves an image file.""" + history = TrainingHistory( + train_losses=[0.5, 0.4, 0.3], + val_losses=[0.6, 0.5, 0.4], + val_accuracies=[0.7, 0.8, 0.85], + best_epoch=2, + epochs_trained=3, + ) + plot_path = tmp_path / "plot.png" + history.plot(output_path=plot_path) + assert plot_path.exists() + + class TestProbabilityCalibrator: """Test the ProbabilityCalibrator class.""" @@ -72,14 +177,20 @@ def sample_dataset(self): ) return CalibrationDataset(metadata=metadata, predictions=[None] * 5) + @pytest.fixture() + def feature_dataset(self): + """Create a FeatureDataset for training tests.""" + n = 100 + np.random.seed(42) + features = np.random.randn(n, 3).astype(np.float32) + labels = np.random.choice([0.0, 1.0], n).astype(np.float32) + return FeatureDataset(features=features, labels=labels) + @pytest.fixture() def labelled_dataset(self): - """Create a dataset with labels for supervised learning.""" - # Create a larger dataset to support early stopping with validation_fraction=0.1 - # Need at least 40 samples to have 4+ validation samples for stratified split + """Create a labelled CalibrationDataset for compute_features tests.""" n_samples = 50 - np.random.seed(42) # For reproducible test data - + np.random.seed(42) metadata = pd.DataFrame( { "confidence": np.random.uniform(0.1, 0.99, n_samples), @@ -94,15 +205,66 @@ def test_initialization(self, calibrator): """Test ProbabilityCalibrator initialisation.""" assert isinstance(calibrator.feature_dict, dict) assert isinstance(calibrator.dependencies, dict) - assert isinstance(calibrator.dependency_reference_counter, dict) assert len(calibrator.feature_dict) == 0 - assert len(calibrator.dependencies) == 0 - - def test_initialization_with_custom_seed(self): - """Test calibrator initialisation with custom seed.""" - calibrator = ProbabilityCalibrator(seed=123) - # The seed should be passed to the MLPClassifier - assert calibrator.classifier.random_state == 123 + assert calibrator.network is None + + def test_initialization_with_params(self): + """Test calibrator initialisation with custom parameters.""" + calibrator = ProbabilityCalibrator( + hidden_dims=(64, 32), + dropout=0.2, + learning_rate=0.01, + seed=123, + ) + assert calibrator.hidden_dims == (64, 32) + assert calibrator.dropout == 0.2 + assert calibrator.learning_rate == 0.01 + assert calibrator.seed == 123 + assert calibrator.val_early_stopping_max_psms == 10000 + assert calibrator.val_subsample_seed is None + + def test_apply_koina_model_input_overrides(self): + """Inference-time Koina constant/column overrides merge into features.""" + calibrator = ProbabilityCalibrator() + calibrator.add_feature(MockKoinaFeature()) + calibrator.apply_koina_model_input_overrides( + model_input_constants={"collision_energies": 30}, + ) + feat = calibrator.feature_dict["mock_koina"] + assert feat.model_input_constants == { + "collision_energies": 30, + } + assert feat.model_input_columns == {"fragmentation_types": "frag_col"} + + def test_constants_win_over_default_columns_in_composed_cfg(self): + """Constants from CLI must not be undone by default column entries in cfg.""" + calibrator = ProbabilityCalibrator() + calibrator.add_feature( + FragmentMatchFeatures( + mz_tolerance=0.02, + mz_tolerance_unit="da", + model_input_columns={ + "collision_energies": "collision_energy", + "fragmentation_types": "frag_type", + }, + ) + ) + calibrator.apply_koina_model_input_overrides( + model_input_constants={ + "collision_energies": 27, + "fragmentation_types": "HCD", + }, + model_input_columns={ + "collision_energies": "collision_energy", + "fragmentation_types": "frag_type", + }, + ) + feat = calibrator.feature_dict["Fragment Match Features"] + assert feat.model_input_constants == { + "collision_energies": 27, + "fragmentation_types": "HCD", + } + assert feat.model_input_columns is None def test_columns_property_empty(self, calibrator): """Test columns property when no features are added.""" @@ -134,9 +296,9 @@ def test_predict_rejects_feature_column_shift( with pytest.raises(ValueError, match="feature dimension mismatch"): calibrator.predict(sample_dataset) - def test_features_property_empty(self, calibrator): - """Test features property when no features are added.""" - assert calibrator.features == [] + def test_feature_names_empty(self, calibrator): + """Test feature_names property when no features are added.""" + assert calibrator.feature_names == [] def test_add_feature_basic(self, calibrator): """Test adding a basic feature without dependencies.""" @@ -144,18 +306,15 @@ def test_add_feature_basic(self, calibrator): calibrator.add_feature(feature) assert "test_feature" in calibrator.feature_dict - assert calibrator.feature_dict["test_feature"] == feature assert calibrator.columns == ["col1", "col2"] - assert calibrator.features == ["test_feature"] + assert calibrator.feature_names == ["test_feature"] def test_add_feature_with_dependencies(self, calibrator): """Test adding a feature with dependencies.""" dependency = MockFeatureDependency("test_dep") feature = MockCalibrationFeature("test_feature", ["col1"], [dependency]) - calibrator.add_feature(feature) - assert "test_feature" in calibrator.feature_dict assert "test_dep" in calibrator.dependencies assert calibrator.dependency_reference_counter["test_dep"] == 1 @@ -169,7 +328,7 @@ def test_add_multiple_features_shared_dependency(self, calibrator): calibrator.add_feature(feature2) assert calibrator.dependency_reference_counter["shared_dep"] == 2 - assert len(calibrator.dependencies) == 1 # Only one dependency instance + assert len(calibrator.dependencies) == 1 def test_add_duplicate_feature_raises_error(self, calibrator): """Test that adding a duplicate feature raises KeyError.""" @@ -183,9 +342,6 @@ def test_remove_feature_basic(self, calibrator): """Test removing a feature without dependencies.""" feature = MockCalibrationFeature("removable_feature") calibrator.add_feature(feature) - - assert "removable_feature" in calibrator.feature_dict - calibrator.remove_feature("removable_feature") assert "removable_feature" not in calibrator.feature_dict @@ -194,8 +350,10 @@ def test_remove_feature_basic(self, calibrator): def test_remove_feature_with_dependencies(self, calibrator): """Test removing a feature with dependencies.""" dependency = MockFeatureDependency("removable_dep") - feature = MockCalibrationFeature("removable_feature", dependencies=[dependency]) - + feature = MockCalibrationFeature( + "removable_feature", + dependencies=[dependency], + ) calibrator.add_feature(feature) calibrator.remove_feature("removable_feature") @@ -211,13 +369,11 @@ def test_remove_feature_shared_dependency(self, calibrator): calibrator.add_feature(feature1) calibrator.add_feature(feature2) - - # Remove one feature - dependency should remain calibrator.remove_feature("feature1") assert "feature1" not in calibrator.feature_dict assert "feature2" in calibrator.feature_dict - assert "shared_dep" in calibrator.dependencies # Still needed by feature2 + assert "shared_dep" in calibrator.dependencies assert calibrator.dependency_reference_counter["shared_dep"] == 1 def test_remove_nonexistent_feature_raises_error(self, calibrator): @@ -225,26 +381,37 @@ def test_remove_nonexistent_feature_raises_error(self, calibrator): with pytest.raises(KeyError): calibrator.remove_feature("nonexistent") - def test_compute_features_unlabelled(self, calibrator, sample_dataset): - """Test computing features for unlabelled dataset.""" + def test_compute_features_mutates_metadata(self, calibrator, sample_dataset): + """Test that compute_features adds feature columns to metadata.""" feature = MockCalibrationFeature("test_feature", ["test_col"]) calibrator.add_feature(feature) - features = calibrator.compute_features(sample_dataset, labelled=False) + calibrator.compute_features(sample_dataset) + + assert "test_col" in sample_dataset.metadata.columns + + def test_extract_feature_matrix_unlabelled(self, calibrator, sample_dataset): + """Test extracting unlabelled feature matrix after compute_features.""" + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + + calibrator.compute_features(sample_dataset) + features = calibrator._extract_feature_matrix(sample_dataset, labelled=False) assert isinstance(features, np.ndarray) assert features.shape[0] == len(sample_dataset.metadata) - assert features.shape[1] == 2 # confidence column + one feature column + assert features.shape[1] == 2 # confidence + one feature column - def test_compute_features_labelled(self, calibrator, labelled_dataset): - """Test computing features for labelled dataset.""" + def test_extract_feature_matrix_labelled(self, calibrator, labelled_dataset): + """Test extracting labelled feature matrix after compute_features.""" feature = MockCalibrationFeature("test_feature", ["test_col"]) calibrator.add_feature(feature) - features, labels = calibrator.compute_features(labelled_dataset, labelled=True) + calibrator.compute_features(labelled_dataset) + features, labels = calibrator._extract_feature_matrix( + labelled_dataset, labelled=True + ) - assert isinstance(features, np.ndarray) - assert isinstance(labels, np.ndarray) assert features.shape[0] == len(labelled_dataset.metadata) assert labels.shape[0] == len(labelled_dataset.metadata) @@ -254,98 +421,397 @@ def test_compute_features_with_dependencies(self, calibrator, sample_dataset): feature = MockCalibrationFeature("test_feature", dependencies=[dependency]) calibrator.add_feature(feature) - calibrator.compute_features(sample_dataset, labelled=False) + calibrator.compute_features(sample_dataset) - # Check that dependency was computed (adds column to dataset) assert f"{dependency.name}_data" in sample_dataset.metadata.columns - def test_fit(self, calibrator, labelled_dataset): - """Test fitting the calibrator.""" - feature = MockCalibrationFeature("test_feature", ["test_col"]) - calibrator.add_feature(feature) + def test_max_epochs_must_be_positive(self): + """max_epochs=0 is rejected at construction time.""" + with pytest.raises(ValueError, match="max_epochs must be at least 1"): + ProbabilityCalibrator(max_epochs=0) + + def test_fit_from_features_returns_history(self, feature_dataset): + """Test that fit_from_features returns a TrainingHistory.""" + calibrator = ProbabilityCalibrator( + max_epochs=3, + hidden_dims=(8,), + seed=42, + ) + history = calibrator.fit_from_features(feature_dataset) + + assert isinstance(history, TrainingHistory) + assert len(history.train_losses) == 3 + assert history.epochs_trained == 3 + assert calibrator.network is not None + + def test_fit_from_features_with_validation(self, feature_dataset): + """Test fit_from_features with an explicit validation dataset.""" + np.random.seed(123) + val_features = np.random.randn(20, 3).astype(np.float32) + val_labels = np.random.choice([0.0, 1.0], 20).astype(np.float32) + val_dataset = FeatureDataset(features=val_features, labels=val_labels) + + calibrator = ProbabilityCalibrator( + max_epochs=5, + hidden_dims=(8,), + n_iter_no_change=3, + seed=42, + ) + history = calibrator.fit_from_features(feature_dataset, val_dataset) - # Should not raise any exception - calibrator.fit(labelled_dataset) + assert history.val_losses is not None + assert history.val_accuracies is not None + assert len(history.val_losses) <= 5 + assert history.final_val_loss is None + assert history.final_val_accuracy is None - # Check that scaler and classifier were fitted (basic smoke test) - assert hasattr(calibrator.scaler, "mean_") # Scaler fitted - assert hasattr(calibrator.classifier, "classes_") # Classifier fitted + def test_fit_from_features_val_subsample_records_full_metrics( + self, feature_dataset + ): + """Large validation sets are subsampled for early stopping; full metrics logged.""" + np.random.seed(123) + n_val = 50 + val_features = np.random.randn(n_val, 3).astype(np.float32) + val_labels = np.random.choice([0.0, 1.0], n_val).astype(np.float32) + val_dataset = FeatureDataset(features=val_features, labels=val_labels) + + calibrator = ProbabilityCalibrator( + max_epochs=2, + hidden_dims=(8,), + n_iter_no_change=10, + seed=42, + val_early_stopping_max_psms=10, + val_subsample_seed=123, + ) + history = calibrator.fit_from_features(feature_dataset, val_dataset) + + assert history.final_val_loss is not None + assert history.final_val_accuracy is not None + assert len(history.val_losses) == 2 + + def test_fit_from_features_sets_normalization(self, feature_dataset): + """Test that fit_from_features computes feature normalization stats.""" + calibrator = ProbabilityCalibrator(max_epochs=1, hidden_dims=(4,)) + calibrator.fit_from_features(feature_dataset) + + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + assert calibrator.feature_mean.shape == (3,) + + def test_end_to_end_fit_predict(self): + """Test the full pipeline: fit -> compute_features (inference) -> predict.""" + n_train = 80 + np.random.seed(42) + train_metadata = pd.DataFrame( + { + "confidence": np.random.uniform(0.1, 0.99, n_train), + "correct": np.random.choice([0, 1], n_train), + } + ) + train_raw = CalibrationDataset( + metadata=train_metadata, + predictions=[None] * n_train, + ) - def test_predict_after_fit(self, calibrator, labelled_dataset, sample_dataset): - """Test prediction after fitting.""" - feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator = ProbabilityCalibrator( + max_epochs=3, + hidden_dims=(8,), + seed=42, + ) + feature = MockCalibrationFeature("mock_feat", ["mock_col"]) calibrator.add_feature(feature) - # Fit first - calibrator.fit(labelled_dataset) + calibrator.fit(train_raw) - # Then predict - calibrator.predict(sample_dataset) + n_pred = 10 + pred_metadata = pd.DataFrame( + { + "confidence": np.random.uniform(0.1, 0.99, n_pred), + } + ) + pred_raw = CalibrationDataset( + metadata=pred_metadata, + predictions=[None] * n_pred, + ) + calibrator.compute_features(pred_raw) + calibrator.predict(pred_raw) - # Check that calibrated confidence was added to dataset - assert "calibrated_confidence" in sample_dataset.metadata.columns + assert "calibrated_confidence" in pred_raw.metadata.columns + probs = pred_raw.metadata["calibrated_confidence"] + assert len(probs) == n_pred + assert all(0.0 <= p <= 1.0 for p in probs) - def test_predict_without_fit_should_fail(self, calibrator, sample_dataset): + def test_predict_without_fit_raises(self, calibrator, sample_dataset): """Test that prediction fails if calibrator hasn't been fitted.""" feature = MockCalibrationFeature("test_feature", ["test_col"]) calibrator.add_feature(feature) - # Should raise an exception since classifier isn't fitted - from sklearn.exceptions import NotFittedError - - with pytest.raises(NotFittedError): + with pytest.raises(RuntimeError, match="not been fitted or loaded"): calibrator.predict(sample_dataset) - def test_classifier_parameters(self): - """Test that classifier is initialised with correct parameters.""" - calibrator = ProbabilityCalibrator(seed=42) + def test_save_strips_koina_runtime_keys(self, tmp_path, feature_dataset): + """Saved config.json must not persist runtime-only Koina settings.""" + calibrator = ProbabilityCalibrator(max_epochs=2, hidden_dims=(8,), seed=42) + calibrator.add_feature( + FragmentMatchFeatures( + mz_tolerance=0.02, + mz_tolerance_unit="da", + model_input_constants={"collision_energies": 27}, + model_input_columns={"fragmentation_types": "frag_type"}, + ) + ) + calibrator.fit_from_features(feature_dataset) + ProbabilityCalibrator.save(calibrator, tmp_path / "koina_model") + + with open(tmp_path / "koina_model" / "config.json") as f: + config = json.load(f) + feature_config = config["features"]["Fragment Match Features"] + for key in KOINA_RUNTIME_CONFIG_KEYS: + assert key not in feature_config + assert feature_config["intensity_model_name"] == "Prosit_2020_intensity_HCD" + + def test_load_strips_legacy_koina_runtime_keys(self, tmp_path, feature_dataset): + """Loading old checkpoints ignores baked-in Koina input presets.""" + calibrator = ProbabilityCalibrator(max_epochs=2, hidden_dims=(8,), seed=42) + calibrator.add_feature( + FragmentMatchFeatures( + mz_tolerance=0.02, + mz_tolerance_unit="da", + model_input_columns={ + "collision_energies": "collision_energy", + "fragmentation_types": "frag_type", + }, + ) + ) + calibrator.fit_from_features(feature_dataset) + ProbabilityCalibrator.save(calibrator, tmp_path / "legacy_model") + + with open(tmp_path / "legacy_model" / "config.json") as f: + config = json.load(f) + feature_config = config["features"]["Fragment Match Features"] + feature_config["model_input_columns"] = { + "collision_energies": "collision_energy", + "fragmentation_types": "frag_type", + } + with open(tmp_path / "legacy_model" / "config.json", "w") as f: + json.dump(config, f, indent=2) + + loaded = ProbabilityCalibrator.load(tmp_path / "legacy_model") + feat = loaded.feature_dict["Fragment Match Features"] + assert feat.model_input_constants is None + assert feat.model_input_columns is None + + def test_save_load_roundtrip(self, tmp_path, feature_dataset): + """Test that save/load produces a working calibrator with correct config.""" + calibrator = ProbabilityCalibrator( + max_epochs=2, + hidden_dims=(8, 4), + seed=42, + ) + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + calibrator.fit_from_features(feature_dataset) - assert calibrator.classifier.random_state == 42 - assert calibrator.classifier.hidden_layer_sizes == (50, 50) - assert calibrator.classifier.learning_rate_init == 0.001 - assert calibrator.classifier.alpha == 0.0001 - assert calibrator.classifier.max_iter == 1000 + ProbabilityCalibrator.save(calibrator, tmp_path / "model") - def test_scaler_initialization(self, calibrator): - """Test that scaler is properly initialised.""" - from sklearn.preprocessing import StandardScaler + loaded = ProbabilityCalibrator.load(tmp_path / "model") - assert isinstance(calibrator.scaler, StandardScaler) + assert loaded.hidden_dims == calibrator.hidden_dims + assert loaded.network is not None + assert loaded.feature_mean is not None + + with open(tmp_path / "model" / "config.json") as f: + config = json.load(f) + assert "features" in config + assert "test_feature" in config["features"] + assert "_target_" in config["features"]["test_feature"] + + def test_save_load_weights_and_normalization_match(self, tmp_path, feature_dataset): + """Test that weights and normalization stats survive save/load exactly.""" + calibrator = ProbabilityCalibrator( + max_epochs=2, + hidden_dims=(8,), + seed=42, + ) + calibrator.fit_from_features(feature_dataset) + + ProbabilityCalibrator.save(calibrator, tmp_path / "model") + loaded = ProbabilityCalibrator.load(tmp_path / "model") + + calibrator.network.cpu().eval() + loaded.network.cpu().eval() + + torch.testing.assert_close( + calibrator.feature_mean.cpu(), + loaded.feature_mean.cpu(), + ) + torch.testing.assert_close( + calibrator.feature_std.cpu(), + loaded.feature_std.cpu(), + ) + + x = torch.randn(5, 3) + x_norm = (x - calibrator.feature_mean.cpu()) / calibrator.feature_std.cpu() + with torch.no_grad(): + original_out = calibrator.network(x_norm) + loaded_out = loaded.network(x_norm) + + torch.testing.assert_close(original_out, loaded_out) + + def test_save_load_then_predict(self, tmp_path): + """Test that a loaded calibrator can compute features and predict on new data.""" + n = 80 + np.random.seed(42) + train_metadata = pd.DataFrame( + { + "confidence": np.random.uniform(0.1, 0.99, n), + "correct": np.random.choice([0, 1], n), + } + ) + train_raw = CalibrationDataset( + metadata=train_metadata, + predictions=[None] * n, + ) + + calibrator = ProbabilityCalibrator( + max_epochs=3, + hidden_dims=(8,), + seed=42, + ) + feature = MockCalibrationFeature("feat", ["feat_col"]) + calibrator.add_feature(feature) + + calibrator.fit(train_raw) + + ProbabilityCalibrator.save(calibrator, tmp_path / "model") + loaded = ProbabilityCalibrator.load(tmp_path / "model") + + pred_metadata = pd.DataFrame({"confidence": [0.9, 0.5, 0.1]}) + pred_ds = CalibrationDataset( + metadata=pred_metadata, + predictions=[None] * 3, + ) + loaded.compute_features(pred_ds) + loaded.predict(pred_ds) + + assert "calibrated_confidence" in pred_ds.metadata.columns + probs = pred_ds.metadata["calibrated_confidence"] + assert all(0.0 <= p <= 1.0 for p in probs) + + def test_early_stopping_triggers(self): + """Test that training stops early when n_iter_no_change is exhausted.""" + np.random.seed(42) + + # Linearly separable training data: label = 1 when x > 0. + x_train = np.linspace(-2, 2, 200).reshape(-1, 1).astype(np.float32) + y_train = (x_train[:, 0] > 0).astype(np.float32) + train_ds = FeatureDataset(features=x_train, labels=y_train) + + # Validation with *inverted* labels: model overfits train, val loss rises. + x_val = np.linspace(-2, 2, 50).reshape(-1, 1).astype(np.float32) + y_val = (x_val[:, 0] <= 0).astype(np.float32) + val_ds = FeatureDataset(features=x_val, labels=y_val) + + calibrator = ProbabilityCalibrator( + max_epochs=50, + hidden_dims=(16,), + learning_rate=0.01, + n_iter_no_change=3, + tol=1e-4, + seed=42, + ) + history = calibrator.fit_from_features(train_ds, val_ds) + + assert history.epochs_trained < 50 + + def test_add_features_plural(self, calibrator): + """Test adding multiple features at once via add_features.""" + feat1 = MockCalibrationFeature("feat_a", ["col_a"]) + feat2 = MockCalibrationFeature("feat_b", ["col_b"]) + calibrator.add_features([feat1, feat2]) + + assert "feat_a" in calibrator.feature_dict + assert "feat_b" in calibrator.feature_dict + assert calibrator.columns == ["col_a", "col_b"] + + def test_save_unfitted_raises(self, tmp_path, calibrator): + """Test that saving an unfitted calibrator raises RuntimeError.""" + with pytest.raises(RuntimeError, match="unfitted"): + ProbabilityCalibrator.save(calibrator, tmp_path / "model") + + def test_load_missing_path_raises(self): + """Test that loading from a nonexistent path raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError): + ProbabilityCalibrator.load("/nonexistent/path/to/model") + + def test_load_legacy_pickle_checkpoint_raises(self, tmp_path): + """Pre-PyTorch calibrator.pkl directories raise a clear error.""" + legacy_dir = tmp_path / "legacy_checkpoint" + legacy_dir.mkdir() + (legacy_dir / "calibrator.pkl").write_bytes(b"fake pickle") + + with pytest.raises( + ValueError, + match="Legacy pickle checkpoint format is no longer supported", + ): + ProbabilityCalibrator.load(legacy_dir) def test_empty_dataset_handling(self, calibrator): """Test handling of empty datasets.""" - empty_metadata = pd.DataFrame({"confidence": []}) # Include confidence column + empty_metadata = pd.DataFrame({"confidence": []}) empty_dataset = CalibrationDataset(metadata=empty_metadata, predictions=[]) feature = MockCalibrationFeature("test_feature") calibrator.add_feature(feature) - # Should handle empty dataset gracefully - features = calibrator.compute_features(empty_dataset, labelled=False) + calibrator.compute_features(empty_dataset) + features = calibrator._extract_feature_matrix(empty_dataset, labelled=False) assert features.shape[0] == 0 - def test_load_legacy_checkpoint_mlpclassifier_raises_error(self, tmp_path): - """Test that loading a legacy checkpoint (MLPClassifier) raises a clear error.""" - # Create a legacy checkpoint: calibrator.pkl containing only MLPClassifier - legacy_path = tmp_path / "legacy_checkpoint" - legacy_path.mkdir(parents=True) - - from sklearn.neural_network import MLPClassifier - - # Save MLPClassifier directly (legacy format) - legacy_classifier = MLPClassifier(random_state=42) - pickle.dump(legacy_classifier, open(legacy_path / "calibrator.pkl", "wb")) - - # Attempt to load - should raise ValueError with helpful message - with pytest.raises(ValueError, match="Legacy checkpoint format detected"): - ProbabilityCalibrator.load(legacy_path) - - # Verify the error message contains helpful information - try: - ProbabilityCalibrator.load(legacy_path) - except ValueError as e: - error_msg = str(e) - assert "Legacy checkpoint format detected" in error_msg - assert "MLPClassifier" in error_msg - assert "cannot correctly infer the trained feature set" in error_msg - assert "retrain" in error_msg.lower() or "retraining" in error_msg.lower() + def test_predict_empty_dataset_raises(self): + """Predict on zero rows should fail with a clear error, not a torch dtype error.""" + n_train = 20 + train_metadata = pd.DataFrame( + { + "confidence": np.linspace(0.9, 0.5, n_train), + "correct": np.ones(n_train, dtype=int), + } + ) + train_raw = CalibrationDataset( + metadata=train_metadata, predictions=[None] * n_train + ) + calibrator = ProbabilityCalibrator(max_epochs=1, hidden_dims=(4,), seed=42) + calibrator.add_feature(MockCalibrationFeature("mock_feat", ["mock_col"])) + calibrator.fit(train_raw) + + empty_dataset = CalibrationDataset( + metadata=train_metadata.iloc[:0].copy(), predictions=[] + ) + calibrator.compute_features(empty_dataset) + + with pytest.raises(ValueError, match="empty dataset"): + calibrator.predict(empty_dataset) + + def test_get_config_on_mock_feature(self): + """Test that get_config returns expected keys.""" + feature = MockCalibrationFeature("test", ["col1", "col2"]) + config = feature.get_config() + + assert "_target_" in config + assert "MockCalibrationFeature" in config["_target_"] + assert config["name"] == "test" + + def test_get_config_converts_omegaconf_to_plain(self): + """Test that get_config resolves DictConfig/ListConfig to plain types.""" + from omegaconf import ListConfig + + feature = MockCalibrationFeature( + "test", + columns=ListConfig(["col1", "col2"]), + dependencies=ListConfig([]), + ) + config = feature.get_config() + + assert isinstance(config["columns"], list) + assert not type(config["columns"]).__module__.startswith("omegaconf") + json.dumps(config) diff --git a/tests/datasets/test_feature_dataset.py b/tests/datasets/test_feature_dataset.py new file mode 100644 index 00000000..0ee9297a --- /dev/null +++ b/tests/datasets/test_feature_dataset.py @@ -0,0 +1,150 @@ +"""Unit tests for winnow.datasets.feature_dataset.""" + +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader + +from winnow.datasets.feature_dataset import FeatureDataset + + +class TestFeatureDataset: + """Test the FeatureDataset class.""" + + def test_from_arrays(self): + """Test construction from numpy arrays.""" + features = np.random.randn(10, 3).astype(np.float32) + labels = np.random.choice([0.0, 1.0], 10).astype(np.float32) + ds = FeatureDataset(features=features, labels=labels) + + assert len(ds) == 10 + x, y = ds[0] + assert x.shape == (3,) + assert y.shape == () + + def test_tensors_are_float32(self): + """Test that features and labels are stored as float32 tensors.""" + ds = FeatureDataset( + features=np.ones((5, 2), dtype=np.float64), + labels=np.ones(5, dtype=np.int32), + ) + assert ds.features.dtype == torch.float32 + assert ds.labels.dtype == torch.float32 + + def test_mismatched_lengths_raises(self): + """Test that mismatched features/labels raises ValueError.""" + with pytest.raises(ValueError, match="same length"): + FeatureDataset( + features=np.zeros((10, 3)), + labels=np.zeros(5), + ) + + def test_dataloader_integration(self): + """Test that FeatureDataset works correctly with PyTorch DataLoader.""" + n = 25 + n_features = 4 + features = np.random.randn(n, n_features).astype(np.float32) + labels = np.random.choice([0.0, 1.0], n).astype(np.float32) + ds = FeatureDataset(features=features, labels=labels) + + loader = DataLoader(ds, batch_size=8, shuffle=True) + total_samples = 0 + for batch_x, batch_y in loader: + assert batch_x.shape[1] == n_features + assert batch_x.dtype == torch.float32 + assert batch_y.dtype == torch.float32 + total_samples += len(batch_x) + + assert total_samples == n + + def test_from_parquet_single_file(self, tmp_path): + """Test loading from a single Parquet file with correct values.""" + import polars as pl + + df = pl.DataFrame( + { + "feature_a": [1.0, 2.0, 3.0], + "feature_b": [4.0, 5.0, 6.0], + "correct": [1.0, 0.0, 1.0], + } + ) + path = tmp_path / "data.parquet" + df.write_parquet(path) + + ds = FeatureDataset.from_parquet(path) + assert len(ds) == 3 + x, y = ds[0] + assert x.shape == (2,) + torch.testing.assert_close(x, torch.tensor([1.0, 4.0])) + assert y.item() == 1.0 + + def test_from_parquet_directory_preserves_values(self, tmp_path): + """Test loading from a directory preserves values across files.""" + import polars as pl + + for i in range(3): + df = pl.DataFrame( + { + "feat": [float(i * 10 + j) for j in range(5)], + "correct": [1.0, 0.0, 1.0, 0.0, 1.0], + } + ) + df.write_parquet(tmp_path / f"part_{i}.parquet") + + ds = FeatureDataset.from_parquet(tmp_path) + assert len(ds) == 15 + + all_feats = ds.features[:, 0].tolist() + assert 0.0 in all_feats + assert 10.0 in all_feats + assert 20.0 in all_feats + + def test_from_parquet_missing_correct_raises(self, tmp_path): + """Test that missing 'correct' column raises ValueError.""" + import polars as pl + + df = pl.DataFrame({"feature_a": [1.0, 2.0]}) + path = tmp_path / "data.parquet" + df.write_parquet(path) + + with pytest.raises(ValueError, match="correct"): + FeatureDataset.from_parquet(path) + + def test_from_parquet_empty_dir_raises(self, tmp_path): + """Test that empty directory raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="No .parquet"): + FeatureDataset.from_parquet(tmp_path) + + def test_from_parquet_feature_columns(self, tmp_path): + """Test explicit feature column selection.""" + import polars as pl + + df = pl.DataFrame( + { + "confidence": [0.9, 0.8], + "feature_a": [1.0, 2.0], + "spectrum_id": [0, 1], + "correct": [1.0, 0.0], + } + ) + path = tmp_path / "data.parquet" + df.write_parquet(path) + + ds = FeatureDataset.from_parquet( + path, feature_columns=["confidence", "feature_a"] + ) + assert ds.features.shape == (2, 2) + torch.testing.assert_close( + ds.features[0], torch.tensor([0.9, 1.0], dtype=torch.float32) + ) + + def test_from_parquet_missing_feature_column_raises(self, tmp_path): + """Test that missing feature_columns raises ValueError.""" + import polars as pl + + df = pl.DataFrame({"feature_a": [1.0], "correct": [1.0]}) + path = tmp_path / "data.parquet" + df.write_parquet(path) + + with pytest.raises(ValueError, match="missing from Parquet"): + FeatureDataset.from_parquet(path, feature_columns=["confidence"]) diff --git a/tests/scripts/test_compute_features.py b/tests/scripts/test_compute_features.py new file mode 100644 index 00000000..ad8b80f1 --- /dev/null +++ b/tests/scripts/test_compute_features.py @@ -0,0 +1,54 @@ +"""Tests for compute-features helper paths.""" + +import pandas as pd + +from winnow.datasets.data_loaders import WinnowDatasetLoader +from winnow.scripts.main import _compute_features_batched_metadata + + +class RecordingCalibrator: + """Minimal calibrator that records feature computation calls.""" + + def __init__(self): + self.compute_features_calls = 0 + + def compute_features(self, dataset): + """Record the call and add a marker feature column.""" + self.compute_features_calls += 1 + dataset.metadata["marker_feature"] = 1.0 + + +def test_compute_features_loads_saved_winnow_dataset_directory(tmp_path): + """Saved Winnow dataset directories should be passed directly to the loader.""" + metadata = pd.DataFrame( + { + "prediction": ["AG", "MG"], + "confidence": [0.9, 0.8], + "mz_array": ["[100.0, 200.0]", "[150.0, 250.0]"], + "intensity_array": ["[1000.0, 2000.0]", "[1500.0, 2500.0]"], + } + ) + metadata.to_csv(tmp_path / "metadata.csv", index=False) + + loader = WinnowDatasetLoader( + residue_masses={ + "A": 71.037114, + "G": 57.021464, + "M": 131.040485, + }, + residue_remapping={}, + ) + calibrator = RecordingCalibrator() + + all_metadata = _compute_features_batched_metadata( + spectrum_path=tmp_path, + predictions_path=None, + data_loader=loader, + calibrator=calibrator, + labelled=False, + ) + + assert calibrator.compute_features_calls == 1 + assert len(all_metadata) == 1 + assert all_metadata[0]["prediction"].tolist() == [["A", "G"], ["M", "G"]] + assert all_metadata[0]["marker_feature"].tolist() == [1.0, 1.0] diff --git a/tests/utils/test_koina_intensity_config.py b/tests/utils/test_koina_intensity_config.py new file mode 100644 index 00000000..b47617ff --- /dev/null +++ b/tests/utils/test_koina_intensity_config.py @@ -0,0 +1,51 @@ +"""Tests for Koina intensity model input configuration helpers.""" + +import pytest +import typer + +from winnow.utils.koina_intensity_config import ( + resolve_feature_model_inputs, + strip_runtime_keys_from_feature_config, + validate_koina_intensity_config, +) + + +def test_resolve_feature_model_inputs_defaults_columns(): + """Null constants/columns resolve to default metadata column names.""" + constants, columns = resolve_feature_model_inputs( + {"collision_energies": None, "fragmentation_types": None}, + {}, + ) + assert constants is None + assert columns == { + "collision_energies": "collision_energy", + "fragmentation_types": "frag_type", + } + + +def test_strip_runtime_keys_from_feature_config(): + """Runtime-only Koina keys are removed from saved feature configs.""" + cfg = { + "_target_": "winnow.calibration.features.fragment_match.FragmentMatchFeatures", + "mz_tolerance": 0.02, + "model_input_constants": {"collision_energies": 27}, + "model_input_columns": {"fragmentation_types": "frag_type"}, + } + stripped = strip_runtime_keys_from_feature_config(cfg) + assert "model_input_constants" not in stripped + assert "model_input_columns" not in stripped + assert stripped["mz_tolerance"] == 0.02 + + +def test_validate_koina_intensity_config_conflicting_sources(): + """Dual CE/frag specification exits with code 1.""" + koina_cfg = { + "input_constants": {"collision_energies": 27}, + "input_columns": {"collision_energies": "collision_energy"}, + } + with pytest.raises(typer.Exit) as exc: + validate_koina_intensity_config( + koina_cfg, + hydra_overrides=["koina.input_columns.collision_energies=collision_energy"], + ) + assert exc.value.exit_code == 1 diff --git a/tests/utils/test_paths.py b/tests/utils/test_paths.py new file mode 100644 index 00000000..bbcf7ee5 --- /dev/null +++ b/tests/utils/test_paths.py @@ -0,0 +1,104 @@ +"""Unit tests for winnow.utils.paths.""" + +import logging +from unittest.mock import patch + +import pytest + +from winnow.utils.paths import resolve_data_path + + +class TestResolveDataPath: + """Test the resolve_data_path function.""" + + def test_local_file(self, tmp_path): + """Test resolution of an existing local file.""" + f = tmp_path / "data.parquet" + f.touch() + result = resolve_data_path(str(f)) + assert result == f.resolve() + + def test_local_directory(self, tmp_path): + """Test resolution of an existing local directory.""" + d = tmp_path / "model_dir" + d.mkdir() + result = resolve_data_path(str(d)) + assert result == d.resolve() + + def test_nonexistent_local_raises(self): + """Test that a nonexistent local path that is not an HF ID raises.""" + with pytest.raises(FileNotFoundError, match="Could not resolve"): + resolve_data_path("/definitely/does/not/exist/anywhere") + + def test_relative_local_path(self, tmp_path, monkeypatch): + """Test resolution of a relative local path.""" + monkeypatch.chdir(tmp_path) + (tmp_path / "local_data").mkdir() + result = resolve_data_path("local_data") + assert result == (tmp_path / "local_data").resolve() + + def test_hf_download_fallback(self, tmp_path): + """Test that HF snapshot_download is called when path doesn't exist locally.""" + fake_download_dir = tmp_path / "hf_cache" / "models--org--repo" + fake_download_dir.mkdir(parents=True) + + with patch( + "winnow.utils.paths.snapshot_download", + return_value=str(fake_download_dir), + ) as mock_dl: + result = resolve_data_path( + "org/repo", + repo_type="model", + cache_dir=tmp_path / "hf_cache", + ) + + mock_dl.assert_called_once_with( + repo_id="org/repo", + repo_type="model", + cache_dir=str(tmp_path / "hf_cache"), + ) + assert result == fake_download_dir + + def test_hf_download_failure_raises_with_context(self): + """Test that HF download failure wraps the original error.""" + with patch( + "winnow.utils.paths.snapshot_download", + side_effect=Exception("401 Unauthorized"), + ): + with pytest.raises(FileNotFoundError, match="401 Unauthorized"): + resolve_data_path("org/nonexistent-model") + + def test_local_path_resembling_hf_warns(self, tmp_path, caplog): + """Test warning when a local path looks like an HF repo ID.""" + hf_like = tmp_path / "org" / "model" + hf_like.mkdir(parents=True) + + with caplog.at_level(logging.WARNING): + result = resolve_data_path(str(hf_like)) + + assert result == hf_like.resolve() + + def test_local_relative_hf_like_path_warns(self, tmp_path, monkeypatch, caplog): + """Test warning when a relative path like 'org/repo' exists locally.""" + monkeypatch.chdir(tmp_path) + (tmp_path / "myorg" / "mymodel").mkdir(parents=True) + + with caplog.at_level(logging.WARNING): + result = resolve_data_path("myorg/mymodel") + + assert result == (tmp_path / "myorg" / "mymodel").resolve() + assert any("HF repo ID" in rec.message for rec in caplog.records) + + def test_cache_dir_none_passes_none(self, tmp_path): + """Test that cache_dir=None passes None to snapshot_download.""" + with patch( + "winnow.utils.paths.snapshot_download", + return_value=str(tmp_path), + ) as mock_dl: + resolve_data_path("org/repo", cache_dir=None) + + mock_dl.assert_called_once_with( + repo_id="org/repo", + repo_type="model", + cache_dir=None, + ) diff --git a/uv.lock b/uv.lock index f56fce74..d1e17e76 100644 --- a/uv.lock +++ b/uv.lock @@ -4,13 +4,23 @@ requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version > '3.11' and python_full_version < '3.12'", - "python_full_version == '3.11'", - "python_full_version < '3.11'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] [[package]] @@ -33,7 +43,9 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch" }, + { name = "torch", version = "2.7.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "torch", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" } wheels = [ @@ -707,7 +719,9 @@ name = "contourpy" version = "1.3.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.11'", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "numpy", marker = "python_full_version < '3.11'" }, @@ -779,12 +793,20 @@ source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version > '3.11' and python_full_version < '3.12'", - "python_full_version == '3.11'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "numpy", marker = "python_full_version >= '3.11'" }, @@ -982,31 +1004,6 @@ toml = [ { name = "tomli", marker = "python_full_version <= '3.11'" }, ] -[[package]] -name = "cuda-bindings" -version = "12.9.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/d8/b546104b8da3f562c1ff8ab36d130c8fe1dd6a045ced80b4f6ad74f7d4e1/cuda_bindings-12.9.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d3c842c2a4303b2a580fe955018e31aea30278be19795ae05226235268032e5", size = 12148218, upload-time = "2025-10-21T14:51:28.855Z" }, - { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" }, - { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, - { url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" }, - { url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" }, - { url = "https://files.pythonhosted.org/packages/d1/af/6dfd8f2ed90b1d4719bc053ff8940e494640fe4212dc3dd72f383e4992da/cuda_bindings-12.9.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8b72ee72a9cc1b531db31eebaaee5c69a8ec3500e32c6933f2d3b15297b53686", size = 11922703, upload-time = "2025-10-21T14:52:03.585Z" }, - { url = "https://files.pythonhosted.org/packages/6c/19/90ac264acc00f6df8a49378eedec9fd2db3061bf9263bf9f39fd3d8377c3/cuda_bindings-12.9.4-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d80bffc357df9988dca279734bc9674c3934a654cab10cadeed27ce17d8635ee", size = 11924658, upload-time = "2025-10-21T14:52:10.411Z" }, -] - -[[package]] -name = "cuda-pathfinder" -version = "1.5.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/93/66/0c02bd330e7d976f83fa68583d6198d76f23581bcbb5c0e98a6148f326e5/cuda_pathfinder-1.5.0-py3-none-any.whl", hash = "sha256:498f90a9e9de36044a7924742aecce11c50c49f735f1bc53e05aa46de9ea4110", size = 49739, upload-time = "2026-03-24T21:14:30.869Z" }, -] - [[package]] name = "cycler" version = "0.12.1" @@ -1693,7 +1690,9 @@ name = "ipython" version = "8.39.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.11'", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, @@ -1718,8 +1717,12 @@ name = "ipython" version = "9.10.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version > '3.11' and python_full_version < '3.12'", - "python_full_version == '3.11'", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, @@ -1746,10 +1749,14 @@ source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, @@ -1815,7 +1822,9 @@ name = "jaxtyping" version = "0.3.7" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.11'", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "wadler-lindig", marker = "python_full_version < '3.11'" }, @@ -1832,12 +1841,20 @@ source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version > '3.11' and python_full_version < '3.12'", - "python_full_version == '3.11'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "wadler-lindig", marker = "python_full_version >= '3.11'" }, @@ -2084,7 +2101,7 @@ dependencies = [ { name = "overrides", marker = "python_full_version < '3.12'" }, { name = "packaging" }, { name = "prometheus-client" }, - { name = "pywinpty", marker = "os_name == 'nt'" }, + { name = "pywinpty", marker = "(python_full_version < '3.12' and os_name == 'nt' and sys_platform == 'linux') or (os_name == 'nt' and sys_platform != 'darwin' and sys_platform != 'linux')" }, { name = "pyzmq" }, { name = "send2trash" }, { name = "terminado" }, @@ -2102,7 +2119,7 @@ name = "jupyter-server-terminals" version = "0.5.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pywinpty", marker = "os_name == 'nt'" }, + { name = "pywinpty", marker = "(python_full_version < '3.12' and os_name == 'nt' and sys_platform == 'linux') or (os_name == 'nt' and sys_platform != 'darwin' and sys_platform != 'linux')" }, { name = "terminado" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f4/a7/bcd0a9b0cbba88986fe944aaaf91bfda603e5a50bda8ed15123f381a3b2f/jupyter_server_terminals-0.5.4.tar.gz", hash = "sha256:bbda128ed41d0be9020349f9f1f2a4ab9952a73ed5f5ac9f1419794761fb87f5", size = 31770, upload-time = "2026-01-14T16:53:20.213Z" } @@ -3306,7 +3323,7 @@ name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, @@ -3317,7 +3334,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, @@ -3344,9 +3361,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-cusparse-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, @@ -3357,7 +3374,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'emscripten') or (python_full_version < '3.12' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, @@ -3373,10 +3390,10 @@ wheels = [ [[package]] name = "nvidia-nccl-cu12" -version = "2.27.5" +version = "2.27.3" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, + { url = "https://files.pythonhosted.org/packages/5c/5b/4e4fff7bad39adf89f735f2bc87248c81db71205b62bcc0d5ca5b606b3c3/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039", size = 322364134, upload-time = "2025-06-03T21:58:04.013Z" }, ] [[package]] @@ -3387,14 +3404,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, ] -[[package]] -name = "nvidia-nvshmem-cu12" -version = "3.4.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, -] - [[package]] name = "nvidia-nvtx-cu12" version = "12.8.90" @@ -4773,7 +4782,9 @@ name = "scikit-learn" version = "1.7.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.11'", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "joblib", marker = "python_full_version < '3.11'" }, @@ -4822,12 +4833,20 @@ source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version > '3.11' and python_full_version < '3.12'", - "python_full_version == '3.11'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "joblib", marker = "python_full_version >= '3.11'" }, @@ -5181,7 +5200,7 @@ version = "0.18.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ptyprocess", marker = "os_name != 'nt'" }, - { name = "pywinpty", marker = "os_name == 'nt'" }, + { name = "pywinpty", marker = "(python_full_version < '3.12' and os_name == 'nt' and sys_platform == 'linux') or (os_name == 'nt' and sys_platform != 'darwin' and sys_platform != 'linux')" }, { name = "tornado" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8a/11/965c6fd8e5cc254f1fe142d547387da17a8ebfd75a3455f637c663fb38a0/terminado-0.18.1.tar.gz", hash = "sha256:de09f2c4b85de4765f7714688fff57d3e75bad1f909b589fde880460c753fd2e", size = 32701, upload-time = "2024-03-12T14:34:39.026Z" } @@ -5266,14 +5285,50 @@ wheels = [ [[package]] name = "torch" -version = "2.10.0" +version = "2.7.1" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", + "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform == 'darwin'", +] dependencies = [ - { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx" }, + { name = "filelock", marker = "sys_platform == 'darwin'" }, + { name = "fsspec", marker = "sys_platform == 'darwin'" }, + { name = "jinja2", marker = "sys_platform == 'darwin'" }, + { name = "networkx", marker = "sys_platform == 'darwin'" }, + { name = "setuptools", marker = "python_full_version >= '3.12' and sys_platform == 'darwin'" }, + { name = "sympy", marker = "sys_platform == 'darwin'" }, + { name = "typing-extensions", marker = "sys_platform == 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/f6/5da3918414e07da9866ecb9330fe6ffdebe15cb9a4c5ada7d4b6e0a6654d/torch-2.7.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:d72acfdb86cee2a32c0ce0101606f3758f0d8bb5f8f31e7920dc2809e963aa7c", size = 68630914, upload-time = "2025-06-04T17:39:31.162Z" }, + { url = "https://files.pythonhosted.org/packages/5b/2b/d36d57c66ff031f93b4fa432e86802f84991477e522adcdffd314454326b/torch-2.7.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aea4fc1bf433d12843eb2c6b2204861f43d8364597697074c8d38ae2507f8730", size = 68640034, upload-time = "2025-06-04T17:39:17.989Z" }, + { url = "https://files.pythonhosted.org/packages/3a/60/04b77281c730bb13460628e518c52721257814ac6c298acd25757f6a175c/torch-2.7.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:787687087412c4bd68d315e39bc1223f08aae1d16a9e9771d95eabbb04ae98fb", size = 68645146, upload-time = "2025-06-04T17:38:52.97Z" }, + { url = "https://files.pythonhosted.org/packages/95/bf/e086ee36ddcef9299f6e708d3b6c8487c1651787bb9ee2939eb2a7f74911/torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:0da4f4dba9f65d0d203794e619fe7ca3247a55ffdcbd17ae8fb83c8b2dc9b585", size = 68925988, upload-time = "2025-06-04T17:38:29.273Z" }, + { url = "https://files.pythonhosted.org/packages/b1/29/beb45cdf5c4fc3ebe282bf5eafc8dfd925ead7299b3c97491900fe5ed844/torch-2.7.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:988b0cbc4333618a1056d2ebad9eb10089637b659eb645434d0809d8d937b946", size = 68645708, upload-time = "2025-06-04T17:34:39.852Z" }, +] + +[[package]] +name = "torch" +version = "2.8.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'linux'", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", +] +dependencies = [ + { name = "filelock", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "fsspec", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "jinja2", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "networkx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -5287,53 +5342,51 @@ dependencies = [ { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "setuptools", marker = "python_full_version >= '3.12'" }, - { name = "sympy" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and sys_platform == 'linux') or (python_full_version >= '3.12' and sys_platform == 'win32')" }, + { name = "sympy", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/5b/30/bfebdd8ec77db9a79775121789992d6b3b75ee5494971294d7b4b7c999bc/torch-2.10.0-2-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:2b980edd8d7c0a68c4e951ee1856334a43193f98730d97408fbd148c1a933313", size = 79411457, upload-time = "2026-02-10T21:44:59.189Z" }, - { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" }, - { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, - { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, - { url = "https://files.pythonhosted.org/packages/16/ee/efbd56687be60ef9af0c9c0ebe106964c07400eade5b0af8902a1d8cd58c/torch-2.10.0-3-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a1ff626b884f8c4e897c4c33782bdacdff842a165fee79817b1dd549fdda1321", size = 915510070, upload-time = "2026-03-11T14:16:39.386Z" }, - { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" }, - { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, - { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, - { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, - { url = "https://files.pythonhosted.org/packages/f4/39/590742415c3030551944edc2ddc273ea1fdfe8ffb2780992e824f1ebee98/torch-2.10.0-3-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b1d5e2aba4eb7f8e87fbe04f86442887f9167a35f092afe4c237dfcaaef6e328", size = 915632474, upload-time = "2026-03-11T14:15:13.666Z" }, - { url = "https://files.pythonhosted.org/packages/b6/8e/34949484f764dde5b222b7fe3fede43e4a6f0da9d7f8c370bb617d629ee2/torch-2.10.0-3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0228d20b06701c05a8f978357f657817a4a63984b0c90745def81c18aedfa591", size = 915523882, upload-time = "2026-03-11T14:14:46.311Z" }, - { url = "https://files.pythonhosted.org/packages/0c/1a/c61f36cfd446170ec27b3a4984f072fd06dab6b5d7ce27e11adb35d6c838/torch-2.10.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:5276fa790a666ee8becaffff8acb711922252521b28fbce5db7db5cf9cb2026d", size = 145992962, upload-time = "2026-01-21T16:24:14.04Z" }, - { url = "https://files.pythonhosted.org/packages/b5/60/6662535354191e2d1555296045b63e4279e5a9dbad49acf55a5d38655a39/torch-2.10.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:aaf663927bcd490ae971469a624c322202a2a1e68936eb952535ca4cd3b90444", size = 915599237, upload-time = "2026-01-21T16:23:25.497Z" }, - { url = "https://files.pythonhosted.org/packages/40/b8/66bbe96f0d79be2b5c697b2e0b187ed792a15c6c4b8904613454651db848/torch-2.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:a4be6a2a190b32ff5c8002a0977a25ea60e64f7ba46b1be37093c141d9c49aeb", size = 113720931, upload-time = "2026-01-21T16:24:23.743Z" }, - { url = "https://files.pythonhosted.org/packages/76/bb/d820f90e69cda6c8169b32a0c6a3ab7b17bf7990b8f2c680077c24a3c14c/torch-2.10.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:35e407430795c8d3edb07a1d711c41cc1f9eaddc8b2f1cc0a165a6767a8fb73d", size = 79411450, upload-time = "2026-01-21T16:25:30.692Z" }, - { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" }, - { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" }, - { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" }, - { url = "https://files.pythonhosted.org/packages/61/d8/15b9d9d3a6b0c01b883787bd056acbe5cc321090d4b216d3ea89a8fcfdf3/torch-2.10.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:b7bd80f3477b830dd166c707c5b0b82a898e7b16f59a7d9d42778dd058272e8b", size = 79423461, upload-time = "2026-01-21T16:24:50.266Z" }, - { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, - { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, - { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, - { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, - { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" }, - { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" }, - { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" }, - { url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" }, - { url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" }, - { url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" }, - { url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" }, - { url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" }, - { url = "https://files.pythonhosted.org/packages/4f/93/716b5ac0155f1be70ed81bacc21269c3ece8dba0c249b9994094110bfc51/torch-2.10.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:bf0d9ff448b0218e0433aeb198805192346c4fd659c852370d5cc245f602a06a", size = 79464992, upload-time = "2026-01-21T16:23:05.162Z" }, - { url = "https://files.pythonhosted.org/packages/69/2b/51e663ff190c9d16d4a8271203b71bc73a16aa7619b9f271a69b9d4a936b/torch-2.10.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:233aed0659a2503b831d8a67e9da66a62c996204c0bba4f4c442ccc0c68a3f60", size = 146018567, upload-time = "2026-01-21T16:22:23.393Z" }, - { url = "https://files.pythonhosted.org/packages/5e/cd/4b95ef7f293b927c283db0b136c42be91c8ec6845c44de0238c8c23bdc80/torch-2.10.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:682497e16bdfa6efeec8cde66531bc8d1fbbbb4d8788ec6173c089ed3cc2bfe5", size = 915721646, upload-time = "2026-01-21T16:21:16.983Z" }, - { url = "https://files.pythonhosted.org/packages/56/97/078a007208f8056d88ae43198833469e61a0a355abc0b070edd2c085eb9a/torch-2.10.0-cp314-cp314-win_amd64.whl", hash = "sha256:6528f13d2a8593a1a412ea07a99812495bec07e9224c28b2a25c0a30c7da025c", size = 113752373, upload-time = "2026-01-21T16:22:13.471Z" }, - { url = "https://files.pythonhosted.org/packages/d8/94/71994e7d0d5238393df9732fdab607e37e2b56d26a746cb59fdb415f8966/torch-2.10.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f5ab4ba32383061be0fb74bda772d470140a12c1c3b58a0cfbf3dae94d164c28", size = 79850324, upload-time = "2026-01-21T16:22:09.494Z" }, - { url = "https://files.pythonhosted.org/packages/e2/65/1a05346b418ea8ccd10360eef4b3e0ce688fba544e76edec26913a8d0ee0/torch-2.10.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:716b01a176c2a5659c98f6b01bf868244abdd896526f1c692712ab36dbaf9b63", size = 146006482, upload-time = "2026-01-21T16:22:18.42Z" }, - { url = "https://files.pythonhosted.org/packages/1d/b9/5f6f9d9e859fc3235f60578fa64f52c9c6e9b4327f0fe0defb6de5c0de31/torch-2.10.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:d8f5912ba938233f86361e891789595ff35ca4b4e2ac8fe3670895e5976731d6", size = 915613050, upload-time = "2026-01-21T16:20:49.035Z" }, - { url = "https://files.pythonhosted.org/packages/66/4d/35352043ee0eaffdeff154fad67cd4a31dbed7ff8e3be1cc4549717d6d51/torch-2.10.0-cp314-cp314t-win_amd64.whl", hash = "sha256:71283a373f0ee2c89e0f0d5f446039bdabe8dbc3c9ccf35f0f784908b0acd185", size = 113995816, upload-time = "2026-01-21T16:22:05.312Z" }, + { url = "https://files.pythonhosted.org/packages/63/28/110f7274254f1b8476c561dada127173f994afa2b1ffc044efb773c15650/torch-2.8.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0be92c08b44009d4131d1ff7a8060d10bafdb7ddcb7359ef8d8c5169007ea905", size = 102052793, upload-time = "2025-08-06T14:53:15.852Z" }, + { url = "https://files.pythonhosted.org/packages/70/1c/58da560016f81c339ae14ab16c98153d51c941544ae568da3cb5b1ceb572/torch-2.8.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:89aa9ee820bb39d4d72b794345cccef106b574508dd17dbec457949678c76011", size = 888025420, upload-time = "2025-08-06T14:54:18.014Z" }, + { url = "https://files.pythonhosted.org/packages/70/87/f69752d0dd4ba8218c390f0438130c166fa264a33b7025adb5014b92192c/torch-2.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:e8e5bf982e87e2b59d932769938b698858c64cc53753894be25629bdf5cf2f46", size = 241363614, upload-time = "2025-08-06T14:53:31.496Z" }, + { url = "https://files.pythonhosted.org/packages/8f/c4/3e7a3887eba14e815e614db70b3b529112d1513d9dae6f4d43e373360b7f/torch-2.8.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:220a06fd7af8b653c35d359dfe1aaf32f65aa85befa342629f716acb134b9710", size = 102073391, upload-time = "2025-08-06T14:53:20.937Z" }, + { url = "https://files.pythonhosted.org/packages/5a/63/4fdc45a0304536e75a5e1b1bbfb1b56dd0e2743c48ee83ca729f7ce44162/torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:c12fa219f51a933d5f80eeb3a7a5d0cbe9168c0a14bbb4055f1979431660879b", size = 888063640, upload-time = "2025-08-06T14:55:05.325Z" }, + { url = "https://files.pythonhosted.org/packages/84/57/2f64161769610cf6b1c5ed782bd8a780e18a3c9d48931319f2887fa9d0b1/torch-2.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:8c7ef765e27551b2fbfc0f41bcf270e1292d9bf79f8e0724848b1682be6e80aa", size = 241366752, upload-time = "2025-08-06T14:53:38.692Z" }, + { url = "https://files.pythonhosted.org/packages/49/0c/2fd4df0d83a495bb5e54dca4474c4ec5f9c62db185421563deeb5dabf609/torch-2.8.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e2fab4153768d433f8ed9279c8133a114a034a61e77a3a104dcdf54388838705", size = 101906089, upload-time = "2025-08-06T14:53:52.631Z" }, + { url = "https://files.pythonhosted.org/packages/99/a8/6acf48d48838fb8fe480597d98a0668c2beb02ee4755cc136de92a0a956f/torch-2.8.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b2aca0939fb7e4d842561febbd4ffda67a8e958ff725c1c27e244e85e982173c", size = 887913624, upload-time = "2025-08-06T14:56:44.33Z" }, + { url = "https://files.pythonhosted.org/packages/af/8a/5c87f08e3abd825c7dfecef5a0f1d9aa5df5dd0e3fd1fa2f490a8e512402/torch-2.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:2f4ac52f0130275d7517b03a33d2493bab3693c83dcfadf4f81688ea82147d2e", size = 241326087, upload-time = "2025-08-06T14:53:46.503Z" }, + { url = "https://files.pythonhosted.org/packages/10/4e/469ced5a0603245d6a19a556e9053300033f9c5baccf43a3d25ba73e189e/torch-2.8.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:2b2f96814e0345f5a5aed9bf9734efa913678ed19caf6dc2cddb7930672d6128", size = 101936856, upload-time = "2025-08-06T14:54:01.526Z" }, + { url = "https://files.pythonhosted.org/packages/16/82/3948e54c01b2109238357c6f86242e6ecbf0c63a1af46906772902f82057/torch-2.8.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:65616ca8ec6f43245e1f5f296603e33923f4c30f93d65e103d9e50c25b35150b", size = 887922844, upload-time = "2025-08-06T14:55:50.78Z" }, + { url = "https://files.pythonhosted.org/packages/e3/54/941ea0a860f2717d86a811adf0c2cd01b3983bdd460d0803053c4e0b8649/torch-2.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:659df54119ae03e83a800addc125856effda88b016dfc54d9f65215c3975be16", size = 241330968, upload-time = "2025-08-06T14:54:45.293Z" }, + { url = "https://files.pythonhosted.org/packages/15/0e/8a800e093b7f7430dbaefa80075aee9158ec22e4c4fc3c1a66e4fb96cb4f/torch-2.8.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:83c13411a26fac3d101fe8035a6b0476ae606deb8688e904e796a3534c197def", size = 102020139, upload-time = "2025-08-06T14:54:39.047Z" }, + { url = "https://files.pythonhosted.org/packages/4a/15/5e488ca0bc6162c86a33b58642bc577c84ded17c7b72d97e49b5833e2d73/torch-2.8.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8f0a9d617a66509ded240add3754e462430a6c1fc5589f86c17b433dd808f97a", size = 887990692, upload-time = "2025-08-06T14:56:18.286Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a8/6a04e4b54472fc5dba7ca2341ab219e529f3c07b6941059fbf18dccac31f/torch-2.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:a7242b86f42be98ac674b88a4988643b9bc6145437ec8f048fea23f72feb5eca", size = 241603453, upload-time = "2025-08-06T14:55:22.945Z" }, +] + +[[package]] +name = "torch" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "filelock", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "fsspec", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "jinja2", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "networkx", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "setuptools", marker = "python_full_version >= '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "sympy", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "typing-extensions", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'" }, ] [[package]] @@ -5385,16 +5438,17 @@ wheels = [ [[package]] name = "triton" -version = "3.6.0" +version = "3.4.0" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "setuptools", marker = "(python_full_version < '3.12' and sys_platform == 'win32') or sys_platform == 'linux'" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/8c/f7/f1c9d3424ab199ac53c2da567b859bcddbb9c9e7154805119f8bd95ec36f/triton-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a6550fae429e0667e397e5de64b332d1e5695b73650ee75a6146e2e902770bea", size = 188105201, upload-time = "2026-01-20T16:00:29.272Z" }, - { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, - { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, - { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, - { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, - { url = "https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803", size = 188273063, upload-time = "2026-01-20T16:01:07.278Z" }, - { url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" }, + { url = "https://files.pythonhosted.org/packages/62/ee/0ee5f64a87eeda19bbad9bc54ae5ca5b98186ed00055281fd40fb4beb10e/triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128", size = 155430069, upload-time = "2025-07-30T19:58:21.715Z" }, + { url = "https://files.pythonhosted.org/packages/7d/39/43325b3b651d50187e591eefa22e236b2981afcebaefd4f2fc0ea99df191/triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b70f5e6a41e52e48cfc087436c8a28c17ff98db369447bcaff3b887a3ab4467", size = 155531138, upload-time = "2025-07-30T19:58:29.908Z" }, + { url = "https://files.pythonhosted.org/packages/d0/66/b1eb52839f563623d185f0927eb3530ee4d5ffe9d377cdaf5346b306689e/triton-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31c1d84a5c0ec2c0f8e8a072d7fd150cab84a9c239eaddc6706c081bfae4eb04", size = 155560068, upload-time = "2025-07-30T19:58:37.081Z" }, + { url = "https://files.pythonhosted.org/packages/30/7b/0a685684ed5322d2af0bddefed7906674f67974aa88b0fae6e82e3b766f6/triton-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00be2964616f4c619193cb0d1b29a99bd4b001d7dc333816073f92cf2a8ccdeb", size = 155569223, upload-time = "2025-07-30T19:58:44.017Z" }, + { url = "https://files.pythonhosted.org/packages/20/63/8cb444ad5cdb25d999b7d647abac25af0ee37d292afc009940c05b82dda0/triton-3.4.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7936b18a3499ed62059414d7df563e6c163c5e16c3773678a3ee3d417865035d", size = 155659780, upload-time = "2025-07-30T19:58:51.171Z" }, ] [[package]] @@ -5597,9 +5651,13 @@ dependencies = [ { name = "koinapy" }, { name = "matchms" }, { name = "matplotlib" }, + { name = "polars" }, + { name = "safetensors" }, { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "tomli" }, + { name = "torch", version = "2.7.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typer" }, ] @@ -5632,8 +5690,12 @@ requires-dist = [ { name = "koinapy", specifier = ">=0.0.10" }, { name = "matchms", specifier = ">=0.31.0" }, { name = "matplotlib", specifier = ">=3.7.0" }, + { name = "polars", specifier = ">=1.39.3" }, + { name = "safetensors", specifier = ">=0.7.0" }, { name = "scikit-learn", specifier = ">=1.3.0" }, { name = "tomli", specifier = ">=2.2.1" }, + { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.6,<2.8" }, + { name = "torch", marker = "sys_platform == 'linux' or sys_platform == 'win32'", specifier = ">=2.6,<2.9" }, { name = "typer", specifier = ">=0.15.2" }, ] diff --git a/winnow/calibration/calibrator.py b/winnow/calibration/calibrator.py index d6c2436d..ce118f27 100644 --- a/winnow/calibration/calibrator.py +++ b/winnow/calibration/calibrator.py @@ -1,83 +1,307 @@ -"""Contains classes and functions for probability recalibration.""" +"""Contains classes and functions for PSM rescoring and probability calibration.""" -from typing import Dict, List, Tuple, Union, Optional +from __future__ import annotations + +import importlib +import json +import logging +from dataclasses import asdict, dataclass, field, fields from pathlib import Path -import pickle +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + import numpy as np -from sklearn.neural_network import MLPClassifier -from sklearn.preprocessing import StandardScaler +from tqdm import tqdm +import torch +import torch.nn as nn from numpy.typing import NDArray -from huggingface_hub import snapshot_download from omegaconf import DictConfig +from safetensors.torch import load_file, save_file +from torch.utils.data import DataLoader from winnow.calibration.calibration_features import ( CalibrationFeatures, FeatureDependency, ) +from winnow.calibration.features.utils import validate_model_input_params from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.datasets.feature_dataset import FeatureDataset +from winnow.utils.koina_intensity_config import ( + resolve_feature_model_inputs, + strip_runtime_keys_from_feature_config, +) +from winnow.utils.paths import resolve_data_path + +logger = logging.getLogger(__name__) + + +def _deserialize_calibration_feature( + feature_config: Dict[str, Any], +) -> CalibrationFeatures: + """Rebuild a feature instance from a saved ``config.json`` entry.""" + feature_config = strip_runtime_keys_from_feature_config(dict(feature_config)) + target = feature_config.pop("_target_") + module_path, class_name = target.rsplit(".", 1) + module = importlib.import_module(module_path) + feature_class = getattr(module, class_name) + return feature_class(**feature_config) + + +def _apply_koina_constant_overrides( + model_input_constants: Dict[str, Any], + model_input_columns: Dict[str, str], + overrides: Optional[Dict[str, Any]], +) -> None: + """Merge constant overrides and drop those keys from column mappings.""" + if not overrides: + return + for key, value in overrides.items(): + if value is None: + continue + model_input_constants[key] = value + model_input_columns.pop(key, None) + + +def _apply_koina_column_overrides( + model_input_constants: Dict[str, Any], + model_input_columns: Dict[str, str], + overrides: Optional[Dict[str, str]], +) -> None: + """Merge column overrides and drop those keys from constants.""" + if not overrides: + return + for key, value in overrides.items(): + if value is None: + continue + # Constants applied first take precedence; ignore default column entries in + # the composed config when a constant override is already set for this key. + if key in model_input_constants: + continue + model_input_columns[key] = value + model_input_constants.pop(key, None) + + +def _feature_config_for_save(feature: CalibrationFeatures) -> Dict[str, Any]: + """Serialise a feature for ``config.json``, omitting runtime-only Koina keys.""" + return strip_runtime_keys_from_feature_config(feature.get_config()) + + +def _merge_koina_overrides_into_single_feature( + feature: Any, + model_input_constants: Optional[Dict[str, Any]], + model_input_columns: Optional[Dict[str, str]], +) -> None: + """Apply Koina constant/column overrides to one feature object.""" + merged_constants = dict(feature.model_input_constants or {}) + merged_columns = dict(feature.model_input_columns or {}) + _apply_koina_constant_overrides( + merged_constants, merged_columns, model_input_constants + ) + _apply_koina_column_overrides(merged_constants, merged_columns, model_input_columns) + resolved_constants, resolved_columns = resolve_feature_model_inputs( + merged_constants or None, merged_columns or None + ) + validate_model_input_params(resolved_constants, resolved_columns) + feature.model_input_constants = resolved_constants + feature.model_input_columns = resolved_columns + + +def _load_saved_features_section( + calibrator: Any, features_section: Dict[str, Any] +) -> None: + """Populate ``calibrator`` with features from a ``config.json`` ``features`` block.""" + for _name, feature_config in features_section.items(): + calibrator.add_feature(_deserialize_calibration_feature(feature_config)) + + +@dataclass +class TrainingHistory: + """Epoch-level training metrics from calibrator fitting. + + Attributes: + train_losses: Training loss at the end of each epoch. + val_losses: Validation loss per epoch (present only when validation data is used). + val_accuracies: Validation accuracy per epoch. + best_epoch: Epoch with the best validation loss (0-indexed), or the last + epoch if no validation data was provided. + epochs_trained: Total number of epochs completed. + final_val_loss: Loss on the full validation set after restoring the best + weights (only when validation was subsampled for early stopping). + final_val_accuracy: Accuracy on the full validation set in that case. + """ + + train_losses: List[float] = field(default_factory=list) + val_losses: Optional[List[float]] = None + val_accuracies: Optional[List[float]] = None + best_epoch: int = 0 + epochs_trained: int = 0 + final_val_loss: Optional[float] = None + final_val_accuracy: Optional[float] = None + + def save(self, path: Union[Path, str]) -> None: + """Save training history to a JSON file.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + json.dump(asdict(self), f, indent=2) + + @classmethod + def load(cls, path: Union[Path, str]) -> TrainingHistory: + """Load training history from a JSON file.""" + with open(path) as f: + data = json.load(f) + known = {f.name for f in fields(cls)} + data = {k: v for k, v in data.items() if k in known} + return cls(**data) + + def plot( + self, + output_path: Optional[Union[Path, str]] = None, + show: bool = False, + ) -> None: + """Plot training and validation curves. + + Args: + output_path: Path to save the plot image. Not saved if ``None``. + show: Whether to call ``plt.show()`` interactively. + """ + import matplotlib.pyplot as plt + + fig, ax = plt.subplots() + epochs = range(1, len(self.train_losses) + 1) + + ax.plot(epochs, self.train_losses, label="Train Loss", linestyle="-") + + if self.val_losses is not None: + ax.plot(epochs, self.val_losses, label="Validation Loss", linestyle="--") + + if self.val_accuracies is not None: + ax2 = ax.twinx() + ax2.plot( + epochs, self.val_accuracies, label="Validation Accuracy", linestyle=":" + ) + ax2.set_ylabel("Validation Accuracy") + ax2.legend(loc="lower left") + + plt.tight_layout() + plt.grid(False) + ax.set_xlabel("Epoch") + ax.set_ylabel("Loss") + ax.set_title("Calibrator Training Progress") + ax.legend(loc="upper left") + + if output_path is not None: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, bbox_inches="tight") + + if show: + plt.show() + + plt.close(fig) + + +class CalibratorNetwork(nn.Module): + """Feed-forward neural network for probability calibration. + + Args: + input_dim: Number of input features. + hidden_dims: Sizes of hidden layers. + dropout: Dropout rate applied after each hidden layer. + """ + + def __init__( + self, + input_dim: int, + hidden_dims: Tuple[int, ...] = (128, 64), + dropout: float = 0.1, + ) -> None: + super().__init__() + layers: List[nn.Module] = [] + prev_dim = input_dim + for dim in hidden_dims: + layers.extend( + [ + nn.Linear(prev_dim, dim), + nn.ReLU(), + nn.Dropout(dropout), + ] + ) + prev_dim = dim + layers.append(nn.Linear(prev_dim, 1)) + self.network = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Return raw logits (scalar per sample).""" + return self.network(x).squeeze(-1) class ProbabilityCalibrator: - """A class for recalibrating probabilities for a de novo peptide sequencing method. + """Recalibrates probabilities for de novo peptide sequencing predictions. - This class provides functionality to recalibrate predicted probabilities by fitting an MLP classifier using various features computed from a calibration dataset. + Uses a PyTorch feed-forward network trained on pre-computed calibration + features. Models are persisted as ``safetensors`` weights plus a + ``config.json`` metadata file. """ def __init__( self, - seed: int = 42, features: Optional[ - Union[List[CalibrationFeatures], Dict[str, CalibrationFeatures], DictConfig] + Union[ + List[CalibrationFeatures], + Dict[str, CalibrationFeatures], + DictConfig, + ] ] = None, - hidden_layer_sizes: Tuple[int, ...] = (50, 50), - learning_rate_init: float = 0.001, - alpha: float = 0.0001, - max_iter: int = 1000, - early_stopping: bool = True, - validation_fraction: float = 0.1, + hidden_dims: Tuple[int, ...] = (128, 64), + dropout: float = 0.1, + learning_rate: float = 1e-3, + weight_decay: float = 1e-4, + max_epochs: int = 100, + batch_size: int = 1024, + n_iter_no_change: int = 10, + tol: float = 1e-4, + seed: int = 42, + val_early_stopping_max_psms: Optional[int] = 10000, + val_subsample_seed: Optional[int] = None, ) -> None: - """Initialise the probability calibrator. + self.hidden_dims = tuple(hidden_dims) + self.dropout = dropout + self.learning_rate = learning_rate + self.weight_decay = weight_decay + if max_epochs < 1: + raise ValueError(f"max_epochs must be at least 1, got {max_epochs}.") + self.max_epochs = max_epochs + self.batch_size = batch_size + self.n_iter_no_change = n_iter_no_change + self.tol = tol + self.seed = seed + self.val_early_stopping_max_psms = val_early_stopping_max_psms + self.val_subsample_seed = val_subsample_seed - Args: - seed (int): Random seed for the classifier. Defaults to 42. - features (Optional[Union[List[CalibrationFeatures], Dict[str, CalibrationFeatures], DictConfig]]): - Features to add to the calibrator. Can be a list or dict of CalibrationFeatures objects. - If None, no features are added. Defaults to None. - hidden_layer_sizes (Tuple[int, ...]): The number of neurons in each hidden layer. Defaults to (50, 50). - learning_rate_init (float): The initial learning rate. Defaults to 0.001. - alpha (float): L2 regularisation parameter. Defaults to 0.0001. - max_iter (int): Maximum number of training iterations. Defaults to 1000. - early_stopping (bool): Whether to use early stopping to terminate training. Defaults to True. - validation_fraction (float): Proportion of training data to use for early stopping validation. Defaults to 0.1. - """ self.feature_dict: Dict[str, CalibrationFeatures] = {} self.dependencies: Dict[str, FeatureDependency] = {} self.dependency_reference_counter: Dict[str, int] = {} - self.classifier = MLPClassifier( - random_state=seed, - hidden_layer_sizes=hidden_layer_sizes, - learning_rate_init=learning_rate_init, - alpha=alpha, - max_iter=max_iter, - early_stopping=early_stopping, - validation_fraction=validation_fraction, - ) - self.scaler = StandardScaler() - # Add features if provided + self.network: Optional[CalibratorNetwork] = None + self.feature_mean: Optional[torch.Tensor] = None + self.feature_std: Optional[torch.Tensor] = None + self._active_feature_columns: Optional[List[str]] = None + if features is not None: if isinstance(features, (dict, DictConfig)): self.add_features(list(features.values())) else: self.add_features(list(features)) + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + @property def columns(self) -> List[str]: - """Returns the list of column names corresponding to the features added to the calibrator. - - Returns: - List[str]: A list of column names representing the features used for calibration. - """ + """Column names corresponding to the features added to the calibrator.""" + if self._active_feature_columns is not None: + return list(self._active_feature_columns) return [ column for feature in self.feature_dict.values() @@ -85,28 +309,101 @@ def columns(self) -> List[str]: ] @property - def features(self) -> List[str]: - """Get the list of features added to the calibrator. - - Returns: - List[str]: The list of feature names - """ + def feature_names(self) -> List[str]: + """Names of features registered with the calibrator.""" return list(self.feature_dict.keys()) + # ------------------------------------------------------------------ + # Static methods + # ------------------------------------------------------------------ + + @staticmethod + def _resolve_device() -> torch.device: + """Auto-detect GPU, warn if falling back to CPU.""" + if torch.cuda.is_available(): + device = torch.device("cuda") + logger.info("Using GPU: %s", torch.cuda.get_device_name(0)) + else: + logger.warning( + "No GPU detected, falling back to CPU. " + "Training may be slow for large datasets." + ) + device = torch.device("cpu") + return device + + @staticmethod + def _reject_legacy_pickle_checkpoint(dir_path: Path) -> None: + """Raise if ``dir_path`` contains a pre-PyTorch pickle checkpoint.""" + legacy_pickle = dir_path / "calibrator.pkl" + config_path = dir_path / "config.json" + if legacy_pickle.exists() and not config_path.exists(): + raise ValueError( + "Legacy pickle checkpoint format is no longer supported. " + "The directory contains calibrator.pkl from a pre-PyTorch " + "version of winnow, but this version expects model.safetensors " + "and config.json. Retrain the calibrator with the current " + "version to produce a compatible checkpoint." + ) + + # ------------------------------------------------------------------ + # Class methods (persistence) + # ------------------------------------------------------------------ + @classmethod def save( - cls, calibrator: "ProbabilityCalibrator", dir_path: Union[Path, str] + cls, + calibrator: ProbabilityCalibrator, + dir_path: Union[Path, str], ) -> None: - """Save the calibrator to a file. + """Save calibrator weights and config to a directory. + + Writes ``model.safetensors`` (network weights, feature mean/std) and + ``config.json`` (architecture, hyperparameters, resolved feature configs). Args: - calibrator (ProbabilityCalibrator): The calibrator to save. - dir_path (Path): The path to the directory where the calibrator checkpoint will be saved. + calibrator: The calibrator to save. + dir_path: Destination directory. """ - if isinstance(dir_path, str): - dir_path = Path(dir_path) + if calibrator.network is None: + raise RuntimeError("Cannot save an unfitted calibrator.") + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + + dir_path = Path(dir_path) dir_path.mkdir(parents=True, exist_ok=True) - pickle.dump(calibrator, open(dir_path / "calibrator.pkl", "wb")) + + tensors: Dict[str, torch.Tensor] = {} + for name, param in calibrator.network.state_dict().items(): + tensors[f"network.{name}"] = param.cpu() + tensors["feature_mean"] = calibrator.feature_mean.cpu() + tensors["feature_std"] = calibrator.feature_std.cpu() + + save_file(tensors, dir_path / "model.safetensors") + + config = { + "input_dim": int(calibrator.feature_mean.shape[0]), + "hidden_dims": list(calibrator.hidden_dims), + "dropout": calibrator.dropout, + "learning_rate": calibrator.learning_rate, + "weight_decay": calibrator.weight_decay, + "max_epochs": calibrator.max_epochs, + "batch_size": calibrator.batch_size, + "n_iter_no_change": calibrator.n_iter_no_change, + "tol": calibrator.tol, + "seed": calibrator.seed, + "val_early_stopping_max_psms": calibrator.val_early_stopping_max_psms, + "val_subsample_seed": calibrator.val_subsample_seed, + "feature_columns": calibrator.columns, + "feature_names": calibrator.feature_names, + "features": { + name: _feature_config_for_save(feature) + for name, feature in calibrator.feature_dict.items() + }, + } + with open(dir_path / "config.json", "w") as f: + json.dump(config, f, indent=2) + + logger.info("Saved calibrator to %s", dir_path) @classmethod def load( @@ -115,65 +412,130 @@ def load( Path, str ] = "InstaDeepAI/winnow-general-model", cache_dir: Optional[Path] = None, - ) -> "ProbabilityCalibrator": - """Load a pretrained calibrator from a local path or Hugging Face repository. If the path is a local directory path, it will be used directly. If it is a Hugging Face repository identifier, it will be downloaded from Hugging Face. + ) -> ProbabilityCalibrator: + """Load a calibrator from a local directory or HuggingFace Hub repo. + + Expects the directory to contain ``model.safetensors`` and + ``config.json``. Args: - pretrained_model_name_or_path (Union[Path, str]): The local directory path (e.g., "./my-model-directory") or the Hugging Face repository identifier (e.g., "InstaDeepAI/winnow-general-model"). - cache_dir (Optional[Path]): Directory to cache the Hugging Face model. + pretrained_model_name_or_path: Local directory path or HuggingFace + repository identifier. + cache_dir: Optional cache directory for HuggingFace downloads. + + Returns: + A fitted ``ProbabilityCalibrator`` ready for prediction. + + Raises: + FileNotFoundError: If the path cannot be resolved. + ValueError: If the directory contains a legacy pickle checkpoint. """ - dir_path = Path(pretrained_model_name_or_path) + dir_path = resolve_data_path( + str(pretrained_model_name_or_path), + repo_type="model", + cache_dir=cache_dir, + ) - # If the path exists locally, use it directly. - if dir_path.exists(): - # Resolve relative paths to absolute, canonical paths - dir_path = dir_path.resolve() - # Otherwise download it from Hugging Face. - else: - # If no cache directory is provided, use the default cache directory. - if cache_dir is None: - cache_dir = Path.home() / ".cache" / "huggingface" / "hub" - # Download from Hugging Face - dir_path = Path( - snapshot_download( - repo_id=pretrained_model_name_or_path, - repo_type="model", - cache_dir=cache_dir, - ) - ) + cls._reject_legacy_pickle_checkpoint(dir_path) - # Load the calibrator object - loaded_obj = pickle.load(open(dir_path / "calibrator.pkl", "rb")) - - # Check if this is a legacy checkpoint (MLPClassifier instead of ProbabilityCalibrator) - if isinstance(loaded_obj, MLPClassifier): - error_msg = ( - "Legacy checkpoint format detected. The checkpoint directory contains " - "an old format where calibrator.pkl contains only the MLPClassifier " - "instead of the full ProbabilityCalibrator object.\n" - "Legacy checkpoints cannot be automatically migrated because they lack " - "the feature and dependency information required by the current version. " - "We cannot correctly infer the trained feature set with old versions.\n" - "To resolve this, retrain the calibrator using the current version with your training dataset. " - "The new format will save the complete ProbabilityCalibrator object including all features and dependencies." - ) - raise ValueError(error_msg) + config_path = dir_path / "config.json" + with open(config_path) as f: + config = json.load(f) - elif not isinstance(loaded_obj, ProbabilityCalibrator): - raise ValueError( - f"Loaded object is of type {type(loaded_obj).__name__}, expected ProbabilityCalibrator." + calibrator = cls( + hidden_dims=tuple(config["hidden_dims"]), + dropout=config["dropout"], + learning_rate=config.get("learning_rate", 1e-3), + weight_decay=config.get("weight_decay", 1e-4), + max_epochs=config.get("max_epochs", 100), + batch_size=config.get("batch_size", 1024), + n_iter_no_change=config.get("n_iter_no_change", config.get("patience", 10)), + tol=config.get("tol", 1e-4), + seed=config.get("seed", 42), + val_early_stopping_max_psms=config.get( + "val_early_stopping_max_psms", 10000 + ), + val_subsample_seed=config.get("val_subsample_seed"), + ) + + _load_saved_features_section(calibrator, config.get("features", {})) + saved_columns = config.get("feature_columns") + if saved_columns: + calibrator._active_feature_columns = list(saved_columns) + + tensors = load_file(dir_path / "model.safetensors") + + calibrator.feature_mean = tensors["feature_mean"] + calibrator.feature_std = tensors["feature_std"] + + input_dim = config["input_dim"] + calibrator.network = CalibratorNetwork( + input_dim=input_dim, + hidden_dims=calibrator.hidden_dims, + dropout=calibrator.dropout, + ) + network_state = { + k.removeprefix("network."): v + for k, v in tensors.items() + if k.startswith("network.") + } + calibrator.network.load_state_dict(network_state) + calibrator.network.eval() + + logger.info( + "Loaded calibrator from %s (input_dim=%d, hidden_dims=%s, features=%s)", + dir_path, + input_dim, + calibrator.hidden_dims, + list(calibrator.feature_dict.keys()), + ) + calibrator._validate_loaded_checkpoint(input_dim) + return calibrator + + def apply_koina_model_input_overrides( + self, + model_input_constants: Optional[Dict[str, Any]] = None, + model_input_columns: Optional[Dict[str, str]] = None, + ) -> None: + """Merge inference-time Koina inputs into intensity-based features. + + Used by ``winnow predict`` from the top-level config keys + ``koina.input_constants`` and ``koina.input_columns`` (same names as + in ``calibrator.yaml``). + + Updates :attr:`model_input_constants` / :attr:`model_input_columns` on + features that define them (e.g. fragment match, chimeric). Per-key + overrides replace any previous constant or column mapping for that key. + + Args: + model_input_constants: Koina input names mapped to constant values + tiled across all rows. + model_input_columns: Koina input names mapped to metadata column + names for per-row values. + """ + for feature in self.feature_dict.values(): + if not hasattr(feature, "model_input_constants") or not hasattr( + feature, "model_input_columns" + ): + continue + _merge_koina_overrides_into_single_feature( + feature, + model_input_constants, + model_input_columns, ) - else: - loaded_obj._validate_fitted_feature_dimensions() - return loaded_obj - def add_feature(self, feature: CalibrationFeatures) -> None: - """Add a feature for the classifier used for calibration. + # ------------------------------------------------------------------ + # Public instance methods + # ------------------------------------------------------------------ - This method ensures that the feature is unique and its dependencies are tracked. + def add_feature(self, feature: CalibrationFeatures) -> None: + """Register a feature for calibration. Args: - feature (CalibrationFeatures): The feature to be added to the calibrator. + feature: The feature to register. + + Raises: + KeyError: If a feature with the same name is already present. """ if feature.name not in self.feature_dict: self.feature_dict[feature.name] = feature @@ -187,22 +549,141 @@ def add_feature(self, feature: CalibrationFeatures) -> None: raise KeyError(f"Feature {feature.name} in feature set.") def add_features(self, features: List[CalibrationFeatures]) -> None: - """Add features for the classifier used for calibration. + """Register multiple features.""" + for feature in features: + self.add_feature(feature) + + def compute_features(self, dataset: CalibrationDataset) -> None: + """Run feature dependencies and feature computation on a dataset. + + Mutates ``dataset.metadata`` in place: adds feature columns and may + drop rows when ``learn_from_missing=False`` on individual features. Args: - features (List[CalibrationFeatures]): A list of features to be added to the calibrator. + dataset: The dataset on which features are computed. """ - for feature in features: - self.add_feature(feature) + for dependency in self.dependencies.values(): + dependency.compute(dataset=dataset) - def remove_feature(self, name: str) -> None: - """Remove a feature for the classifier used for calibration. + for feature in self.feature_dict.values(): + feature.prepare(dataset=dataset) + feature.compute(dataset=dataset) + + def fit( + self, + train_dataset: CalibrationDataset, + val_dataset: Optional[CalibrationDataset] = None, + progress_bar: bool = True, + ) -> TrainingHistory: + """Compute features and train the calibrator. - This method also removes any dependencies that are no longer required. + This is the primary training entry point. It runs + :meth:`compute_features` on the supplied datasets, extracts the + numeric feature matrix and labels, and trains the calibrator network. + + Both ``train_dataset`` and ``val_dataset`` are mutated in place + (feature columns are added, and rows may be dropped when individual + features have ``learn_from_missing=False``). + + Args: + train_dataset: Labelled training dataset. + val_dataset: Optional labelled validation dataset for early + stopping. Subsampling for large sets matches + :meth:`fit_from_features` (see ``val_early_stopping_max_psms``). + progress_bar: Whether to display a progress bar during training. + + Returns: + Epoch-level training metrics. + """ + self.compute_features(train_dataset) + features, labels = self._extract_feature_matrix(train_dataset, labelled=True) + train_fd = FeatureDataset( + features=np.asarray(features), labels=np.asarray(labels) + ) + + val_fd = None + if val_dataset is not None: + self.compute_features(val_dataset) + val_features, val_labels = self._extract_feature_matrix( + val_dataset, labelled=True + ) + val_fd = FeatureDataset( + features=np.asarray(val_features), + labels=np.asarray(val_labels), + ) + + return self._fit_from_features(train_fd, val_fd, progress_bar=progress_bar) + + def fit_from_features( + self, + train_dataset: FeatureDataset, + val_dataset: Optional[FeatureDataset] = None, + progress_bar: bool = True, + epoch_callback: Optional[Callable[[int, float], None]] = None, + ) -> TrainingHistory: + """Train the calibrator from pre-computed feature arrays. + + Use this for the two-phase workflow where features have already been + computed and saved to Parquet, then reloaded via + :meth:`FeatureDataset.from_parquet`. Args: - name (str): The name of the feature to be removed. + train_dataset: Training features and labels. + val_dataset: Optional held-out validation set for early stopping. + When it has more than ``val_early_stopping_max_psms`` rows, a + random subset (``val_subsample_seed``, defaulting to ``seed``) + is evaluated each epoch; after training, metrics on the full + validation set are stored in ``TrainingHistory.final_val_*``. + progress_bar: Whether to display a progress bar during training. + epoch_callback: Optional callback invoked after each validation + step with ``(epoch, val_loss)``. Useful for external + hyperparameter optimisation frameworks (e.g. Optuna pruning). + Only called when *val_dataset* is provided. + + Returns: + Epoch-level training metrics. """ + return self._fit_from_features( + train_dataset, val_dataset, progress_bar, epoch_callback + ) + + @torch.no_grad() + def predict(self, dataset: CalibrationDataset) -> None: + """Predict calibrated probabilities and store them on the dataset. + + The ``calibrated_confidence`` column is added to ``dataset.metadata``. + + Args: + dataset: The CalibrationDataset containing the features for which predictions are made. + """ + if self.network is None: + raise RuntimeError( + "Calibrator has not been fitted or loaded. Call fit() or load() first." + ) + assert self.feature_mean is not None + assert self.feature_std is not None + + if len(dataset) == 0: + raise ValueError( + "Cannot predict on an empty dataset: no spectra remain after " + "feature computation. Check feature validity filters and " + "learn_from_missing settings." + ) + + self._validate_fitted_feature_dimensions() + features = self._extract_feature_matrix(dataset, labelled=False) + device = next(self.network.parameters()).device + x = torch.as_tensor(features, dtype=torch.float32, device=device) + x = (x - self.feature_mean) / self.feature_std + + self.network.eval() + logits = self.network(x) + probs = torch.sigmoid(logits).cpu().numpy() + + dataset.metadata["calibrated_confidence"] = probs.tolist() + + def remove_feature(self, name: str) -> None: + """Remove a feature and its orphaned dependencies.""" feature = self.feature_dict.pop(name) for dependency in feature.dependencies: self.dependency_reference_counter[dependency.name] -= 1 @@ -210,48 +691,31 @@ def remove_feature(self, name: str) -> None: self.dependency_reference_counter.pop(dependency.name) self.dependencies.pop(dependency.name) - def fit(self, dataset: CalibrationDataset) -> None: - """Fit the MLP classifier using the given calibration dataset. - - This method computes the features from the dataset, prepares the labels, and trains an MLP classifier for recalibrating probabilities. - - Args: - dataset (CalibrationDataset): The dataset used for training the classifier. - """ - features, labels = self.compute_features(dataset=dataset, labelled=True) - # Fit and transform features with scaler - features_scaled = self.scaler.fit_transform(features) - self.classifier.fit(features_scaled, labels) + # ------------------------------------------------------------------ + # Private instance methods + # ------------------------------------------------------------------ - def compute_features( + def _extract_feature_matrix( self, dataset: CalibrationDataset, labelled: bool ) -> Union[ NDArray[np.float64], Tuple[NDArray[np.float64], NDArray[np.float64]], ]: - """Compute the features for the dataset, including any dependencies and feature calculations. + """Pull numeric feature columns (and optionally labels) from metadata. - This method handles both labelled and unlabelled datasets. It computes the necessary features and returns them for model training or prediction. + Must be called *after* :meth:`compute_features` has populated the + feature columns on ``dataset.metadata``. Args: - dataset (CalibrationDataset): The dataset from which features are computed. - labelled (bool): Whether the dataset contains labels for supervised learning. + dataset: Dataset whose metadata already contains the feature + columns. + labelled: If ``True``, also return the ``correct`` column as + labels. Returns: - Union[ - NDArray[np.float64], - Tuple[NDArray[np.float64], NDArray[np.float64]] - ]: - - If `labelled` is True: A tuple containing the computed feature matrix and the corresponding labels. - - If `labelled` is False: Only the computed feature matrix. + Feature matrix, or ``(features, labels)`` when + ``labelled=True``. """ - for dependency in self.dependencies.values(): - dependency.compute(dataset=dataset) - - for feature in self.feature_dict.values(): - feature.prepare(dataset=dataset) - feature.compute(dataset=dataset) - feature_columns = [dataset.confidence_column] feature_columns.extend(self.columns) features = dataset.metadata[feature_columns] @@ -259,34 +723,327 @@ def compute_features( if labelled: labels = dataset.metadata["correct"] return features.values, labels.values - else: - return features.values + return features.values - def predict(self, dataset: CalibrationDataset) -> None: - """Predict the calibrated probabilities for a given dataset. + def _early_stopping_validation_split( + self, + val_dataset: FeatureDataset, + ) -> Tuple[FeatureDataset, bool]: + """Return validation data for early stopping, optionally subsampled.""" + max_psms = self.val_early_stopping_max_psms + n_val = len(val_dataset) + if max_psms is None or max_psms <= 0 or n_val <= max_psms: + return val_dataset, False - This method computes the features and uses the trained classifier to predict the calibrated probabilities for the dataset. The calibrated probabilities are stored in the dataset under the "calibrated_confidence" column. + subsample_seed = ( + self.val_subsample_seed + if self.val_subsample_seed is not None + else self.seed + ) + rng = np.random.default_rng(subsample_seed) + idx = rng.permutation(n_val)[:max_psms] + val_early = FeatureDataset( + features=val_dataset.features[idx].cpu().numpy(), + labels=val_dataset.labels[idx].cpu().numpy(), + ) + logger.info( + "Using %d of %d validation PSMs for early stopping " + "(val_early_stopping_max_psms=%s, val_subsample_seed=%s).", + len(val_early), + n_val, + max_psms, + subsample_seed, + ) + return val_early, True + + def _record_full_validation_if_subsampled( + self, + val_was_subsampled: bool, + val_dataset: Optional[FeatureDataset], + history: TrainingHistory, + criterion: nn.Module, + device: torch.device, + ) -> None: + """After training, evaluate on full validation when early stopping used a subset.""" + if not val_was_subsampled or val_dataset is None: + return + assert self.network is not None + final_loss, final_acc = self._evaluate(val_dataset, criterion, device) + history.final_val_loss = final_loss + history.final_val_accuracy = final_acc + logger.info( + "Full validation set metrics after training (n=%d): " + "loss=%.6f, accuracy=%.4f", + len(val_dataset), + final_loss, + final_acc, + ) + + def _fit_from_features( + self, + train_dataset: FeatureDataset, + val_dataset: Optional[FeatureDataset] = None, + progress_bar: bool = True, + epoch_callback: Optional[Callable[[int, float], None]] = None, + ) -> TrainingHistory: + """Train the calibrator network on pre-computed features. Args: - dataset (CalibrationDataset): The dataset for which predictions are made. + train_dataset: Training features and labels. + val_dataset: Optional held-out validation set for early stopping. + progress_bar: Whether to display a progress bar during training. + epoch_callback: Optional callback invoked after each validation + step with ``(epoch, val_loss)``. Only called when + *val_dataset* is provided. + + Returns: + Epoch-level training metrics. """ - self._validate_fitted_feature_dimensions() - features = self.compute_features(dataset=dataset, labelled=False) - # Transform features with scaler - features_scaled = self.scaler.transform(features) - correct_probs = self.classifier.predict_proba(features_scaled) - dataset.metadata["calibrated_confidence"] = correct_probs[:, 1].tolist() + torch.manual_seed(self.seed) + device = self._resolve_device() + + input_dim = train_dataset.features.shape[1] + self.network = CalibratorNetwork( + input_dim=input_dim, + hidden_dims=self.hidden_dims, + dropout=self.dropout, + ).to(device) + + self.feature_mean = train_dataset.features.mean(dim=0).to(device) + self.feature_std = train_dataset.features.std(dim=0).clamp(min=1e-8).to(device) + + train_loader = DataLoader( + train_dataset, + batch_size=self.batch_size, + shuffle=True, + drop_last=False, + ) + + optimizer = torch.optim.Adam( + self.network.parameters(), + lr=self.learning_rate, + weight_decay=self.weight_decay, + ) + criterion = nn.BCEWithLogitsLoss() + + val_early: Optional[FeatureDataset] = None + val_was_subsampled = False + if val_dataset is not None: + val_early, val_was_subsampled = self._early_stopping_validation_split( + val_dataset + ) + + has_val = val_early is not None + history = TrainingHistory( + val_losses=[] if has_val else None, + val_accuracies=[] if has_val else None, + ) + best_val_loss = float("inf") + best_state: Optional[Dict[str, torch.Tensor]] = None + epochs_without_improvement = 0 + + for epoch in tqdm( + range(self.max_epochs), + disable=not progress_bar, + desc="Training calibrator", + unit="epoch", + ): + avg_train_loss = self._train_one_epoch( + train_loader, + optimizer, + criterion, + device, + ) + history.train_losses.append(avg_train_loss) + + if val_early is None: + history.best_epoch = epoch + continue + + should_stop, best_val_loss, best_state, epochs_without_improvement = ( + self._validation_step( + val_early, + criterion, + device, + history, + epoch, + best_val_loss, + best_state, + epochs_without_improvement, + ) + ) + + if epoch_callback is not None: + assert history.val_losses is not None + epoch_callback(epoch, history.val_losses[-1]) + + if should_stop: + break + + history.epochs_trained = len(history.train_losses) + + if best_state is not None: + self.network.load_state_dict(best_state) + self.network.to(device) + + self._record_full_validation_if_subsampled( + val_was_subsampled, + val_dataset, + history, + criterion, + device, + ) + + if self._feature_input_dim() != input_dim: + self._active_feature_columns = [ + f"feature_{i}" for i in range(input_dim - 1) + ] + + return history + + @torch.no_grad() + def _evaluate( + self, + dataset: FeatureDataset, + criterion: nn.Module, + device: torch.device, + ) -> Tuple[float, float]: + """Compute loss and accuracy on a dataset.""" + assert self.network is not None + assert self.feature_mean is not None + assert self.feature_std is not None + + self.network.eval() + loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False) + total_loss = 0.0 + correct = 0 + total = 0 + for features, labels in loader: + features = (features.to(device) - self.feature_mean) / self.feature_std + labels = labels.to(device) + logits = self.network(features) + total_loss += criterion(logits, labels).item() * len(labels) + preds = (logits > 0.0).float() + correct += (preds == labels).sum().item() + total += len(labels) + + return total_loss / max(total, 1), correct / max(total, 1) + + def _snapshot_state(self) -> Dict[str, torch.Tensor]: + """Return a CPU copy of the current network state dict.""" + assert self.network is not None + return {k: v.cpu().clone() for k, v in self.network.state_dict().items()} + + def _train_one_epoch( + self, + train_loader: DataLoader, + optimizer: torch.optim.Optimizer, + criterion: nn.Module, + device: torch.device, + ) -> float: + """Run one training epoch, return the average loss.""" + assert self.network is not None + assert self.feature_mean is not None + assert self.feature_std is not None + + self.network.train() + epoch_loss = 0.0 + n_batches = 0 + for features, labels in train_loader: + features = (features.to(device) - self.feature_mean) / self.feature_std + labels = labels.to(device) + optimizer.zero_grad() + logits = self.network(features) + loss = criterion(logits, labels) + loss.backward() + optimizer.step() + epoch_loss += loss.item() + n_batches += 1 + + return epoch_loss / max(n_batches, 1) + + def _validation_step( + self, + val_dataset: FeatureDataset, + criterion: nn.Module, + device: torch.device, + history: TrainingHistory, + epoch: int, + best_val_loss: float, + best_state: Optional[Dict[str, torch.Tensor]], + epochs_without_improvement: int, + ) -> Tuple[bool, float, Optional[Dict[str, torch.Tensor]], int]: + """Evaluate on validation set and check early stopping. + + Stops after ``n_iter_no_change`` consecutive epochs where validation loss + did not improve by at least ``tol`` over the running best (sklearn-style). + + Returns: + Tuple of (should_stop, best_val_loss, best_state, + epochs_without_improvement). + """ + val_loss, val_acc = self._evaluate(val_dataset, criterion, device) + + assert history.val_losses is not None + assert history.val_accuracies is not None + history.val_losses.append(val_loss) + history.val_accuracies.append(val_acc) + + # Match sklearn SGD early stopping: count epochs where loss did not beat + # ``best_val_loss - tol`` (see ``loss > best_loss - tol`` in sklearn). + if val_loss <= best_val_loss - self.tol: + best_val_loss = val_loss + best_state = self._snapshot_state() + history.best_epoch = epoch + epochs_without_improvement = 0 + else: + epochs_without_improvement += 1 + + if epochs_without_improvement >= self.n_iter_no_change: + return True, best_val_loss, best_state, epochs_without_improvement + + return False, best_val_loss, best_state, epochs_without_improvement def _feature_input_dim(self) -> int: """Return the number of model inputs (confidence plus feature columns).""" return 1 + len(self.columns) + def _validate_loaded_checkpoint(self, input_dim: int) -> None: + """Raise if saved weights and config feature metadata are inconsistent.""" + if self.feature_mean is None: + return + + trained_dim = int(self.feature_mean.shape[0]) + if trained_dim != input_dim: + raise ValueError( + "Calibrator checkpoint is corrupt: config input_dim " + f"({input_dim}) does not match saved normalisation stats " + f"({trained_dim} feature(s))." + ) + + if self._active_feature_columns is None: + return + + expected_dim = 1 + len(self._active_feature_columns) + if trained_dim == expected_dim: + return + + raise ValueError( + "Calibrator feature dimension mismatch: the saved model was trained " + f"with {trained_dim} input feature(s) but the checkpoint lists " + f"{len(self._active_feature_columns)} feature column(s) " + f"({self._active_feature_columns}). " + "This usually means the calibrator checkpoint is incompatible with " + "the installed version of winnow. Retrain the calibrator or use a " + "checkpoint published for this version." + ) + def _validate_fitted_feature_dimensions(self) -> None: - """Raise if the current feature set does not match the fitted scaler/classifier.""" - if not hasattr(self.scaler, "n_features_in_"): + """Raise if the current feature set does not match the fitted network.""" + if self.feature_mean is None: return - trained_dim = self.scaler.n_features_in_ + trained_dim = int(self.feature_mean.shape[0]) expected_dim = self._feature_input_dim() if trained_dim == expected_dim: return diff --git a/winnow/calibration/features/base.py b/winnow/calibration/features/base.py index 428335f8..8549d7f4 100644 --- a/winnow/calibration/features/base.py +++ b/winnow/calibration/features/base.py @@ -1,5 +1,6 @@ from abc import ABCMeta, abstractmethod -from typing import Dict, List, Optional +import inspect +from typing import Any, Dict, List, Optional import pandas as pd from sklearn.neural_network import MLPRegressor @@ -75,6 +76,20 @@ class and the name of the feature in the feature dataframe. str: The feature identifier """ + @staticmethod + def _to_plain(value: Any) -> Any: + """Convert values to plain JSON-serialisable Python types.""" + if isinstance(value, set): + return sorted(value) + try: + from omegaconf import DictConfig, ListConfig, OmegaConf + + if isinstance(value, (DictConfig, ListConfig)): + return OmegaConf.to_container(value, resolve=True) + except ImportError: + pass + return value + @abstractmethod def prepare(self, dataset: CalibrationDataset) -> None: """Prepare the dataset for feature computation. @@ -96,3 +111,27 @@ def compute(self, dataset: CalibrationDataset) -> None: dataset (CalibrationDataset): The dataset on which to compute the feature. """ pass + + def get_config(self) -> Dict[str, Any]: + """Return a serialisable dict of constructor parameters. + + The returned dict includes a ``_target_`` key with the fully-qualified + class path, plus every ``__init__`` parameter and its current value. + This is sufficient to reconstruct the feature instance. + + OmegaConf containers (``DictConfig``, ``ListConfig``) are resolved to + plain Python types so that the result is always JSON-serialisable. + + Returns: + Dict whose values are JSON-serialisable. + """ + sig = inspect.signature(self.__class__.__init__) + config: Dict[str, Any] = { + "_target_": (f"{self.__class__.__module__}.{self.__class__.__qualname__}"), + } + for param_name in sig.parameters: + if param_name == "self": + continue + if hasattr(self, param_name): + config[param_name] = self._to_plain(getattr(self, param_name)) + return config diff --git a/winnow/calibration/features/retention_time.py b/winnow/calibration/features/retention_time.py index ae9d2f70..df4cb123 100644 --- a/winnow/calibration/features/retention_time.py +++ b/winnow/calibration/features/retention_time.py @@ -1,12 +1,12 @@ -from typing import List, Optional +from typing import Dict, List, Optional, Set, Union +from sklearn.linear_model import LinearRegression +from pathlib import Path +import torch +from safetensors.torch import save_file, load_file import pandas as pd import numpy as np import warnings import koinapy -from typing import Dict, Set, Union -from pathlib import Path -import pickle -from sklearn.linear_model import LinearRegression from winnow.calibration.features.base import CalibrationFeatures, FeatureDependency from winnow.datasets.calibration_dataset import CalibrationDataset @@ -18,8 +18,8 @@ class RetentionTimeFeature(CalibrationFeatures): Uses a Koina iRT model to predict indexed retention times (iRT) for high-confidence peptides and trains a per-experiment linear regressor to map observed retention times - to predicted iRT values. Default behaviour is to always re-fit the regressor at both training and inference time, - but this can be overridden by passing a checkpoint file to `load_regressors` at inference time. + to predicted iRT values. The regressor is always re-fitted at both training and + inference time using self-supervised data (no database labels needed). """ def __init__( @@ -65,6 +65,7 @@ def __init__( self.max_peptide_length = max_peptide_length self.irt_predictors: Dict[str, LinearRegression] = {} self._loaded_experiment_names: Set[str] = set() + self._skipped_experiments: List[str] = [] @property def dependencies(self) -> List[FeatureDependency]: @@ -101,6 +102,18 @@ def columns(self) -> List[str]: columns.append("is_missing_irt_error") return columns + @staticmethod + def _sequence_key(row: pd.Series) -> tuple: + """Hashable key for a predicted peptide sequence.""" + prediction = row["prediction"] + if isinstance(prediction, list): + return tuple(prediction) + if "prediction_untokenised" in row.index and pd.notna( + row.get("prediction_untokenised") + ): + return (str(row["prediction_untokenised"]),) + return (str(prediction),) + def check_valid_irt_prediction(self, dataset: CalibrationDataset) -> pd.Series: """Check which predictions are valid for iRT prediction. @@ -137,37 +150,56 @@ def check_valid_irt_prediction(self, dataset: CalibrationDataset) -> pd.Series: return is_valid_irt_prediction def save_regressors(self, path: Union[Path, str]) -> None: - """Save fitted per-experiment regressors to a pickle file. + """Save fitted per-experiment regressors to a safetensors file. + + Each ``LinearRegression`` is stored as two tensors keyed by + ``{experiment_name}/coef`` and ``{experiment_name}/intercept``. Args: - path: File path for the output pickle. + path: File path for the output ``.safetensors`` file. """ path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "wb") as f: - pickle.dump(self.irt_predictors, f) + tensors: Dict[str, torch.Tensor] = {} + for exp_name, reg in self.irt_predictors.items(): + tensors[f"{exp_name}/coef"] = torch.as_tensor( + reg.coef_, dtype=torch.float64 + ) + tensors[f"{exp_name}/intercept"] = torch.as_tensor( + np.atleast_1d(reg.intercept_), dtype=torch.float64 + ) + save_file(tensors, path) def load_regressors(self, path: Union[Path, str]) -> None: - """Load per-experiment regressors from a pickle file. + """Load per-experiment regressors from a safetensors file. Loaded regressors are merged into ``self.irt_predictors`` and pinned so that subsequent ``prepare()`` calls skip re-fitting for those experiments. Args: - path: File path to the pickle containing saved regressors. + path: File path to the ``.safetensors`` file containing saved + regressors. """ - with open(Path(path), "rb") as f: - loaded: Dict[str, LinearRegression] = pickle.load(f) - self.irt_predictors.update(loaded) - self._loaded_experiment_names.update(loaded.keys()) + tensors = load_file(Path(path)) + experiments: Dict[str, Dict[str, torch.Tensor]] = {} + for key, tensor in tensors.items(): + exp_name, param = key.rsplit("/", 1) + experiments.setdefault(exp_name, {})[param] = tensor + + for exp_name, params in experiments.items(): + reg = LinearRegression() + reg.coef_ = params["coef"].numpy() + reg.intercept_ = params["intercept"].numpy().item() + self.irt_predictors[exp_name] = reg + self._loaded_experiment_names.update(experiments.keys()) def prepare(self, dataset: CalibrationDataset) -> None: """Fit per-experiment RT->iRT linear regressors. For each experiment in the dataset (identified by the ``experiment_name`` column), fits a ``LinearRegression`` on the top ``train_fraction`` of spectra by confidence. - Experiments pinned via ``load_regressors`` are skipped; all others are re-fitted - on every call. + Experiments that already have a regressor (e.g., loaded via ``load_regressors``) + are skipped. If ``experiment_name`` is absent, a single global regressor is fitted with a warning. @@ -176,8 +208,11 @@ def prepare(self, dataset: CalibrationDataset) -> None: dataset: The dataset containing peptide sequences and retention times. Raises: - ValueError: If ``retention_time`` column is missing, or if any experiment - has fewer than ``min_train_points`` valid training spectra. + ValueError: If ``retention_time`` column is missing from the dataset. + + Note: + Experiments with fewer than ``min_train_points`` valid training spectra + are skipped with a warning. Their iRT features will be imputed as zero. """ if "retention_time" not in dataset.metadata.columns: raise ValueError( @@ -185,32 +220,27 @@ def prepare(self, dataset: CalibrationDataset) -> None: "This is required for iRT features computation." ) - if "experiment_name" not in dataset.metadata.columns: - if "__global__" in self._loaded_experiment_names: - return + # Reset per dataset so directory-mode reuse does not carry stale skips + # from an earlier file into a later file with the same experiment_name. + self._skipped_experiments = [] + + experiments_to_fit = self._resolve_experiments_to_fit(dataset) + if not experiments_to_fit: + return + + per_exp_train = self._collect_training_data(experiments_to_fit) + if not per_exp_train: + return + + if "__global__" in per_exp_train: warnings.warn( - "No 'experiment_name' column found. Fitting a single global " - "RT->iRT regressor. For multi-experiment data, ensure each " - "spectrum file includes an 'experiment_name' column or use " - "MGF format (which derives it from the filename).", + "No 'experiment_name' column found. Fitting one global " + "RT->iRT regressor shared by all spectra.\n" + "For multi-experiment data, add an 'experiment_name' column or " + "use MGF format (experiment name derived from the filename) so " + "each run can have its own mapper.", stacklevel=2, ) - experiments_to_fit = {"__global__": dataset.metadata} - else: # Always re-fit for experiments if not loaded from checkpoint - experiments_to_fit = { - str(exp_name): group - for exp_name, group in dataset.metadata.groupby("experiment_name") - if str(exp_name) not in self._loaded_experiment_names - } - if not experiments_to_fit: - return - - # Select training data per experiment, validate counts, collect into - # a single DataFrame for one batched Koina call. - per_exp_train: Dict[str, pd.DataFrame] = {} - for exp_name, group in experiments_to_fit.items(): - train_data = self._select_training_data(group, exp_name) - per_exp_train[exp_name] = train_data all_train = pd.concat(per_exp_train.values(), ignore_index=True) @@ -235,7 +265,7 @@ def prepare(self, dataset: CalibrationDataset) -> None: y = train_data["iRT"].values regressor = LinearRegression() regressor.fit(x, y) - self.irt_predictors[str(exp_name)] = regressor + self.irt_predictors[exp_name] = regressor def compute(self, dataset: CalibrationDataset) -> None: """Compute the iRT error feature for each spectrum. @@ -247,41 +277,24 @@ def compute(self, dataset: CalibrationDataset) -> None: Args: dataset: The dataset containing peptide sequences and retention times. """ - # Check which predictions are valid for iRT prediction - is_valid_irt_prediction = self.check_valid_irt_prediction(dataset) - dataset.metadata["is_missing_irt_error"] = ~is_valid_irt_prediction + self._mark_missing_spectra(dataset) - if not self.learn_from_missing: - # Filter invalid entries from the dataset in place so that they are dropped entirely - # (not imputed with zeros) and downstream features also do not see them. - n_invalid = (~is_valid_irt_prediction).sum() - if n_invalid > 0: - warnings.warn( - f"Filtered {n_invalid} spectra that do not satisfy the validity constraints " - f"for the Koina iRT model '{self.irt_model_name}' " - f"(learn_from_missing=False). Constraints applied:\n" - f" - Retention time data required\n" - f" - max_peptide_length={self.max_peptide_length} residue tokens\n" - f" - unsupported_residues: {self.unsupported_residues[:3]}{'...' if len(self.unsupported_residues) > 3 else ''}\n" - f"Set learn_from_missing=True to impute missing features instead of filtering.", - stacklevel=2, - ) - _filtered = dataset.filter_entries( - metadata_predicate=lambda row: row["is_missing_irt_error"] - ) - dataset.metadata = _filtered.metadata - dataset.predictions = _filtered.predictions - # All remaining rows are valid — the reindex below will find every spectrum_id - # in predictions with no NaN fill needed. + if dataset.metadata.empty: + dataset.metadata["irt_error"] = pd.Series(dtype=float) + return original_indices = dataset.metadata.index - # Filter out invalid spectra for iRT prediction valid_irt_input = dataset.filter_entries( metadata_predicate=lambda row: row["is_missing_irt_error"] ) - # Prepare input data + if valid_irt_input.metadata.empty: + dataset.metadata["iRT"] = np.nan + dataset.metadata["predicted iRT"] = np.nan + dataset.metadata["irt_error"] = 0.0 + return + inputs = pd.DataFrame() inputs["peptide_sequences"] = np.array( [ @@ -295,18 +308,116 @@ def compute(self, dataset: CalibrationDataset) -> None: predictions = koina_model.predict(inputs) predictions["spectrum_id"] = predictions.index - # Match computed metadata to valid spectra and impute missing values for invalid spectra - # Reindex to match dataset.metadata.index and fill missing values with NaN dataset.metadata.index = dataset.metadata["spectrum_id"] dataset.metadata["iRT"] = predictions["irt"].reindex( dataset.metadata["spectrum_id"], fill_value=np.nan ) - # Apply per-experiment regressors + self._apply_regressors(dataset) + + # Revert to original indices + dataset.metadata.index = original_indices + + # Compute iRT error + # Set zeros for rows where "irt" is missing + dataset.metadata["irt_error"] = np.abs( + dataset.metadata["predicted iRT"] - dataset.metadata["iRT"] + ).fillna(0.0) + + def _resolve_experiments_to_fit( + self, dataset: CalibrationDataset + ) -> Optional[Dict[str, pd.DataFrame]]: + """Determine which experiments need a regressor fitted. + + Returns: + A dict mapping experiment name to its metadata subset, or None if + there is nothing to fit. + """ + if "experiment_name" not in dataset.metadata.columns: + if "__global__" in self._loaded_experiment_names: + return None + if "__global__" not in self.irt_predictors: + return {"__global__": dataset.metadata} + return None + + experiments_to_fit = { + str(exp_name): group + for exp_name, group in dataset.metadata.groupby("experiment_name") + if str(exp_name) not in self._loaded_experiment_names + } + return experiments_to_fit or None + + def _collect_training_data( + self, experiments_to_fit: Dict[str, pd.DataFrame] + ) -> Dict[str, pd.DataFrame]: + """Select training data per experiment, skipping those with insufficient data.""" + per_exp_train: Dict[str, pd.DataFrame] = {} + for exp_name, group in experiments_to_fit.items(): + try: + train_data = self._select_training_data(group, exp_name) + except ValueError as e: + outcome = ( + "dropped" if not self.learn_from_missing else "imputed as zero" + ) + if exp_name == "__global__": + header = "Skipping global RT->iRT regressor fit:" + else: + header = ( + f"Skipping RT->iRT regressor fit for experiment '{exp_name}':" + ) + warnings.warn( + f"{header}\n{e}\n" + f"Affected spectra will be {outcome} " + f"(learn_from_missing={self.learn_from_missing}).", + stacklevel=2, + ) + self._skipped_experiments.append(exp_name) + continue + per_exp_train[exp_name] = train_data + return per_exp_train + + def _mark_missing_spectra(self, dataset: CalibrationDataset) -> None: + """Mark spectra that cannot produce iRT features as missing. + + Sets the ``is_missing_irt_error`` column on ``dataset.metadata``. When + ``learn_from_missing`` is False, also drops those rows in place. + """ + is_valid = self.check_valid_irt_prediction(dataset) + dataset.metadata["is_missing_irt_error"] = ~is_valid + + if self._skipped_experiments: + if "experiment_name" in dataset.metadata.columns: + skipped_mask = dataset.metadata["experiment_name"].isin( + self._skipped_experiments + ) + dataset.metadata.loc[skipped_mask, "is_missing_irt_error"] = True + elif "__global__" in self._skipped_experiments: + dataset.metadata["is_missing_irt_error"] = True + + if not self.learn_from_missing: + n_invalid = dataset.metadata["is_missing_irt_error"].sum() + if n_invalid > 0: + warnings.warn( + f"Filtered {n_invalid} spectra that do not satisfy the validity " + f"constraints for the Koina iRT model '{self.irt_model_name}' " + f"(learn_from_missing=False).\n" + "Set learn_from_missing=True to impute missing features instead.", + stacklevel=3, + ) + _filtered = dataset.filter_entries( + metadata_predicate=lambda row: row["is_missing_irt_error"] + ) + dataset.metadata = _filtered.metadata + dataset.predictions = _filtered.predictions + + def _apply_regressors(self, dataset: CalibrationDataset) -> None: + """Apply per-experiment (or global) regressors to predict iRT from RT.""" if "experiment_name" in dataset.metadata.columns: predicted_irt = pd.Series(np.nan, index=dataset.metadata.index) for exp_name, group in dataset.metadata.groupby("experiment_name"): - regressor = self.irt_predictors[str(exp_name)] + regressor = self.irt_predictors.get(str(exp_name)) + if regressor is None: + continue predicted_irt.loc[group.index] = regressor.predict( group["retention_time"].values.reshape(-1, 1) ) @@ -316,14 +427,9 @@ def compute(self, dataset: CalibrationDataset) -> None: "__global__" ].predict(dataset.metadata["retention_time"].values.reshape(-1, 1)) - # Revert to original indices - dataset.metadata.index = original_indices - - # Compute iRT error - # Set zeros for rows where "iRT" is missing - dataset.metadata["irt_error"] = np.abs( - dataset.metadata["predicted iRT"] - dataset.metadata["iRT"] - ).fillna(0.0) + def _count_unique_sequences(self, metadata: pd.DataFrame) -> int: + """Count distinct predicted sequences in a metadata subset.""" + return len({self._sequence_key(row) for _, row in metadata.iterrows()}) def _select_training_data( self, metadata: pd.DataFrame, experiment_name: str @@ -359,14 +465,37 @@ def _select_training_data( n_train = max(1, int(self.train_fraction * len(train_data))) train_data = train_data.iloc[:n_train] + n_unique_peptides = self._count_unique_sequences(train_data) + if n_unique_peptides < 2: + raise ValueError( + "Cannot fit iRT calibration regressor —\n" + f" training pool has only {n_unique_peptides} unique peptide(s) after applying train_fraction={self.train_fraction}.\n" + " Koina iRT prediction models output one iRT per peptide sequence, so RT->iRT regression requires at least 2 distinct training peptides.\n" + " Increase calibrator.irt_calibration.train_fraction or provide more peptide diversity." + ) + + if 2 < n_unique_peptides < self.min_train_points: + warnings.warn( + f"Experiment '{experiment_name}': iRT calibration pool (top {self.train_fraction:.0%}, {len(train_data)} PSMs):\n" + f" Only {n_unique_peptides} unique peptide(s), below min_train_points={self.min_train_points}.\n" + f" The RT->iRT regressor fit may be unreliable.\n" + f" Consider increasing calibrator.irt_calibration.train_fraction or de-duplicating peptides.", + stacklevel=2, + ) + + if train_data["retention_time"].nunique() < 2: + raise ValueError( + "Cannot fit iRT calibration regressor —\n" + f" training pool has only {train_data['retention_time'].nunique()} unique retention time value(s) after applying train_fraction={self.train_fraction}.\n" + " Increase calibrator.irt_calibration.train_fraction." + ) + if len(train_data) < self.min_train_points: raise ValueError( - f"Experiment '{experiment_name}': insufficient data for iRT " - f"calibration. After applying train_fraction={self.train_fraction}, " - f"only {len(train_data)} valid training points remain " - f"(from {len(metadata)} total spectra), but " - f"min_train_points={self.min_train_points} are required. " - f"Adjust train_fraction, min_train_points, or provide more data." + "Insufficient data for iRT calibration.\n" + f" After applying train_fraction={self.train_fraction}, only {len(train_data)} valid training points remain (from {len(metadata)} total spectra),\n" + f" but min_train_points={self.min_train_points} are required.\n" + " Adjust train_fraction, min_train_points, or provide more data." ) return train_data @@ -377,6 +506,7 @@ def __getstate__(self) -> dict: state.pop("irt_predictors", None) state.pop("irt_predictor", None) state.pop("_loaded_experiment_names", None) + state.pop("_skipped_experiments", None) return state def __setstate__(self, state: dict) -> None: @@ -384,6 +514,7 @@ def __setstate__(self, state: dict) -> None: self.__dict__.update(state) self.irt_predictors = {} self._loaded_experiment_names = set() + self._skipped_experiments = [] if "min_train_points" not in self.__dict__: self.min_train_points = 10 warnings.warn("min_train_points not found in state, setting to 10") diff --git a/winnow/configs/calibrator.yaml b/winnow/configs/calibrator.yaml index f5414c02..7de66de8 100644 --- a/winnow/configs/calibrator.yaml +++ b/winnow/configs/calibrator.yaml @@ -1,15 +1,30 @@ # --- Calibrator configuration --- +defaults: + - koina calibrator: _target_: winnow.calibration.calibrator.ProbabilityCalibrator + # Network architecture + hidden_dims: [128, 64] + dropout: 0.1 + + # Training hyperparameters + learning_rate: 0.001 + weight_decay: 0.0001 + max_epochs: 100 + batch_size: 1024 + # Early stopping: stop after this many epochs in a row where + # validation loss did not improve by at least `tol` over the best so far. + n_iter_no_change: 10 + tol: 0.0001 seed: 42 - hidden_layer_sizes: [50, 50] # The number of neurons in each hidden layer of the MLP classifier. - learning_rate_init: 0.001 # The initial learning rate for the MLP classifier. - alpha: 0.0001 # L2 regularisation parameter for the MLP classifier. - max_iter: 1000 # Maximum number of training iterations for the MLP classifier. - early_stopping: true # Whether to use early stopping to terminate training. - validation_fraction: 0.1 # Proportion of training data to use for early stopping validation. + # If the validation set is larger than this, randomly subsample (with + # val_subsample_seed, defaulting to seed) for per-epoch early stopping only. + # After training, loss/accuracy are reported on the full validation set. + # Set to null to disable subsampling (use all validation PSMs each epoch). + val_early_stopping_max_psms: null + val_subsample_seed: null features: mass_error: @@ -37,46 +52,3 @@ calibrator: irt_model_name: ${koina.irt_model} # The name of the Koina iRT model to use. max_peptide_length: ${koina.constraints.max_peptide_length} # Maximum peptide length accepted by the Koina iRT model. unsupported_residues: ${koina.constraints.unsupported_residues} # Residues unsupported by the configured Koina iRT model. - -# Koina model configuration — shared settings for all intensity- and iRT-based features. -koina: - # Model names - intensity_model: Prosit_2025_intensity_22PTM - irt_model: Prosit_2025_irt_22PTM - - # Model inputs — applied to FragmentMatchFeatures and ChimericFeatures. - # To use a constant value tiled across all rows, specify it under input_constants. - # To use per-row values from a metadata column, add the column mapping under input_columns. - # Each input key must appear in at most one of these two dicts. - input_constants: - collision_energies: 27 - fragmentation_types: HCD - input_columns: {} - - # Model constraints — adjust to match the capabilities of your Koina models. - # See docs/configuration.md for guidance on choosing these values. - constraints: - max_precursor_charge: 6 # Maximum precursor charge accepted by the intensity models. - max_peptide_length: 30 # Maximum peptide length (residue token count) accepted by the intensity/iRT models. - # Residues unsupported by the configured Koina models. - # These residues must be specified using UNIMOD PTM IDs. - # Check that all PTMs unsupported by your selected Koina models but supported by Winnow are included. - unsupported_residues: - # Residue modifications (amino acid + modification) - # - "N[UNIMOD:7]" # Deamidated asparagine - # - "Q[UNIMOD:7]" # Deamidated glutamine - # - "R[UNIMOD:7]" # Arginine citrullination - # - "P[UNIMOD:35]" # Proline hydroxylation - # - "S[UNIMOD:21]" # Phosphorylated serine - # - "T[UNIMOD:21]" # Phosphorylated threonine - # - "Y[UNIMOD:21]" # Phosphorylated tyrosine - # - "C[UNIMOD:312]" # Cysteinylation - # - "E[UNIMOD:27]" # Pyro-glutamine - # - "Q[UNIMOD:28]" # Pyro-glutamine - # N-terminal modifications (standalone tokens) - # - "[UNIMOD:1]" # N-terminal acetylation - - "[UNIMOD:5]" # N-terminal carbamylation - - "[UNIMOD:385]" # N-terminal ammonia loss - - "(+25.98)" # Carbamylation & NH3 loss (legacy notation) - # Unsupported residues - # - "C" # Unmodified cysteine (must be explicitly carbamidomethylated for some Koina models) diff --git a/winnow/configs/compute_features.yaml b/winnow/configs/compute_features.yaml index a79ddb2d..b7603928 100644 --- a/winnow/configs/compute_features.yaml +++ b/winnow/configs/compute_features.yaml @@ -10,8 +10,8 @@ dataset: spectrum_path_or_directory: examples/example_data/spectra.mgf predictions_path: examples/example_data/predictions.csv -dataset_output_path: results/metadata.csv -# Drop rows with empty or non-list predictions (same helper as train/predict). -filter_empty_predictions: true +metadata_output_path: results/metadata.csv +# Optional: write a lean numeric Parquet for model training. +# training_matrix_output_path: results/training_matrix.parquet # If true, the dataset must include ground-truth sequence labels. labelled: true diff --git a/winnow/configs/diagnose_calibration.yaml b/winnow/configs/diagnose_calibration.yaml index c91ff037..56e5033e 100644 --- a/winnow/configs/diagnose_calibration.yaml +++ b/winnow/configs/diagnose_calibration.yaml @@ -2,6 +2,7 @@ defaults: - _self_ - residues + - koina - data_loader: instanovo dataset: diff --git a/winnow/configs/koina.yaml b/winnow/configs/koina.yaml new file mode 100644 index 00000000..a3c4da1f --- /dev/null +++ b/winnow/configs/koina.yaml @@ -0,0 +1,47 @@ +# --- Koina model configuration --- +# Shared by train, compute-features, predict, and diagnose-calibration. +# Intensity/iRT features interpolate ${koina.*} at train time; at inference the +# checkpoint restores feature params and only input_constants / input_columns +# are merged from this block (see winnow predict). + +koina: + # Model names + intensity_model: Prosit_2025_intensity_22PTM + irt_model: Prosit_2025_irt_22PTM + + # Model inputs — applied to FragmentMatchFeatures and ChimericFeatures. + # To use a constant value tiled across all rows, specify it under input_constants. + # To use per-row values from a metadata column, add the column mapping under input_columns. + # Each input key must appear in at most one of these two dicts. + input_constants: + collision_energies: 27 + fragmentation_types: HCD + input_columns: {} + + # Model constraints — adjust to match the capabilities of your Koina models. + # See docs/configuration.md for guidance on choosing these values. + constraints: + max_precursor_charge: 6 # Maximum precursor charge accepted by the intensity models. + max_peptide_length: 30 # Maximum peptide length (residue token count) accepted by the intensity/iRT models. + # Residues unsupported by the configured Koina models. + # These residues must be specified using UNIMOD PTM IDs. + # Check that all PTMs unsupported by your selected Koina models but supported by Winnow are included. + unsupported_residues: + # Residue modifications (amino acid + modification) + # - "N[UNIMOD:7]" # Deamidated asparagine + # - "Q[UNIMOD:7]" # Deamidated glutamine + # - "R[UNIMOD:7]" # Arginine citrullination + # - "P[UNIMOD:35]" # Proline hydroxylation + # - "S[UNIMOD:21]" # Phosphorylated serine + # - "T[UNIMOD:21]" # Phosphorylated threonine + # - "Y[UNIMOD:21]" # Phosphorylated tyrosine + # - "C[UNIMOD:312]" # Cysteinylation + # - "E[UNIMOD:27]" # Pyro-glutamine + # - "Q[UNIMOD:28]" # Pyro-glutamine + # N-terminal modifications (standalone tokens) + # - "[UNIMOD:1]" # N-terminal acetylation + - "[UNIMOD:5]" # N-terminal carbamylation + - "[UNIMOD:385]" # N-terminal ammonia loss + - "(+25.98)" # Carbamylation & NH3 loss (legacy notation) + # Unsupported residues + # - "C" # Unmodified cysteine (must be explicitly carbamidomethylated for some Koina models) diff --git a/winnow/configs/predict.yaml b/winnow/configs/predict.yaml index f7bd83ec..bfebfc90 100644 --- a/winnow/configs/predict.yaml +++ b/winnow/configs/predict.yaml @@ -2,6 +2,7 @@ defaults: - _self_ - residues + - koina - data_loader: instanovo # Options: instanovo, mztab, pointnovo, winnow - fdr_method: nonparametric # Options: nonparametric, database_grounded @@ -29,6 +30,10 @@ calibrator: # Useful for within-experiment use cases where the unlabelled data has an unreliable confidence distribution # for training a new RT->iRT regressor. irt_regressor_path: null + # Optional predict-time RT->iRT regressor fitting overrides (when irt_regressor_path is null): + # irt_calibration: + # train_fraction: 0.3 + # min_train_points: 10 fdr_control: # FDR settings: diff --git a/winnow/configs/train.yaml b/winnow/configs/train.yaml index 4ad37dd2..2732f4f6 100644 --- a/winnow/configs/train.yaml +++ b/winnow/configs/train.yaml @@ -7,8 +7,24 @@ defaults: # --- Pipeline Execution Configuration --- +# Two-phase training: set features_path to skip raw data loading and train +# directly from pre-computed feature Parquets (produced by compute-features). +# Accepts a single .parquet file or a directory of .parquet files. +# Leave as null for single-phase training (load raw data + compute features). +features_path: null + +# Explicit validation set: a Parquet file or directory of Parquets. +# When set, this overrides validation_fraction. +val_features_path: null + +# Automatic validation split fraction (used when val_features_path is null). +# WARNING: random splits may leak peptides or experiment artifacts between +# train and validation. For rigorous evaluation, hold out entire projects +# via val_features_path instead. +validation_fraction: 0.1 + +# Single-phase dataset config (ignored when features_path is set): dataset: - # Dataset paths: # Path to the spectrum data file or to folder containing saved internal Winnow dataset. spectrum_path_or_directory: examples/example_data/spectra.mgf # Path to the beam predictions file. @@ -24,3 +40,5 @@ dataset_output_path: results/calibrated_dataset.csv # from the same experiment(s), where the inference confidence distribution may be # unreliable for training a new RT->iRT regressor. irt_regressor_output_path: null +# Optional: save epoch-level training history as JSON for analysis and plotting. +training_history_path: null diff --git a/winnow/datasets/feature_dataset.py b/winnow/datasets/feature_dataset.py new file mode 100644 index 00000000..f674c48b --- /dev/null +++ b/winnow/datasets/feature_dataset.py @@ -0,0 +1,116 @@ +"""PyTorch Dataset wrapper for pre-computed calibration features.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Sequence + +import numpy as np +import polars as pl +import torch +from torch.utils.data import Dataset + +logger = logging.getLogger(__name__) + + +class FeatureDataset(Dataset): + """Wraps numpy feature/label arrays as a PyTorch Dataset. + + Each sample is a ``(features_tensor, label_tensor)`` pair. The + dataset can be constructed from in-memory arrays or loaded from + Parquet files via :meth:`from_parquet`. + + Args: + features: 2-D array of shape ``(n_samples, n_features)``. + labels: 1-D array of shape ``(n_samples,)``. + """ + + def __init__(self, features: np.ndarray, labels: np.ndarray) -> None: + if len(features) != len(labels): + raise ValueError( + f"features ({len(features)}) and labels ({len(labels)}) " + f"must have the same length" + ) + self.features = torch.as_tensor(features, dtype=torch.float32) + self.labels = torch.as_tensor(labels, dtype=torch.float32) + + @classmethod + def from_parquet( + cls, + path: str | Path, + feature_columns: Sequence[str] | None = None, + ) -> FeatureDataset: + """Load features from a single Parquet file or a directory of Parquets. + + If ``path`` is a directory, all ``*.parquet`` files inside it are + read and concatenated. The ``correct`` column is used as the + label. + + When ``feature_columns`` is provided, only those columns are used + as features (in the given order). Otherwise all numeric columns + except ``correct`` are used, which is only appropriate when the + Parquet contains exclusively feature columns. + + Args: + path: A ``.parquet`` file or a directory containing + ``*.parquet`` files. + feature_columns: Ordered list of column names to use as + features. If ``None``, all numeric columns (excluding + ``correct``) are used. + + Returns: + A new ``FeatureDataset`` instance. + + Raises: + FileNotFoundError: If no Parquet files are found at ``path``. + """ + path = Path(path) + if path.is_dir(): + parquet_files = sorted(path.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError(f"No .parquet files found in directory {path}") + df = pl.concat([pl.read_parquet(f) for f in parquet_files]) + else: + df = pl.read_parquet(path) + + if "correct" not in df.columns: + raise ValueError( + f"Parquet at {path} must contain a 'correct' column " + f"for labels. Found columns: {df.columns}" + ) + + labels = df["correct"].to_numpy().astype(np.float32) + + if feature_columns is not None: + missing = [c for c in feature_columns if c not in df.columns] + if missing: + raise ValueError( + f"Feature columns missing from Parquet: {missing}. " + f"Available: {df.columns}" + ) + features = df.select(feature_columns).to_numpy().astype(np.float32) + else: + import polars.selectors + + features = ( + df.drop("correct") + .select(polars.selectors.numeric()) + .to_numpy() + .astype(np.float32) + ) + logger.warning( + "No feature_columns specified; using all %d numeric columns " + "from %s. Pass feature_columns explicitly to avoid including " + "metadata columns.", + features.shape[1], + path, + ) + + return cls(features=features, labels=labels) + + def __len__(self) -> int: + return len(self.labels) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + return self.features[idx], self.labels[idx] diff --git a/winnow/scripts/main.py b/winnow/scripts/main.py index 56ad465a..71c63db0 100644 --- a/winnow/scripts/main.py +++ b/winnow/scripts/main.py @@ -10,11 +10,17 @@ from typing import Union, Tuple, Optional, List, TYPE_CHECKING, Annotated import typer import logging -from rich.logging import RichHandler +from pathlib import Path + +from winnow.utils.rich_console import STDERR_CONSOLE, notebook_safe_rich_handler + +import polars as pl +import pandas as pd +import numpy as np # Lazy imports for heavy dependencies - only imported when actually needed if TYPE_CHECKING: - import pandas as pd + from winnow.calibration.calibrator import ProbabilityCalibrator from winnow.datasets.calibration_dataset import CalibrationDataset from winnow.fdr.nonparametric import NonParametricFDRControl from winnow.fdr.database_grounded import DatabaseGroundedFDRControl @@ -25,7 +31,7 @@ # Prevent duplicate messages by disabling propagation and using only RichHandler logger.propagate = False if not logger.handlers: - logger.addHandler(RichHandler()) + logger.addHandler(notebook_safe_rich_handler()) # Typer CLI setup @@ -56,6 +62,112 @@ def print_config(cfg) -> None: formatter.print_config(cfg) +def _handle_koina_intensity_config( + cfg, + calibrator: Optional["ProbabilityCalibrator"] = None, + *, + hydra_overrides: Optional[List[str]] = None, + execute: bool = True, +) -> None: + """Validate and apply Koina intensity-model predict-time settings.""" + from winnow.utils.koina_intensity_config import ( + apply_koina_intensity_config, + validate_koina_intensity_config, + ) + + koina_cfg = cfg.get("koina") + validate_koina_intensity_config(koina_cfg, hydra_overrides) + + if not execute or calibrator is None: + return + + apply_koina_intensity_config(calibrator, koina_cfg, logger) + + +def _handle_irt_calibration_config( + cfg, + calibrator: Optional["ProbabilityCalibrator"] = None, + *, + hydra_overrides: Optional[List[str]] = None, + execute: bool = True, +) -> None: + """Validate and apply predict-time iRT regressor calibration overrides.""" + from winnow.utils.irt_calibration_config import ( + apply_irt_calibration_config, + validate_irt_calibration_config, + ) + + validate_irt_calibration_config(cfg, hydra_overrides) + + if not execute or calibrator is None: + return + + irt_calibration_cfg = cfg.calibrator.get("irt_calibration") + apply_irt_calibration_config( + calibrator, + irt_calibration_cfg, + hydra_overrides=hydra_overrides, + logger=logger, + ) + + +def _require_spectrum_path(spectrum_path_or_directory: Optional[str]) -> Path: + """Validate that a dataset path was supplied and return it as a :class:`Path`. + + Raises: + typer.Exit: If no path was provided (exit code 1). + """ + if ( + spectrum_path_or_directory is None + or not str(spectrum_path_or_directory).strip() + ): + STDERR_CONSOLE.print( + "[bold red]Error:[/bold red] No dataset supplied.\n\n" + "Set dataset.spectrum_path_or_directory to your spectrum " + "parquet file or a directory of internal Winnow datasets.\n\n" + "[bold cyan]Example:[/bold cyan]\n" + " [dim]winnow predict data_loader=winnow " + "dataset.spectrum_path_or_directory=/path/to/winnow_dataset[/dim]\n" + " [dim]winnow predict data_loader=instanovo " + "dataset.spectrum_path_or_directory=/path/to/data.parquet " + "dataset.predictions_path=/path/to/instanovo_predictions.csv[/dim]\n\n" + "Edit predict.yaml or pass overrides on the command line. " + "Run [dim]winnow config predict[/dim] to inspect the resolved config.", + ) + raise typer.Exit(code=1) + return Path(spectrum_path_or_directory) + + +def _require_spectra_after_features(n_spectra: int) -> None: + """Abort the CLI when feature computation removed every spectrum. + + Raises: + typer.Exit: If ``n_spectra`` is zero (exit code 1). + """ + if n_spectra > 0: + return + + lines = [ + "[bold red]Error:[/bold red] All spectra were removed during feature " + "computation; nothing left to calibrate.", + "", + "Check the warnings above. Common causes:", + " • iRT calibration skipped an experiment (e.g. only one " + "peptide in the top train_fraction pool, insufficient RT spread, or too " + "few calibration points) while the iRT feature has " + "learn_from_missing=False", + " • Koina validity filters (peptide length, precursor charge, " + "unsupported residues) with learn_from_missing=False on " + "intensity or iRT features", + "", + "Try increasing calibrator.irt_calibration.train_fraction, " + "setting learn_from_missing=True on affected features, or " + "fixing input data.", + ] + STDERR_CONSOLE.print("\n".join(lines)) + raise typer.Exit(code=1) + + def filter_dataset(dataset: CalibrationDataset) -> CalibrationDataset: """Filter out rows whose predictions are empty or contain unsupported PSMs. @@ -101,12 +213,23 @@ def apply_fdr_control( def check_if_labelled(dataset: CalibrationDataset) -> None: """Check if the dataset contains a ground-truth column.""" - if "sequence" not in dataset.metadata.columns: + if not _has_ground_truth_sequences(dataset.metadata): raise ValueError( "Database-grounded FDR control can only be performed on annotated data." ) +def _has_ground_truth_sequences(metadata: pd.DataFrame) -> bool: + """Return True when metadata contains at least one tokenised ground-truth sequence.""" + if "sequence" not in metadata.columns: + return False + return ( + metadata["sequence"] + .apply(lambda value: isinstance(value, list) and len(value) > 0) + .any() + ) + + def separate_metadata_and_predictions( dataset_metadata: pd.DataFrame, fdr_control: Union[NonParametricFDRControl, DatabaseGroundedFDRControl], @@ -134,10 +257,8 @@ def separate_metadata_and_predictions( # NonParametricFDRControl adds psm_pep column if isinstance(fdr_control, NonParametricFDRControl): preds_and_fdr_metrics_cols.append("psm_pep") - if "sequence" in dataset_metadata.columns: - preds_and_fdr_metrics_cols.append("sequence") - preds_and_fdr_metrics_cols.append("num_matches") - preds_and_fdr_metrics_cols.append("correct") + if _has_ground_truth_sequences(dataset_metadata): + preds_and_fdr_metrics_cols.extend(["sequence", "num_matches", "correct"]) dataset_preds_and_fdr_metrics = dataset_metadata[ ["spectrum_id"] + preds_and_fdr_metrics_cols ] @@ -150,22 +271,33 @@ def train_entry_point( execute: bool = True, config_dir: Optional[str] = None, ) -> None: - """The main training pipeline entry point. + """Train a calibrator from raw data or pre-computed feature Parquets. + + Two modes are supported, controlled by the presence of ``features_path`` + in the config: + + 1. **Single-phase** (``features_path`` not set): load raw spectra and + call ``calibrator.fit(CalibrationDataset)`` which computes features + and trains in one step. + 2. **Two-phase** (``features_path`` set): load pre-computed features from + Parquet file(s) into a ``FeatureDataset`` and call + ``calibrator.fit_from_features()``. + + In two-phase mode, validation data can be supplied explicitly via + ``val_features_path``, or split out automatically with + ``validation_fraction`` (default 0.1). Args: overrides: Optional list of config overrides. - execute: If False, only print the configuration and return without executing the pipeline. - config_dir: Optional path to custom config directory. If provided, configs in this - directory take precedence over package configs. Files not in custom dir will use package defaults (file-by-file resolution). + execute: If False, only print the configuration and return. + config_dir: Optional path to custom config directory. """ from hydra import initialize_config_dir, compose from hydra.utils import instantiate from winnow.utils.config_path import get_primary_config_dir - # Get primary config directory (custom if provided, otherwise package/dev) primary_config_dir = get_primary_config_dir(config_dir) - # Initialise Hydra with primary config directory with initialize_config_dir( config_dir=str(primary_config_dir), version_base="1.3", @@ -174,6 +306,7 @@ def train_entry_point( cfg = compose(config_name="train", overrides=overrides) if not execute: + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=False) print_config(cfg) return @@ -182,38 +315,249 @@ def train_entry_point( logger.info("Starting training pipeline.") logger.info(f"Training configuration: {cfg}") - # Load dataset - Hydra creates the DatasetLoader object + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=True) + + calibrator = instantiate(cfg.calibrator) + features_path = cfg.get("features_path") + + if features_path is not None: + train_dataset, val_dataset = _load_feature_datasets( + cfg, features_path, calibrator + ) + logger.info( + "Training on %d samples%s.", + len(train_dataset), + f" (val={len(val_dataset)})" if val_dataset is not None else "", + ) + history = calibrator.fit_from_features(train_dataset, val_dataset) + else: + annotated_dataset, dataset_output_path = _load_and_prepare_dataset( + cfg, instantiate + ) + # Compute features on the full dataset first so that + # annotated_dataset.metadata contains all feature columns when + # saved below. Splitting happens afterwards to avoid computing + # features twice (once per split). + calibrator.compute_features(annotated_dataset) + train_data, val_data = _maybe_split_calibration_dataset(cfg, annotated_dataset) + logger.info( + "Training on %d samples%s.", + len(train_data), + f" (val={len(val_data)})" if val_data is not None else "", + ) + history = _fit_from_calibration_datasets(calibrator, train_data, val_data) + if dataset_output_path: + logger.info(f"Saving training metadata to {dataset_output_path}") + annotated_dataset.save_metadata(dataset_output_path) + + logger.info(f"Saving model to {cfg.model_output_dir}") + ProbabilityCalibrator.save(calibrator, cfg.model_output_dir) + + _save_training_artifacts(cfg, calibrator, history) + + logger.info("Training pipeline completed successfully.") + + +def _load_feature_datasets(cfg, features_path, calibrator): + """Load pre-computed feature Parquets for two-phase training. + + Args: + cfg: Resolved Hydra config. + features_path: Path (file or directory) to training features. + calibrator: Instantiated calibrator whose ``columns`` define the + feature columns to select from the Parquet files. + + Returns: + Tuple of (train_dataset, val_dataset). val_dataset may be ``None``. + """ + from winnow.datasets.feature_dataset import FeatureDataset + from winnow.utils.paths import resolve_data_path + + feature_columns = ["confidence"] + calibrator.columns + + resolved = resolve_data_path(str(features_path)) + logger.info(f"Loading pre-computed features from {resolved}") + train_dataset = FeatureDataset.from_parquet(resolved, feature_columns) + + val_features_path = cfg.get("val_features_path") + if val_features_path is not None: + val_resolved = resolve_data_path(str(val_features_path)) + logger.info(f"Loading validation features from {val_resolved}") + val_dataset = FeatureDataset.from_parquet(val_resolved, feature_columns) + return train_dataset, val_dataset + + return _maybe_split_validation(cfg, train_dataset) + + +def _load_and_prepare_dataset(cfg, instantiate): + """Single-phase: load raw data and filter, ready for calibrator.fit(). + + Args: + cfg: Resolved Hydra config. + instantiate: Hydra's instantiate function. + + Returns: + Tuple of (dataset, dataset_output_path). ``dataset_output_path`` + may be ``None``. + """ logger.info("Loading dataset.") data_loader = instantiate(cfg.data_loader) - # Extract dataset loading parameters and convert to dict for flexible kwargs dataset_params = dict(cfg.dataset) - # Rename config keys to match the Protocol interface dataset_params["data_path"] = dataset_params.pop("spectrum_path_or_directory") dataset_params["predictions_path"] = dataset_params.pop("predictions_path", None) annotated_dataset = data_loader.load(**dataset_params) - logger.info(f"Loaded: {len(annotated_dataset.metadata)} spectra") logger.info("Filtering dataset for empty predictions.") annotated_dataset = filter_dataset(annotated_dataset) - logger.info(f"After filtering: {len(annotated_dataset.metadata)} spectra") - # Instantiate the calibrator from the config - logger.info("Instantiating calibrator from config.") - calibrator = instantiate(cfg.calibrator) + dataset_output_path = cfg.get("dataset_output_path") + return annotated_dataset, dataset_output_path - # Fit the calibrator to the dataset - logger.info("Fitting calibrator to dataset.") - calibrator.fit(annotated_dataset) - # Save the model - logger.info(f"Saving model to {cfg.model_output_dir}") - ProbabilityCalibrator.save(calibrator, cfg.model_output_dir) +def _random_validation_indices( + n: int, validation_fraction: float, seed: int +) -> Tuple[np.ndarray, np.ndarray]: + """Return shuffled (train_indices, val_indices) for a validation split.""" + n_val = max(1, int(n * validation_fraction)) + rng = np.random.default_rng(seed) + indices = rng.permutation(n) + return indices[: n - n_val], indices[n - n_val :] + + +def _maybe_split_validation(cfg, train_dataset): + """Optionally split a random validation set from the training data. + + Args: + cfg: Resolved Hydra config. + train_dataset: The full training FeatureDataset. + + Returns: + Tuple of (train_dataset, val_dataset). val_dataset is ``None`` when + ``validation_fraction`` is 0 or absent. + """ + from winnow.datasets.feature_dataset import FeatureDataset + + validation_fraction = cfg.get("validation_fraction", 0.1) + + if not validation_fraction or validation_fraction <= 0: + return train_dataset, None + + logger.warning( + "Using automatic validation split (%.0f%%). This performs a random " + "split that may leak peptides or experiment artifacts between train " + "and validation sets. For rigorous evaluation, provide separate " + "train/val Parquet files via features_path and val_features_path.", + validation_fraction * 100, + ) + + n = len(train_dataset) + seed = cfg.get("calibrator", {}).get("seed", 42) + train_idx, val_idx = _random_validation_indices(n, validation_fraction, seed) + + train_split = FeatureDataset( + features=train_dataset.features[train_idx].numpy(), + labels=train_dataset.labels[train_idx].numpy(), + ) + val_split = FeatureDataset( + features=train_dataset.features[val_idx].numpy(), + labels=train_dataset.labels[val_idx].numpy(), + ) + + return train_split, val_split + + +def _maybe_split_calibration_dataset(cfg, dataset): + """Optionally split a random validation set from a CalibrationDataset. + + Args: + cfg: Resolved Hydra config. + dataset: The full CalibrationDataset. + + Returns: + Tuple of (train_dataset, val_dataset). val_dataset is ``None`` when + ``validation_fraction`` is 0 or absent. + """ + from winnow.datasets.calibration_dataset import CalibrationDataset + + validation_fraction = cfg.get("validation_fraction", 0.1) + + if not validation_fraction or validation_fraction <= 0: + return dataset, None + + logger.warning( + "Using automatic validation split (%.0f%%). This performs a random " + "split that may leak peptides or experiment artifacts between train " + "and validation sets. For rigorous evaluation, provide separate " + "train/val Parquet files via features_path and val_features_path.", + validation_fraction * 100, + ) - # Save per-experiment iRT regressors if configured + n = len(dataset) + seed = cfg.get("calibrator", {}).get("seed", 42) + train_idx, val_idx = _random_validation_indices(n, validation_fraction, seed) + + train_meta = dataset.metadata.iloc[train_idx].reset_index(drop=True) + val_meta = dataset.metadata.iloc[val_idx].reset_index(drop=True) + + train_preds = None + val_preds = None + if dataset.predictions is not None: + train_preds = [dataset.predictions[i] for i in train_idx] + val_preds = [dataset.predictions[i] for i in val_idx] + + return ( + CalibrationDataset(metadata=train_meta, predictions=train_preds), + CalibrationDataset(metadata=val_meta, predictions=val_preds), + ) + + +def _fit_from_calibration_datasets(calibrator, train_data, val_data): + """Extract features from already-featurised CalibrationDatasets and train. + + This bridges the gap between single-phase feature computation (which + happens on the full dataset before splitting) and the calibrator's + ``fit_from_features`` method that expects ``FeatureDataset`` objects. + + Args: + calibrator: Instantiated calibrator with feature columns already + present in the datasets' metadata. + train_data: Training CalibrationDataset with feature columns. + val_data: Optional validation CalibrationDataset with feature columns, + or ``None``. + + Returns: + TrainingHistory from the calibrator. + """ + import numpy as np + from winnow.datasets.feature_dataset import FeatureDataset + + features, labels = calibrator._extract_feature_matrix(train_data, labelled=True) + train_fd = FeatureDataset(features=np.asarray(features), labels=np.asarray(labels)) + + val_fd = None + if val_data is not None: + val_features, val_labels = calibrator._extract_feature_matrix( + val_data, labelled=True + ) + val_fd = FeatureDataset( + features=np.asarray(val_features), labels=np.asarray(val_labels) + ) + + return calibrator.fit_from_features(train_fd, val_fd) + + +def _save_training_artifacts(cfg, calibrator, history): + """Save iRT regressors and training history if configured. + + Args: + cfg: Resolved Hydra config. + calibrator: The fitted calibrator. + history: TrainingHistory returned by fit(). + """ irt_regressor_output_path = cfg.get("irt_regressor_output_path") if irt_regressor_output_path: from winnow.calibration.calibration_features import RetentionTimeFeature @@ -223,12 +567,153 @@ def train_entry_point( logger.info(f"Saving iRT regressors to {irt_regressor_output_path}") rt_feature.save_regressors(irt_regressor_output_path) - # Save the training dataset results - logger.info(f"Final dataset: {len(annotated_dataset)} spectra") - logger.info(f"Saving training dataset results to {cfg.dataset_output_path}") - annotated_dataset.save_metadata(cfg.dataset_output_path) + training_history_path = cfg.get("training_history_path") + if training_history_path: + logger.info(f"Saving training history to {training_history_path}") + history.save(training_history_path) - logger.info("Training pipeline completed successfully.") + +def _discover_experiment_files(directory) -> list[Path]: + """Return sorted list of spectrum files in a directory. + + Args: + directory: Path to a directory containing spectrum files. + + Returns: + Sorted list of paths with supported extensions (.parquet, .ipc, .mgf). + + Raises: + FileNotFoundError: If the directory contains no supported files. + """ + directory = Path(directory) + supported = {".parquet", ".ipc", ".mgf"} + files = sorted(f for f in directory.iterdir() if f.suffix in supported) + if not files: + raise FileNotFoundError( + f"No spectrum files found in {directory}. " + f"Supported extensions: {', '.join(sorted(supported))}" + ) + return files + + +def _compute_features_directory( + spectrum_path, + predictions_path, + data_loader, + calibrator, + labelled: bool, +) -> list[pd.DataFrame]: + """Process a directory of experiment files one at a time.""" + experiment_files = _discover_experiment_files(spectrum_path) + logger.info( + f"Directory mode: found {len(experiment_files)} experiment file(s) " + f"in {spectrum_path}" + ) + all_metadata: list[pd.DataFrame] = [] + for file_path in experiment_files: + logger.info(f"Processing experiment file: {file_path.name}") + dataset = data_loader.load( + data_path=file_path, + predictions_path=predictions_path, + ) + logger.info(f" Loaded: {len(dataset.metadata)} spectra") + dataset = filter_dataset(dataset) + logger.info(f" After filtering: {len(dataset.metadata)} spectra") + if labelled and "sequence" not in dataset.metadata.columns: + raise ValueError( + f"Labelled dataset must contain a 'sequence' column " + f"(missing in {file_path.name})." + ) + calibrator.compute_features(dataset) + all_metadata.append(dataset.metadata) + logger.info( + f" {file_path.name}: {len(dataset.metadata)} spectra after features" + ) + return all_metadata + + +def _compute_features_single_file( + spectrum_path, + predictions_path, + data_loader, + calibrator, + labelled: bool, +) -> list[pd.DataFrame]: + """Process a single spectrum file.""" + logger.info(f"Single-file mode: {spectrum_path}") + dataset = data_loader.load( + data_path=spectrum_path, + predictions_path=predictions_path, + ) + logger.info(f"Loaded: {len(dataset.metadata)} spectra") + + if labelled and "sequence" not in dataset.metadata.columns: + raise ValueError( + "Labelled dataset must contain a 'sequence' column with " + "ground-truth sequences." + ) + + dataset = filter_dataset(dataset) + logger.info(f"After filtering: {len(dataset.metadata)} spectra") + + calibrator.compute_features(dataset) + return [dataset.metadata] + + +def _compute_features_batched_metadata( + spectrum_path: Path, + predictions_path, + data_loader, + calibrator, + labelled: bool, +) -> list[pd.DataFrame]: + """Run feature computation over a directory (one file at a time) or a single spectrum file.""" + if not spectrum_path.exists(): + raise FileNotFoundError(f"Spectrum path does not exist: {spectrum_path}") + if spectrum_path.is_dir(): + # Catch Winnow dataset directory + if (spectrum_path / "metadata.csv").is_file(): + return _compute_features_single_file( + spectrum_path, predictions_path, data_loader, calibrator, labelled + ) + return _compute_features_directory( + spectrum_path, + predictions_path, + data_loader, + calibrator, + labelled, + ) + return _compute_features_single_file( + spectrum_path, + predictions_path, + data_loader, + calibrator, + labelled, + ) + + +def _write_training_matrix(metadata, calibrator, confidence_column, output_path): + """Write a lean numeric training matrix to Parquet. + + Args: + metadata: Combined metadata DataFrame. + calibrator: The calibrator (used to get feature column names). + confidence_column: Name of the confidence column. + output_path: Destination Parquet path. + """ + feature_columns = [confidence_column] + feature_columns.extend(calibrator.columns) + keep_cols = list(feature_columns) + if "correct" in metadata.columns: + keep_cols.append("correct") + + training_df = pl.from_pandas(metadata[keep_cols]) + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + training_df.write_parquet(output_path) + logger.info( + f"Saved training matrix ({len(training_df)} rows, " + f"{len(training_df.columns)} cols) to {output_path}" + ) def compute_features_entry_point( @@ -236,9 +721,15 @@ def compute_features_entry_point( execute: bool = True, config_dir: Optional[str] = None, ) -> None: - """Load a dataset, compute calibration features into metadata, and save CSV. + """Load a dataset, compute calibration features into metadata, and save CSV or parquet. + + Supports two input modes via ``dataset.spectrum_path_or_directory``: - Does not fit the MLP or write a calibrator checkpoint; reuses ``calibrator.features`` from Hydra like training. + * **Directory**: each file in the directory is treated as a separate + experiment and processed one at a time to limit RAM usage. + * **Single file**: loaded in one shot. If the file contains an + ``experiment_name`` column the data is processed per-group; otherwise + the filename stem is used as the experiment name. Args: overrides: Optional list of config overrides. @@ -259,44 +750,53 @@ def compute_features_entry_point( cfg = compose(config_name="compute_features", overrides=overrides) if not execute: + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=False) print_config(cfg) return + spectrum_path = _require_spectrum_path(cfg.dataset.spectrum_path_or_directory) + + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=True) + logger.info("Starting compute-features pipeline.") logger.info(f"Compute-features configuration: {cfg}") labelled = bool(cfg.labelled) - - logger.info("Loading dataset.") data_loader = instantiate(cfg.data_loader) + calibrator = instantiate(cfg.calibrator) + predictions_path = cfg.dataset.get("predictions_path") + + all_metadata = _compute_features_batched_metadata( + spectrum_path, + predictions_path, + data_loader, + calibrator, + labelled, + ) - dataset_params = dict(cfg.dataset) - dataset_params["data_path"] = dataset_params.pop("spectrum_path_or_directory") - dataset_params["predictions_path"] = dataset_params.pop("predictions_path", None) + combined_metadata = pd.concat(all_metadata, ignore_index=True) + logger.info(f"Total spectra after feature computation: {len(combined_metadata)}") - dataset = data_loader.load(**dataset_params) + from winnow.datasets.calibration_dataset import CalibrationDataset - logger.info(f"Loaded: {len(dataset.metadata)} spectra") + combined_dataset = CalibrationDataset(metadata=combined_metadata) - if labelled and "sequence" not in dataset.metadata.columns: - raise ValueError( - "Labelled dataset must contain a 'sequence' column with ground-truth sequences." + metadata_output_path = cfg.get( + "metadata_output_path", cfg.get("dataset_output_path") + ) + if metadata_output_path: + logger.info(f"Saving metadata to {metadata_output_path}") + combined_dataset.save_metadata(metadata_output_path) + + training_matrix_output_path = cfg.get("training_matrix_output_path") + if training_matrix_output_path: + _write_training_matrix( + combined_metadata, + calibrator, + combined_dataset.confidence_column, + training_matrix_output_path, ) - if cfg.filter_empty_predictions: - logger.info("Filtering dataset for empty predictions.") - dataset = filter_dataset(dataset) - logger.info(f"After filtering: {len(dataset.metadata)} spectra") - - logger.info("Instantiating calibrator from config.") - calibrator = instantiate(cfg.calibrator) - - logger.info(f"Computing features with labelled={labelled}.") - calibrator.compute_features(dataset, labelled=labelled) - - logger.info(f"Saving dataset with features to {cfg.dataset_output_path}") - dataset.save_metadata(cfg.dataset_output_path) - logger.info("Compute-features pipeline completed successfully.") @@ -330,9 +830,16 @@ def predict_entry_point( cfg = compose(config_name="predict", overrides=overrides) if not execute: + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=False) + _handle_irt_calibration_config(cfg, hydra_overrides=overrides, execute=False) print_config(cfg) return + spectrum_path = _require_spectrum_path(cfg.dataset.spectrum_path_or_directory) + + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=True) + _handle_irt_calibration_config(cfg, hydra_overrides=overrides, execute=True) + from winnow.calibration.calibrator import ProbabilityCalibrator from winnow.datasets.calibration_dataset import CalibrationDataset from winnow.fdr.database_grounded import DatabaseGroundedFDRControl @@ -340,25 +847,6 @@ def predict_entry_point( logger.info("Starting prediction pipeline.") logger.info(f"Prediction configuration: {cfg}") - # Load dataset - Hydra creates the DatasetLoader object - logger.info("Loading dataset.") - data_loader = instantiate(cfg.data_loader) - - # Extract dataset loading parameters and convert to dict for flexible kwargs - dataset_params = dict(cfg.dataset) - # Rename config keys to match the Protocol interface - dataset_params["data_path"] = dataset_params.pop("spectrum_path_or_directory") - dataset_params["predictions_path"] = dataset_params.pop("predictions_path", None) - - dataset = data_loader.load(**dataset_params) - - logger.info(f"Loaded: {len(dataset.metadata)} spectra") - - logger.info("Filtering dataset for empty predictions.") - dataset = filter_dataset(dataset) - - logger.info(f"After filtering: {len(dataset.metadata)} spectra") - # Load trained calibrator logger.info("Loading trained calibrator.") calibrator = ProbabilityCalibrator.load( @@ -366,6 +854,13 @@ def predict_entry_point( cache_dir=cfg.calibrator.cache_dir, ) + _handle_koina_intensity_config( + cfg, calibrator, hydra_overrides=overrides, execute=True + ) + _handle_irt_calibration_config( + cfg, calibrator, hydra_overrides=overrides, execute=True + ) + # Load pre-fitted iRT regressors if configured irt_regressor_path = cfg.calibrator.get("irt_regressor_path") if irt_regressor_path: @@ -376,6 +871,26 @@ def predict_entry_point( logger.info(f"Loading iRT regressors from {irt_regressor_path}") rt_feature.load_regressors(irt_regressor_path) + # Load dataset - Hydra creates the DatasetLoader object + logger.info("Loading dataset.") + data_loader = instantiate(cfg.data_loader) + + predictions_path = cfg.dataset.get("predictions_path") + + all_metadata = _compute_features_batched_metadata( + spectrum_path, + predictions_path, + data_loader, + calibrator, + labelled=False, + ) + + combined_metadata = pd.concat(all_metadata, ignore_index=True) + logger.info(f"Total spectra after feature computation: {len(combined_metadata)}") + _require_spectra_after_features(len(combined_metadata)) + + dataset = CalibrationDataset(metadata=combined_metadata) + # Calibrate scores logger.info("Calibrating scores.") calibrator.predict(dataset) @@ -445,9 +960,14 @@ def diagnose_calibration_entry_point( cfg = compose(config_name="diagnose_calibration", overrides=overrides) if not execute: + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=False) + _handle_irt_calibration_config(cfg, hydra_overrides=overrides, execute=False) print_config(cfg) return + _handle_koina_intensity_config(cfg, hydra_overrides=overrides, execute=True) + _handle_irt_calibration_config(cfg, hydra_overrides=overrides, execute=True) + diagnostics = cfg.diagnostics label_column = ( diagnostics.label_column if diagnostics.label_column is not None else None @@ -481,6 +1001,15 @@ def diagnose_calibration_entry_point( pretrained_model_name_or_path=cfg.calibrator.pretrained_model_name_or_path, cache_dir=cfg.calibrator.cache_dir, ) + _handle_koina_intensity_config( + cfg, calibrator, hydra_overrides=overrides, execute=True + ) + _handle_irt_calibration_config( + cfg, calibrator, hydra_overrides=overrides, execute=True + ) + logger.info("Computing calibration features.") + calibrator.compute_features(dataset) + _require_spectra_after_features(len(dataset)) logger.info("Calibrating scores.") calibrator.predict(dataset) @@ -598,8 +1127,7 @@ def train( " [dim]winnow compute-features[/dim] # Uses config/compute_features.yaml\n\n" "[bold cyan]Override parameters:[/bold cyan]\n" " [dim]winnow compute-features dataset_output_path=results/my_features.csv[/dim]\n" - " [dim]winnow compute-features run_prepare=false[/dim] # Skip feature.prepare()\n" - " [dim]winnow compute-features filter_empty_predictions=false[/dim]\n\n" + " [dim]winnow compute-features run_prepare=false[/dim] # Skip feature.prepare()\n\n" "[bold cyan]Custom config directory:[/bold cyan]\n" " [dim]winnow compute-features --config-dir /path/to/configs[/dim]\n" " [dim]winnow compute-features -cp ./my_configs[/dim]\n\n" diff --git a/winnow/utils/hydra_overrides.py b/winnow/utils/hydra_overrides.py new file mode 100644 index 00000000..7a8c126b --- /dev/null +++ b/winnow/utils/hydra_overrides.py @@ -0,0 +1,18 @@ +"""Shared helpers for interpreting Hydra CLI override paths.""" + +from __future__ import annotations + +from typing import List, Optional, Set + + +def hydra_override_keys(overrides: Optional[List[str]]) -> Set[str]: + """Normalised Hydra override paths (without leading ``+`` / ``~``).""" + keys: Set[str] = set() + if not overrides: + return keys + for item in overrides: + if "=" not in item: + continue + path = item.split("=", 1)[0].lstrip("~+") + keys.add(path) + return keys diff --git a/winnow/utils/irt_calibration_config.py b/winnow/utils/irt_calibration_config.py new file mode 100644 index 00000000..b85cea38 --- /dev/null +++ b/winnow/utils/irt_calibration_config.py @@ -0,0 +1,163 @@ +"""Predict-time RT→iRT linear regressor calibration settings (RetentionTimeFeature).""" + +from __future__ import annotations + +from typing import Any, List, Optional, Set + +import typer + +from winnow.utils.hydra_overrides import hydra_override_keys + +IRT_CALIBRATION_PREFIX = "calibrator.irt_calibration." +IRT_CALIBRATION_FIELDS = ("train_fraction", "min_train_points") + + +def _irt_calibration_block(cfg: Any) -> Any: + calibrator_cfg = cfg.get("calibrator") + if calibrator_cfg is None: + return None + return calibrator_cfg.get("irt_calibration") + + +def explicit_irt_calibration_fields( + irt_calibration_cfg: Any, + hydra_overrides: Optional[List[str]] = None, +) -> Set[str]: + """Field names the user explicitly set via config or Hydra overrides.""" + explicit: Set[str] = set() + override_keys = hydra_override_keys(hydra_overrides) + for field in IRT_CALIBRATION_FIELDS: + hydra_path = f"{IRT_CALIBRATION_PREFIX}{field}" + if hydra_path in override_keys: + explicit.add(field) + continue + if ( + irt_calibration_cfg is not None + and irt_calibration_cfg.get(field) is not None + ): + explicit.add(field) + return explicit + + +def validate_irt_calibration_config( + cfg: Any, + hydra_overrides: Optional[List[str]] = None, +) -> None: + """Exit if ``irt_calibration`` overrides conflict with ``irt_regressor_path``.""" + calibrator_cfg = cfg.get("calibrator") or {} + irt_regressor_path = calibrator_cfg.get("irt_regressor_path") + if not irt_regressor_path: + return + + irt_calibration_cfg = _irt_calibration_block(cfg) + explicit = explicit_irt_calibration_fields(irt_calibration_cfg, hydra_overrides) + if not explicit: + return + + from winnow.utils.rich_console import STDERR_CONSOLE + + fields = ", ".join(f"calibrator.irt_calibration.{f}" for f in sorted(explicit)) + lines = [ + "[bold red]Error:[/bold red] Cannot set iRT calibration training parameters " + "when calibrator.irt_regressor_path is set.", + "", + f"Conflicting override(s): {fields}", + "", + "Pre-loaded regressors are used as-is. Unset " + "calibrator.irt_regressor_path to re-fit with custom " + "calibration settings, for example:", + " [dim]calibrator.irt_regressor_path=null " + "calibrator.irt_calibration.train_fraction=0.3[/dim]", + ] + STDERR_CONSOLE.print("\n".join(lines)) + raise typer.Exit(code=1) + + +def _get_retention_time_feature(calibrator: Any): + from winnow.calibration.features.retention_time import RetentionTimeFeature + + rt_feature = calibrator.feature_dict.get("iRT Feature") + if isinstance(rt_feature, RetentionTimeFeature): + return rt_feature + return None + + +def _resolved_irt_calibration( + irt_calibration_cfg: Any, + model_defaults: dict[str, Any], +) -> dict[str, Any]: + resolved = dict(model_defaults) + if irt_calibration_cfg is None: + return resolved + for field in IRT_CALIBRATION_FIELDS: + value = irt_calibration_cfg.get(field) + if value is not None: + resolved[field] = value + return resolved + + +def _log_irt_calibration_override( + field: str, + new_value: Any, + old_value: Any, + *, + explicit: Set[str], + logger: Any, +) -> None: + if field not in explicit or logger is None: + return + if new_value != old_value: + logger.warning( + "Overriding iRT regressor training setting '%s' for predict: " + "model default %s -> %s.", + field, + old_value, + new_value, + ) + return + logger.warning( + "Explicit iRT calibration override for '%s' matches the loaded model " + "default (%s).", + field, + old_value, + ) + + +def apply_irt_calibration_config( + calibrator: Any, + irt_calibration_cfg: Any, + hydra_overrides: Optional[List[str]] = None, + logger: Any = None, +) -> None: + """Apply optional predict-time overrides to RetentionTimeFeature calibration params.""" + rt_feature = _get_retention_time_feature(calibrator) + if rt_feature is None: + if logger is not None: + logger.debug( + "No RetentionTimeFeature on calibrator; skipping iRT calibration config." + ) + return + + explicit = explicit_irt_calibration_fields(irt_calibration_cfg, hydra_overrides) + model_defaults = { + "train_fraction": rt_feature.train_fraction, + "min_train_points": rt_feature.min_train_points, + } + resolved = _resolved_irt_calibration(irt_calibration_cfg, model_defaults) + + for field in IRT_CALIBRATION_FIELDS: + _log_irt_calibration_override( + field, + resolved[field], + model_defaults[field], + explicit=explicit, + logger=logger, + ) + setattr(rt_feature, field, resolved[field]) + + if logger is not None: + logger.info( + "iRT calibration settings: train_fraction=%s, min_train_points=%s", + rt_feature.train_fraction, + rt_feature.min_train_points, + ) diff --git a/winnow/utils/koina_intensity_config.py b/winnow/utils/koina_intensity_config.py new file mode 100644 index 00000000..f00a8853 --- /dev/null +++ b/winnow/utils/koina_intensity_config.py @@ -0,0 +1,209 @@ +"""Koina intensity-model runtime configuration for predict / train / compute-features. + +Collision energy and fragmentation type resolution applies to intensity Koina models +(``FragmentMatchFeatures``, ``ChimericFeatures``). +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Set, Tuple + +import typer + +from winnow.utils.hydra_overrides import hydra_override_keys + +KOINA_RUNTIME_CONFIG_KEYS = frozenset( + { + "model_input_constants", + "model_input_columns", + } +) + +KOINA_INPUT_KEYS = frozenset({"collision_energies", "fragmentation_types"}) + +DEFAULT_KOINA_INPUT_COLUMNS: Dict[str, str] = { + "collision_energies": "collision_energy", + "fragmentation_types": "frag_type", +} + + +def resolve_feature_model_inputs( + model_input_constants: Optional[Dict[str, Any]], + model_input_columns: Optional[Dict[str, str]], +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, str]]]: + """Resolve CE/frag sources for feature construction or checkpoint merge. + + Priority per key: non-null constant > non-null column > default column name. + """ + active_constants = _active_non_null(model_input_constants) + active_columns = _active_non_null(model_input_columns) + for key in list(active_columns): + if key in active_constants: + del active_columns[key] + for key, default_column in DEFAULT_KOINA_INPUT_COLUMNS.items(): + if key not in active_constants and key not in active_columns: + active_columns[key] = default_column + return ( + active_constants if active_constants else None, + active_columns if active_columns else None, + ) + + +def strip_runtime_keys_from_feature_config( + feature_config: Dict[str, Any], +) -> Dict[str, Any]: + """Return a copy of a saved feature config with runtime-only keys removed.""" + return { + key: value + for key, value in feature_config.items() + if key not in KOINA_RUNTIME_CONFIG_KEYS + } + + +def parse_koina_intensity_config( + koina_cfg: Any, +) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, str]]]: + """Extract input_constants / input_columns dicts from a Hydra koina config block.""" + if koina_cfg is None: + return None, None + from omegaconf import DictConfig, ListConfig, OmegaConf + + constants_cfg = koina_cfg.get("input_constants") + columns_cfg = koina_cfg.get("input_columns") + + def _to_plain_dict(value): + if value is None: + return None + if isinstance(value, (DictConfig, ListConfig)): + return OmegaConf.to_container(value, resolve=True) + return value + + constants = _to_plain_dict(constants_cfg) + columns = _to_plain_dict(columns_cfg) + return constants, columns + + +def _active_non_null(d: Optional[Dict[str, Any]]) -> Dict[str, Any]: + if not d: + return {} + return {k: v for k, v in d.items() if v is not None} + + +def user_specified_constant( + input_constants: Optional[Dict[str, Any]], key: str +) -> bool: + """True when ``input_constants[key]`` is explicitly non-null.""" + return _active_non_null(input_constants).get(key) is not None + + +def user_specified_column( + input_columns: Optional[Dict[str, str]], + key: str, + override_keys: Set[str], +) -> bool: + """True when the user intentionally set a column mapping for ``key``.""" + columns = _active_non_null(input_columns) + value = columns.get(key) + if value is None: + return False + if f"koina.input_columns.{key}" in override_keys: + return True + return value != DEFAULT_KOINA_INPUT_COLUMNS.get(key) + + +def validate_koina_intensity_config( + koina_cfg: Any, + hydra_overrides: Optional[List[str]] = None, +) -> None: + """Exit with a pretty error if CE/frag are set as both constant and column. + + Shipped default columns with null constants are not treated as a dual specification. + """ + constants, columns = parse_koina_intensity_config(koina_cfg) + override_keys = hydra_override_keys(hydra_overrides) + conflicts = [ + key + for key in sorted(KOINA_INPUT_KEYS) + if user_specified_constant(constants, key) + and user_specified_column(columns, key, override_keys) + ] + if not conflicts: + return + + from winnow.utils.rich_console import STDERR_CONSOLE + + lines = [ + "[bold red]Error:[/bold red] Koina model input(s) cannot be set in both " + "koina.input_constants and koina.input_columns:", + "", + ] + for key in conflicts: + lines.append(f" • {key}") + lines.extend( + [ + "", + "Use exactly one source per key, for example:", + " [dim]koina.input_constants.collision_energies=27[/dim]", + " [dim]koina.input_columns.collision_energies=collision_energy[/dim]", + ] + ) + STDERR_CONSOLE.print("\n".join(lines)) + raise typer.Exit(code=1) + + +def format_resolved_koina_setting( + key: str, + constants: Optional[Dict[str, Any]], + columns: Optional[Dict[str, str]], +) -> str: + """``key=value`` fragment for resolved Koina inputs (matches iRT settings log style).""" + active_constants = _active_non_null(constants) + active_columns = _active_non_null(columns) + if key in active_constants: + value = active_constants[key] + if isinstance(value, str): + return f"{key}={value!r}" + return f"{key}={value}" + if key in active_columns: + return f"{key}={active_columns[key]!r}" + default_col = DEFAULT_KOINA_INPUT_COLUMNS.get(key) + if default_col is not None: + return f"{key}={default_col!r}" + return f"{key}=unset" + + +def log_resolved_koina_intensity_config(calibrator: Any, logger: Any) -> None: + """Log how CE/frag inputs are resolved on intensity-based features.""" + for feature in calibrator.feature_dict.values(): + if not ( + hasattr(feature, "model_input_constants") + and hasattr(feature, "model_input_columns") + ): + continue + settings = ", ".join( + format_resolved_koina_setting( + key, feature.model_input_constants, feature.model_input_columns + ) + for key in sorted(KOINA_INPUT_KEYS) + ) + if not settings: + continue + logger.info("Koina input settings: %s", settings) + return + + +def apply_koina_intensity_config( + calibrator: Any, + koina_cfg: Any, + logger: Any, +) -> None: + """Apply CE/frag overrides from predict config to a calibrator.""" + if koina_cfg is None: + return + + constants, columns = parse_koina_intensity_config(koina_cfg) + calibrator.apply_koina_model_input_overrides( + model_input_constants=constants, + model_input_columns=columns, + ) + log_resolved_koina_intensity_config(calibrator, logger) diff --git a/winnow/utils/paths.py b/winnow/utils/paths.py new file mode 100644 index 00000000..7c59abb4 --- /dev/null +++ b/winnow/utils/paths.py @@ -0,0 +1,75 @@ +"""Utilities for resolving local paths or HuggingFace Hub repository IDs.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +from huggingface_hub import snapshot_download + +logger = logging.getLogger(__name__) + + +def resolve_data_path( + path_or_repo_id: str, + repo_type: str = "model", + cache_dir: Path | None = None, +) -> Path: + """Resolve a local path or HuggingFace Hub repo ID to a local path. + + Resolution order (logged at INFO level): + 1. Check if ``path_or_repo_id`` exists as a local path (file or dir). + 2. If not, attempt HuggingFace Hub download via ``snapshot_download()``. + 3. If both fail, raise with a clear error message. + + Warning: + If a local directory happens to match an HF repo ID pattern + (e.g. ``myorg/mydata`` exists locally), the local path takes + priority. This is logged as a warning. + + Args: + path_or_repo_id: A local file/directory path or a HuggingFace + repository identifier (e.g. ``"InstaDeepAI/winnow-features"``). + repo_type: HuggingFace repo type (``"model"``, ``"dataset"``, etc.). + cache_dir: Optional directory for caching HuggingFace downloads. + + Returns: + Resolved local ``Path``. + + Raises: + FileNotFoundError: If the path does not exist locally and cannot + be downloaded from HuggingFace Hub. + """ + local = Path(path_or_repo_id) + if local.exists(): + resolved = local.resolve() + if "/" in path_or_repo_id and not path_or_repo_id.startswith("/"): + logger.warning( + "Local path %s matches an HF repo ID pattern but exists " + "locally. Using local path. To force HF download, rename or " + "remove the local directory.", + resolved, + ) + logger.info("Resolved path locally: %s", resolved) + return resolved + + try: + logger.info( + "Path %s not found locally, downloading from HuggingFace Hub...", + path_or_repo_id, + ) + downloaded = Path( + snapshot_download( + repo_id=path_or_repo_id, + repo_type=repo_type, + cache_dir=str(cache_dir) if cache_dir else None, + ) + ) + logger.info("Downloaded to: %s", downloaded) + return downloaded + except Exception as exc: + raise FileNotFoundError( + f"Could not resolve '{path_or_repo_id}' as a local path or " + f"HuggingFace Hub repository (repo_type={repo_type!r}). " + f"HuggingFace error: {exc}" + ) from exc diff --git a/winnow/utils/rich_console.py b/winnow/utils/rich_console.py new file mode 100644 index 00000000..02ed2f1f --- /dev/null +++ b/winnow/utils/rich_console.py @@ -0,0 +1,19 @@ +"""Rich helpers that render in terminals and Jupyter notebooks.""" + +from __future__ import annotations + +from rich.console import Console +from rich.logging import RichHandler + +# Shared stderr console for CLI error messages. +STDERR_CONSOLE = Console(stderr=True) + + +def notebook_safe_rich_handler(**kwargs: object) -> RichHandler: + """Return a Rich log handler that uses ANSI colours, not OSC-8 hyperlinks. + + OSC-8 file links render as escape noise in many notebook frontends; plain + ANSI styling (including dimmed paths) still works there. + """ + kwargs.setdefault("enable_link_path", False) + return RichHandler(console=STDERR_CONSOLE, **kwargs)