Calibration API

The winnow.calibration module implements confidence calibration for peptide-spectrum matches using machine learning-based feature extraction and neural network classification.

Classes

ProbabilityCalibrator

The main calibration model that transforms raw confidence scores into calibrated probabilities using a PyTorch neural network (CalibratorNetwork) with various peptide and spectral features.

from winnow.calibration import ProbabilityCalibrator
from winnow.calibration.calibration_features import (
    MassErrorDaFeature, FragmentMatchFeatures, BeamFeatures
)
from winnow.datasets.calibration_dataset import CalibrationDataset
residue_masses = {
            "G": 57.021464,
            "A": 71.037114,
            "P": 97.052764,
            "E": 129.042593,
            "T": 101.047670,
            "I": 113.084064,
            "D": 115.026943,
            "R": 156.101111,
            "O": 237.147727,
            "N": 114.042927,
            "S": 87.032028,
            "M": 131.040485,
            "L": 113.084064,
        }

# Create and configure calibrator
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 directly from a labelled CalibrationDataset
calibrator.fit(train_dataset)

# Make predictions on new data
calibrator.predict(test_dataset)

# Save/load trained models (safetensors + config.json)
ProbabilityCalibrator.save(calibrator, Path("calibrator_checkpoint"))

# Load models - supports multiple sources
# 1. Load default pretrained model from Hugging Face
loaded_calibrator = ProbabilityCalibrator.load()

# 2. Load a custom Hugging Face model
loaded_calibrator = ProbabilityCalibrator.load("my-org/my-custom-model")

# 3. Load from local directory
loaded_calibrator = ProbabilityCalibrator.load("calibrator_checkpoint")

Key Features:

• 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/load using safetensors (weights) and config.json (architecture, normalisation stats, feature definitions)
• Two-phase Training: Supports training from pre-computed Parquet via FeatureDataset.from_parquet(), select_for(calibrator), and fit_from_features()
• GPU Support: Automatic GPU detection with CPU fallback during training; inference runs on CPU.

Main Methods:

• add_feature(feature): Add a calibration feature
• 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 FeatureDataset whose .columns match calibrator.columns
• set_training_feature_columns(columns): Optionally restrict the MLP to a subset of registered feature columns (compute still runs the full feature set)
• predict(dataset): Generate calibrated confidence scores
• 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
  • Hugging Face: Pass a repository ID string (e.g., "my-org/my-model")
  • Local: Pass a str or Path object pointing to a model directory
  • Models from Hugging Face are automatically cached in ~/.cache/huggingface/hub

Calibration Features

The calibrator uses a feature-based approach where multiple feature extractors compute signals from the peptide-spectrum match data. See the Calibration Features documentation for:

• The CalibrationFeatures base class for creating custom features
• Built-in features: Mass Error Features, Beam Features, Fragment Match Features, Chimeric Features, Retention Time Feature, Sequence Features, Token Score Features
• Feature dependencies and how they work
• Handling missing features (learn vs filter strategies)

Column schemas

Several APIs expose a .columns attribute which provide differing information.

Object	What it describes
CalibrationFeatures.columns	Which metadata fields this feature module writes when it runs compute()
ProbabilityCalibrator.columns	Which of those fields (plus confidence, prepended automatically) the MLP uses
FeatureDataset.columns	Which non-confidence fields are in the training matrix (layout is always [confidence, *columns])

By default, ProbabilityCalibrator.columns is the full set of columns from every registered feature. After a successful fit or load it freezes to the trained schema. Before fit you can narrow it with set_training_feature_columns([...]) or Hydra calibrator.training_feature_columns=[...] for ablations; all features will be computed and outputted, but only the chosen columns train the network.

For two-phase training from a pre-computed Parquet file, load with FeatureDataset.from_parquet, set the training subset on the calibrator if needed, then call select_for(calibrator) so the matrix matches calibrator.columns before fit_from_features. Names are checked when you set them against an already-built feature registry, and again at fit/extract time.

Workflow

Training workflow

1. Create Calibrator: Initialise ProbabilityCalibrator
2. Add Features: Use add_feature() to include desired calibration features
3. Optional subset: set_training_feature_columns([...]) to train on a subset of registry columns
4. Fit Model: Call fit() with a labelled CalibrationDataset
5. Save Model: Use save() to persist trained calibrator

For the two-phase workflow (compute features once, save a matrix, train later):

1. Call compute_features(dataset) to populate metadata columns
2. Export labelled Parquet containing at least confidence, the feature columns you may train on, and correct via the CLI training_matrix_output_path (lean) or a full metadata export (wide).
3. Reload with FeatureDataset.from_parquet(path) (loads all numeric/boolean feature columns; confidence is required and placed at index 0; FeatureDataset.columns lists the non-confidence names)
4. Optionally set calibrator.set_training_feature_columns([...]) (e.g. for ablations)
5. Align with train_dataset = wide.select_for(calibrator) so train_dataset.columns == list(calibrator.columns)
6. Train with fit_from_features(train_dataset, val_dataset=val_dataset)

Prediction workflow

1. Load Calibrator: Use load() to restore trained model from a Hugging Face repository or a local directory

   # Option 1: Use default pretrained model
   calibrator = ProbabilityCalibrator.load()
   
   # Option 2: Use custom Hugging Face model
   calibrator = ProbabilityCalibrator.load("my-org/my-custom-model")
   
   # Option 3: Use local model
   calibrator = ProbabilityCalibrator.load("./my_calibrator")
   

2. Predict: Call predict() with an unlabelled CalibrationDataset
3. Access Results: Calibrated scores stored in dataset's "calibrated_confidence" column

For detailed examples and usage patterns, refer to the examples notebook.