feat: record training history

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
 

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

winnow/calibration/__init__.py

@@ -0,0 +1,3 @@
+from winnow.calibration.calibrator import ProbabilityCalibrator, TrainingHistory
+
+__all__ = ["ProbabilityCalibrator", "TrainingHistory"]