diff --git a/pearsonify/__init__.py b/pearsonify/__init__.py index 3e0058d..4d8fb73 100644 --- a/pearsonify/__init__.py +++ b/pearsonify/__init__.py @@ -1,3 +1,3 @@ __version__ = "1.0.1" -from .wrapper import Pearsonify \ No newline at end of file +from .wrapper import Pearsonify diff --git a/pearsonify/wrapper.py b/pearsonify/wrapper.py index b2573c5..ad9df74 100644 --- a/pearsonify/wrapper.py +++ b/pearsonify/wrapper.py @@ -1,10 +1,12 @@ -import numpy as np import matplotlib.pyplot as plt +import numpy as np from sklearn.base import BaseEstimator +from sklearn.utils.validation import NotFittedError, check_is_fitted + from .utils import ( - compute_pearson_residuals, - compute_confidence_intervals, calculate_coverage, + compute_confidence_intervals, + compute_pearson_residuals, ) @@ -24,7 +26,15 @@ def __init__(self, estimator: BaseEstimator, alpha=0.05): def fit(self, X_train, y_train, X_cal, y_cal): """Fit the model and compute Pearson residual-based quantile from calibration data.""" # Train the model if it's not already fitted - self.estimator.fit(X_train, y_train) + try: + check_is_fitted(self.estimator) + if not callable(getattr(self.estimator, "predict_proba", None)): + raise TypeError("The estimator must have a callable 'predict_proba' method.") + except TypeError as e: + raise TypeError(f"Estimator validation failed: {e}") from e + except NotFittedError: + # Attempt to fit the estimator if not already fitted + self.estimator.fit(X_train, y_train) # Compute residuals on calibration set y_cal_pred_proba = self.estimator.predict_proba(X_cal)[:, 1]