Feat PyTorch calibrator and two-phase training pipeline

.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

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:
@@ -134,9 +148,9 @@ train-sample:
 	model_output_dir=models/new_model \
 	dataset_output_path=results/metadata.csv \
 	calibrator.features.retention_time_feature.train_fraction=0.3 \
-	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:

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))
 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

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
 
 # Configure individual features
 winnow train calibrator.features.fragment_match_features.mz_tolerance=0.01
@@ -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.ipc 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.ipc \
+    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.<name>` 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,16 +376,36 @@ 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
 # Train with custom calibrator settings
 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.hidden_dims=[128,64,32] \
+    calibrator.learning_rate=0.01 \
+    calibrator.max_epochs=200 \
     calibrator.features.fragment_match_features.mz_tolerance=0.01
 
 # Predict with database-grounded FDR
@@ -385,7 +421,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