diff --git a/docs/api/calibration.md b/docs/api/calibration.md index 42a1a4c1..531ca258 100644 --- a/docs/api/calibration.md +++ b/docs/api/calibration.md @@ -24,8 +24,14 @@ calibrator.add_feature(MassErrorFeature(residue_masses=RESIDUE_MASSES)) calibrator.add_feature(FragmentMatchFeatures(mz_tolerance=0.02)) calibrator.add_feature(BeamFeatures()) -# Train the calibrator -calibrator.fit(training_dataset) +# Train the calibrator and get training history +training_history = calibrator.fit(training_dataset) +print(f"Final training loss: {training_history.final_training_loss:.6f}") +if training_history.final_validation_score is not None: + print(f"Final validation score: {training_history.final_validation_score:.6f}") + +# Plot training progress +training_history.plot(output_path="training_progress.png") # Make predictions calibrator.predict(test_dataset) @@ -55,15 +61,56 @@ loaded_calibrator = ProbabilityCalibrator.load("calibrator_checkpoint") **Main Methods:** - `add_feature(feature)`: Add a calibration feature -- `fit(dataset)`: Train the calibrator on a labelled dataset +- `fit(dataset)`: Train the calibrator on a labelled dataset. Returns a `TrainingHistory` object containing training metrics. - `predict(dataset)`: Generate calibrated confidence scores - `save(calibrator, path)`: Save trained model to disk - `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` + - 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` + +### TrainingHistory + +A dataclass containing training metrics from calibrator fitting. The training history is automatically saved to a JSON file during training (configurable via `training_history_path`). + +```python +from winnow.calibration import TrainingHistory + +# Returned from calibrator.fit() +training_history = calibrator.fit(training_dataset) + +# Access training metrics +print(training_history.loss_curve) # List of training loss at each iteration +print(training_history.validation_scores) # List of validation scores (if early_stopping=True) +print(training_history.final_training_loss) # Final training loss value +print(training_history.final_validation_score) # Final validation score (if early_stopping=True) +print(training_history.n_iter) # Number of training iterations + +# Save training history to JSON +training_history.save("training_history.json") + +# Load training history from JSON (e.g., for later analysis or plotting) +loaded_history = TrainingHistory.load("training_history.json") + +# Plot training progress +loaded_history.plot(output_path="training_plot.png") +``` + +**Attributes:** + +- `loss_curve`: List of training loss values at each iteration +- `validation_scores`: List of validation classification accuracy scores at each iteration (only if `early_stopping=True`) +- `final_training_loss`: The final training loss value +- `final_validation_score`: The final validation score (only if `early_stopping=True`) +- `n_iter`: Number of iterations the solver ran + +**Methods:** + +- `save(path)`: Save training history to a JSON file +- `load(path)`: Load training history from a JSON file (class method) +- `plot(output_path, show)`: Plot training and validation loss curves ### CalibrationFeatures @@ -245,6 +292,7 @@ rt_feat = RetentionTimeFeature(hidden_dim=10, train_fraction=0.1, learn_from_mis ### Prediction workflow 1. **Load Calibrator**: Use `load()` to restore trained model from a Hugging Face repository or a local directory + ```python # Option 1: Use default pretrained model calibrator = ProbabilityCalibrator.load() @@ -255,6 +303,7 @@ rt_feat = RetentionTimeFeature(hidden_dim=10, train_fraction=0.1, learn_from_mis # Option 3: Use local model calibrator = ProbabilityCalibrator.load("./my_calibrator") ``` + 2. **Predict**: Call `predict()` with unlabelled `CalibrationDataset` 3. **Access Results**: Calibrated scores stored in dataset's "calibrated_confidence" column diff --git a/tests/calibration/test_calibrator.py b/tests/calibration/test_calibrator.py index 5455b577..41f4e253 100644 --- a/tests/calibration/test_calibrator.py +++ b/tests/calibration/test_calibrator.py @@ -1,10 +1,12 @@ """Unit tests for winnow ProbabilityCalibrator.""" +import warnings import pickle import numpy as np import pandas as pd import pytest -from winnow.calibration.calibrator import ProbabilityCalibrator +from sklearn.exceptions import ConvergenceWarning +from winnow.calibration.calibrator import ProbabilityCalibrator, TrainingHistory from winnow.calibration.calibration_features import ( CalibrationFeatures, FeatureDependency, @@ -238,13 +240,25 @@ def test_fit(self, calibrator, labelled_dataset): feature = MockCalibrationFeature("test_feature", ["test_col"]) calibrator.add_feature(feature) - # Should not raise any exception - calibrator.fit(labelled_dataset) + # Should not raise any exception and return TrainingHistory + history = calibrator.fit(labelled_dataset) # Check that scaler and classifier were fitted (basic smoke test) assert hasattr(calibrator.scaler, "mean_") # Scaler fitted assert hasattr(calibrator.classifier, "classes_") # Classifier fitted + # Check that TrainingHistory is returned with expected attributes + assert isinstance(history, TrainingHistory) + assert isinstance(history.loss_curve, list) + assert len(history.loss_curve) > 0 + assert isinstance(history.final_training_loss, float) + assert history.n_iter > 0 + + # Early stopping is True by default, so validation scores should be available + assert history.validation_scores is not None + assert isinstance(history.validation_scores, list) + assert history.final_validation_score is not None + def test_predict_after_fit(self, calibrator, labelled_dataset, sample_dataset): """Test prediction after fitting.""" feature = MockCalibrationFeature("test_feature", ["test_col"]) @@ -323,3 +337,106 @@ def test_load_legacy_checkpoint_mlpclassifier_raises_error(self, tmp_path): 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_fit_without_early_stopping(self, labelled_dataset): + """Test fitting the calibrator without early stopping.""" + calibrator = ProbabilityCalibrator(seed=42, early_stopping=False, max_iter=10) + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + history = calibrator.fit(labelled_dataset) + + # Check that TrainingHistory is returned + assert isinstance(history, TrainingHistory) + assert isinstance(history.loss_curve, list) + assert len(history.loss_curve) > 0 + assert isinstance(history.final_training_loss, float) + + # Without early stopping, validation scores should be None + assert history.validation_scores is None + assert history.final_validation_score is None + + def test_plot_training_history(self, tmp_path, labelled_dataset): + """Test plotting training history.""" + calibrator = ProbabilityCalibrator(seed=42) + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + + history = calibrator.fit(labelled_dataset) + + # Test saving plot to file + plot_path = tmp_path / "training_plot.png" + history.plot(output_path=plot_path) + + # Check that the plot file was created + assert plot_path.exists() + assert plot_path.stat().st_size > 0 + + def test_plot_training_history_no_early_stopping(self, tmp_path, labelled_dataset): + """Test plotting training history without validation scores.""" + calibrator = ProbabilityCalibrator(seed=42, early_stopping=False, max_iter=10) + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + history = calibrator.fit(labelled_dataset) + + # Test saving plot to file (should work even without validation scores) + plot_path = tmp_path / "training_plot_no_val.png" + history.plot(output_path=plot_path) + + # Check that the plot file was created + assert plot_path.exists() + assert plot_path.stat().st_size > 0 + + def test_training_history_save_load(self, tmp_path, labelled_dataset): + """Test saving and loading training history.""" + calibrator = ProbabilityCalibrator(seed=42) + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + + history = calibrator.fit(labelled_dataset) + + # Save the history + history_path = tmp_path / "training_history.json" + history.save(history_path) + + # Check that the file was created + assert history_path.exists() + + # Load the history + loaded_history = TrainingHistory.load(history_path) + + # Check that the loaded history matches the original + assert loaded_history.loss_curve == history.loss_curve + assert loaded_history.validation_scores == history.validation_scores + assert loaded_history.final_training_loss == history.final_training_loss + assert loaded_history.final_validation_score == history.final_validation_score + assert loaded_history.n_iter == history.n_iter + + def test_training_history_save_load_no_early_stopping( + self, tmp_path, labelled_dataset + ): + """Test saving and loading training history without early stopping.""" + calibrator = ProbabilityCalibrator(seed=42, early_stopping=False, max_iter=10) + feature = MockCalibrationFeature("test_feature", ["test_col"]) + calibrator.add_feature(feature) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + history = calibrator.fit(labelled_dataset) + + # Save the history + history_path = tmp_path / "training_history_no_val.json" + history.save(history_path) + + # Load the history + loaded_history = TrainingHistory.load(history_path) + + # Check that validation_scores is None + assert loaded_history.validation_scores is None + assert loaded_history.final_validation_score is None + assert loaded_history.loss_curve == history.loss_curve diff --git a/winnow/calibration/__init__.py b/winnow/calibration/__init__.py index e69de29b..8d50ef3e 100644 --- a/winnow/calibration/__init__.py +++ b/winnow/calibration/__init__.py @@ -0,0 +1,3 @@ +from winnow.calibration.calibrator import ProbabilityCalibrator, TrainingHistory + +__all__ = ["ProbabilityCalibrator", "TrainingHistory"] diff --git a/winnow/calibration/calibrator.py b/winnow/calibration/calibrator.py index 9bccc544..abdf4385 100644 --- a/winnow/calibration/calibrator.py +++ b/winnow/calibration/calibrator.py @@ -2,6 +2,7 @@ from typing import Dict, List, Tuple, Union, Optional from pathlib import Path +from dataclasses import dataclass import pickle import numpy as np from sklearn.neural_network import MLPClassifier @@ -17,6 +18,148 @@ from winnow.datasets.calibration_dataset import CalibrationDataset +@dataclass +class TrainingHistory: + """Container for training history metrics from calibrator fitting. + + Attributes: + loss_curve: List of training loss values at each iteration. + validation_scores: List of validation scores at each iteration (only if early_stopping=True). + final_training_loss: The final training loss value. + final_validation_score: The final validation score (only if early_stopping=True). + n_iter: Number of iterations the solver ran. + """ + + loss_curve: List[float] + validation_scores: Optional[List[float]] + final_training_loss: float + final_validation_score: Optional[float] + n_iter: int + + def save(self, path: Union[Path, str]) -> None: + """Save the training history to a JSON file. + + Args: + path: Path to save the JSON file. + """ + import json + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + data = { + "loss_curve": self.loss_curve, + "validation_scores": self.validation_scores, + "final_training_loss": self.final_training_loss, + "final_validation_score": self.final_validation_score, + "n_iter": self.n_iter, + } + + with open(path, "w") as f: + json.dump(data, f, indent=2) + + @classmethod + def load(cls, path: Union[Path, str]) -> "TrainingHistory": + """Load training history from a JSON file. + + Args: + path: Path to the JSON file. + + Returns: + TrainingHistory: The loaded training history. + """ + import json + + with open(path) as f: + data = json.load(f) + + return cls( + loss_curve=data["loss_curve"], + validation_scores=data.get("validation_scores"), + final_training_loss=data["final_training_loss"], + final_validation_score=data.get("final_validation_score"), + n_iter=data["n_iter"], + ) + + def plot( + self, + output_path: Optional[Union[Path, str]] = None, + show: bool = False, + ) -> None: + """Plot the training and validation loss curves. + + Creates a visualization of the training progress showing the loss curve + and validation scores (if available) over training iterations. + + Args: + output_path (Optional[Union[Path, str]]): Path to save the plot image. + If None, the plot is not saved. Defaults to None. + show (bool): Whether to display the plot interactively. Defaults to False. + """ + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(figsize=(10, 6)) + + iterations = range(1, len(self.loss_curve) + 1) + ax.plot( + iterations, + self.loss_curve, + label="Training Loss", + color="#2563eb", + linewidth=2, + ) + + if self.validation_scores is not None: + # Validation scores are accuracy-like (higher is better), so we plot them on a secondary axis + ax2 = ax.twinx() + ax2.plot( + iterations, + self.validation_scores, + label="Validation Score", + color="#dc2626", + linewidth=2, + linestyle="--", + ) + ax2.set_ylabel("Validation Score", color="#dc2626", fontsize=12) + ax2.tick_params(axis="y", labelcolor="#dc2626") + ax2.legend(loc="upper right") + + ax.set_xlabel("Iteration", fontsize=12) + ax.set_ylabel("Training Loss", color="#2563eb", fontsize=12) + ax.tick_params(axis="y", labelcolor="#2563eb") + ax.set_title("Calibrator Training Progress", fontsize=14, fontweight="bold") + ax.legend(loc="upper left") + ax.grid(True, alpha=0.3) + + # Add text annotation with final metrics + final_text = f"Final Training Loss: {self.final_training_loss:.6f}" + if self.final_validation_score is not None: + final_text += f"\nFinal Validation Score: {self.final_validation_score:.6f}" + final_text += f"\nIterations: {self.n_iter}" + + ax.text( + 0.02, + 0.02, + final_text, + transform=ax.transAxes, + fontsize=10, + verticalalignment="bottom", + bbox={"boxstyle": "round", "facecolor": "white", "alpha": 0.8}, + ) + + plt.tight_layout() + + if output_path is not None: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, dpi=150, bbox_inches="tight") + + if show: + plt.show() + + plt.close(fig) + + class ProbabilityCalibrator: """A class for recalibrating probabilities for a de novo peptide sequencing method. @@ -209,19 +352,47 @@ 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: + def fit(self, dataset: CalibrationDataset) -> TrainingHistory: """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. + + Returns: + TrainingHistory: A dataclass containing training metrics including loss curves + and validation scores (if early_stopping is enabled). """ 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) + # Extract training history from the fitted classifier + loss_curve = list(self.classifier.loss_curve_) + final_training_loss = loss_curve[-1] if loss_curve else float("nan") + + # Validation scores are only available if early_stopping was enabled + validation_scores: Optional[List[float]] = None + final_validation_score: Optional[float] = None + if ( + hasattr(self.classifier, "validation_scores_") + and self.classifier.validation_scores_ + ): + validation_scores = list(self.classifier.validation_scores_) + final_validation_score = ( + validation_scores[-1] if validation_scores else None + ) + + return TrainingHistory( + loss_curve=loss_curve, + validation_scores=validation_scores, + final_training_loss=final_training_loss, + final_validation_score=final_validation_score, + n_iter=self.classifier.n_iter_, + ) + def compute_features( self, dataset: CalibrationDataset, labelled: bool ) -> Union[ diff --git a/winnow/configs/train.yaml b/winnow/configs/train.yaml index 839d3f89..33b9f086 100644 --- a/winnow/configs/train.yaml +++ b/winnow/configs/train.yaml @@ -19,3 +19,4 @@ dataset: # Output paths: model_output_dir: models/new_model dataset_output_path: results/calibrated_dataset.csv +training_history_path: results/training_history.json diff --git a/winnow/scripts/main.py b/winnow/scripts/main.py index a1d8d7d6..a4201a1b 100644 --- a/winnow/scripts/main.py +++ b/winnow/scripts/main.py @@ -207,7 +207,19 @@ def train_entry_point( # Fit the calibrator to the dataset logger.info("Fitting calibrator to dataset.") - calibrator.fit(annotated_dataset) + training_history = calibrator.fit(annotated_dataset) + + # Log final training metrics + logger.info(f"Training completed in {training_history.n_iter} iterations.") + logger.info(f"Final training loss: {training_history.final_training_loss:.6f}") + if training_history.final_validation_score is not None: + logger.info( + f"Final validation score: {training_history.final_validation_score:.6f}" + ) + + # Save training history + logger.info(f"Saving training history to {cfg.training_history_path}") + training_history.save(cfg.training_history_path) # Save the model logger.info(f"Saving model to {cfg.model_output_dir}")