Implement loss decomposition metrics and wrappers for early stopping

.gitignore

@@ -164,5 +164,8 @@ cython_debug/
 .DS_Store
 .idea/
 
+# Cursor IDE
+.cursor/
+
 # Local issue drafts
 .github/ISSUES/

README.md

@@ -1,78 +1,101 @@
 # Splinator 📈
 
-**Probablistic Calibration with Regression Splines**
+**Probability Calibration for Python**
 
-[scikit-learn](https://scikit-learn.org) compatible
+A scikit-learn compatible toolkit for measuring and improving probability calibration.
 
-[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+[![PyPI version](https://img.shields.io/pypi/v/splinator)](https://pypi.org/project/splinator/)
+[![Downloads](https://static.pepy.tech/badge/splinator)](https://pepy.tech/project/splinator)
+[![Downloads/Month](https://static.pepy.tech/badge/splinator/month)](https://pepy.tech/project/splinator)
 [![Documentation Status](https://readthedocs.org/projects/splinator/badge/?version=latest)](https://splinator.readthedocs.io/en/latest/)
 [![Build](https://img.shields.io/github/actions/workflow/status/affirm/splinator/.github/workflows/python-package.yml)](https://github.com/affirm/splinator/actions)
 
 ## Installation
 
-`pip install splinator`
+```bash
+pip install splinator
+```
 
-## Algorithm
+## What's Inside
 
-Supported models:
+| Category | Components |
+|----------|------------|
+| **Calibrators** | `LinearSplineLogisticRegression` (piecewise), `TemperatureScaling` (single param) |
+| **Refinement Metrics** | `spline_refinement_loss`, `ts_refinement_loss` |
+| **Decomposition** | `logloss_decomposition`, `brier_decomposition` |
+| **Calibration Metrics** | ECE, Spiegelhalter's z |
 
-- Linear Spline Logistic Regression
+## Quick Start
 
-Supported metrics:
+```python
+from splinator import LinearSplineLogisticRegression, TemperatureScaling
 
-- Spiegelhalter’s z statistic
-- Expected Calibration Error (ECE)
+# Piecewise linear calibration (flexible, monotonic)
+spline = LinearSplineLogisticRegression(n_knots=10, monotonicity='increasing')
+spline.fit(scores.reshape(-1, 1), y_true)
+calibrated = spline.predict_proba(scores.reshape(-1, 1))[:, 1]
 
-\[1\] You can find more information in the [Linear Spline Logistic
-Regression](https://github.com/Affirm/splinator/wiki/Linear-Spline-Logistic-Regression).
+# Temperature scaling (simple, single parameter)
+ts = TemperatureScaling()
+ts.fit(probs.reshape(-1, 1), y_true)
+calibrated = ts.predict(probs.reshape(-1, 1))
+```
 
-\[2\] Additional readings
+## Calibration Metrics
 
-- Zhang, Jian, and Yiming Yang. [Probabilistic score estimation with
-    piecewise logistic
-    regression](https://pal.sri.com/wp-content/uploads/publications/radar/2004/icml04zhang.pdf).
-    Proceedings of the twenty-first international conference on Machine
-    learning. 2004.
-- Guo, Chuan, et al. "On calibration of modern neural networks." International conference on machine learning. PMLR, 2017.
+```python
+from splinator import (
+    expected_calibration_error,
+    spiegelhalters_z_statistic,
+    logloss_decomposition,      # Log loss → refinement + calibration
+    brier_decomposition,        # Brier score → refinement + calibration
+    spline_refinement_loss,     # Log loss after piecewise spline
+)
 
+# Assess calibration quality
+ece = expected_calibration_error(y_true, probs)
+z_stat = spiegelhalters_z_statistic(y_true, probs)
 
-## Examples
+# Decompose log loss into fixable vs irreducible parts
+decomp = logloss_decomposition(y_true, probs)
+print(f"Refinement (irreducible): {decomp['refinement_loss']:.4f}")
+print(f"Calibration (fixable):    {decomp['calibration_loss']:.4f}")
 
-| comparison                                     | notebook                                                                                                                                                           |
-|------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| scikit-learn's sigmoid and isotonic regression | [![colab1](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/Affirm/splinator/blob/main/examples/calibrator_model_comparison.ipynb)    |
-| pyGAM’s spline model                           | [![colab2](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/Affirm/splinator/blob/main/examples/spline_model_comparison.ipynb) |
+# Refinement using splinator's piecewise calibrator
+spline_ref = spline_refinement_loss(y_val, probs, n_knots=5)
+```
 
-## Development
+## XGBoost / LightGBM Integration
 
-The dependencies are managed by [uv](https://github.com/astral-sh/uv).
+Use calibration-aware metrics for early stopping:
 
-```bash
-# Install uv (if not already installed)
-curl -LsSf https://astral.sh/uv/install.sh | sh
+```python
+from splinator import ts_refinement_loss
+from splinator.metric_wrappers import make_metric_wrapper
 
-# Create virtual environment and install dependencies
-uv sync --dev
+metric = make_metric_wrapper(ts_refinement_loss, framework='xgboost')
+model = xgb.train(params, dtrain, custom_metric=metric, early_stopping_rounds=10, ...)
+```
 
-# Run tests
-uv run pytest tests -v
+## Examples
 
-# Run type checking
-uv run mypy src/splinator
-```
+| Notebook | Description |
+|----------|-------------|
+| [calibrator_model_comparison](examples/calibrator_model_comparison.ipynb) | Compare with sklearn calibrators |
+| [spline_model_comparison](examples/spline_model_comparison.ipynb) | Compare with pyGAM |
+| [ts_refinement_xgboost](examples/ts_refinement_xgboost.py) | Early stopping with refinement loss |
 
-## Example Usage
+## References
 
-``` python
-from splinator.estimators import LinearSplineLogisticRegression
-import numpy as np
+- Zhang, J. & Yang, Y. (2004). [Probabilistic score estimation with piecewise logistic regression](https://pal.sri.com/wp-content/uploads/publications/radar/2004/icml04zhang.pdf). ICML.
+- Guo, C., Pleiss, G., Sun, Y. & Weinberger, K. Q. (2017). [On calibration of modern neural networks](https://arxiv.org/abs/1706.04599). ICML.
+- Berta, E., Holzmüller, D., Jordan, M. I. & Bach, F. (2025). [Rethinking Early Stopping: Refine, Then Calibrate](https://arxiv.org/abs/2501.19195). arXiv:2501.19195.
 
-# random synthetic dataset
-n_samples = 100
-rng = np.random.RandomState(0)
-X = rng.normal(loc=100, size=(n_samples, 2))
-y = np.random.randint(2, size=n_samples)
+See also: [probmetrics](https://github.com/dholzmueller/probmetrics) (PyTorch calibration by the refinement paper authors)
 
-lslr = LinearSplineLogisticRegression(n_knots=10)
-lslr.fit(X, y)
+## Development
+
+```bash
+curl -LsSf https://astral.sh/uv/install.sh | sh  # Install uv
+uv sync --dev && uv run pytest tests -v          # Setup and test
 ```