From eccef9f784a5d54365f863552e3a84753ee171df Mon Sep 17 00:00:00 2001 From: Sarp Tan Doven <119648987+sarptandoven@users.noreply.github.com> Date: Wed, 24 Jun 2026 02:28:31 -0700 Subject: [PATCH 1/2] fix: set classifier metadata --- ngboost/api.py | 97 ++++++++++++++++++++++++++++++++++++++++++++- tests/test_basic.py | 42 ++++++++++++++++++++ 2 files changed, 137 insertions(+), 2 deletions(-) diff --git a/ngboost/api.py b/ngboost/api.py index 8ca65f9..efde683 100644 --- a/ngboost/api.py +++ b/ngboost/api.py @@ -1,7 +1,8 @@ "The NGBoost library API" # pylint: disable=too-many-arguments -from sklearn.base import BaseEstimator +from sklearn.base import BaseEstimator, ClassifierMixin +from sklearn.preprocessing import LabelEncoder from sklearn.utils import check_array from ngboost.distns import ( @@ -122,7 +123,7 @@ def __setstate__(self, state_dict): super().__setstate__(state_dict) -class NGBClassifier(NGBoost, BaseEstimator): +class NGBClassifier(ClassifierMixin, NGBoost, BaseEstimator): """ Constructor for NGBoost classification models. @@ -154,6 +155,12 @@ class NGBClassifier(NGBoost, BaseEstimator): tol : numerical tolerance to be used in optimization random_state : seed for reproducibility. See https://stackoverflow.com/questions/28064634/random-state-pseudo-random-number-in-scikit-learn + validation_fraction: Proportion of training data to set + aside as validation data for early stopping. + early_stopping_rounds: The number of consecutive boosting iterations during which the + loss has to increase before the algorithm stops early. + Set to None to disable early stopping and validation. + None enables running over the full data set. Output: An NGBClassifier object that can be fit. """ @@ -173,6 +180,8 @@ def __init__( verbose_eval=100, tol=1e-4, random_state=None, + validation_fraction=0.1, + early_stopping_rounds=None, ): assert issubclass( Dist, ClassificationDistn @@ -190,9 +199,93 @@ def __init__( verbose_eval, tol, random_state, + validation_fraction, + early_stopping_rounds, ) self._estimator_type = "classifier" + def _encode_labels(self, Y): + return self._le.transform(Y) + + # pylint: disable=too-many-positional-arguments,attribute-defined-outside-init + def fit( + self, + X, + Y, + X_val=None, + Y_val=None, + sample_weight=None, + val_sample_weight=None, + train_loss_monitor=None, + val_loss_monitor=None, + early_stopping_rounds=None, + ): + self._le = LabelEncoder().fit(Y) + self.classes_ = self._le.classes_ + Y = self._encode_labels(Y) + if Y_val is not None: + Y_val = self._le.transform(Y_val) + self.base_models = [] + self.scalings = [] + self.col_idxs = [] + return NGBoost.partial_fit( + self, + X, + Y, + X_val=X_val, + Y_val=Y_val, + sample_weight=sample_weight, + val_sample_weight=val_sample_weight, + train_loss_monitor=train_loss_monitor, + val_loss_monitor=val_loss_monitor, + early_stopping_rounds=early_stopping_rounds, + ) + + # pylint: disable=too-many-positional-arguments,attribute-defined-outside-init + def partial_fit( + self, + X, + Y, + X_val=None, + Y_val=None, + sample_weight=None, + val_sample_weight=None, + train_loss_monitor=None, + val_loss_monitor=None, + early_stopping_rounds=None, + ): + if not hasattr(self, "classes_"): + self._le = LabelEncoder().fit(Y) + self.classes_ = self._le.classes_ + Y = self._encode_labels(Y) + if Y_val is not None: + Y_val = self._le.transform(Y_val) + return NGBoost.partial_fit( + self, + X, + Y, + X_val=X_val, + Y_val=Y_val, + sample_weight=sample_weight, + val_sample_weight=val_sample_weight, + train_loss_monitor=train_loss_monitor, + val_loss_monitor=val_loss_monitor, + early_stopping_rounds=early_stopping_rounds, + ) + + def predict(self, X, max_iter=None): + return self.classes_[super().predict(X, max_iter=max_iter)] + + def staged_predict(self, X, max_iter=None): + return [ + self.classes_[pred] for pred in super().staged_predict(X, max_iter=max_iter) + ] + + def score(self, X, Y): # for sklearn + return self.Manifold(self.pred_param(check_array(X)).T).total_score( + self._le.transform(Y) + ) + def predict_proba(self, X, max_iter=None): """ Probability prediction of Y at the points X=x diff --git a/tests/test_basic.py b/tests/test_basic.py index ddf4bce..526d2e4 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -20,6 +20,48 @@ def predict(self, X): return np.full(X.shape[0], self.prediction_) +def test_classifier_sets_sklearn_classes_and_encodes_labels(breast_cancer_data): + from sklearn.base import is_classifier # pylint: disable=import-outside-toplevel + from sklearn.metrics import ( # pylint: disable=import-outside-toplevel + RocCurveDisplay, + ) + from sklearn.model_selection import ( # pylint: disable=import-outside-toplevel + cross_val_score, + ) + + x_train, x_test, y_train, y_test = breast_cancer_data + y_labels = ["malignant" if y == 0 else "benign" for y in y_train] + y_test_labels = ["malignant" if y == 0 else "benign" for y in y_test] + + ngb = NGBClassifier( + Dist=Bernoulli, + n_estimators=2, + verbose=False, + random_state=0, + ) + + assert is_classifier(ngb) + assert isinstance(clone(ngb), NGBClassifier) + + ngb.fit(x_train, y_labels) + + assert list(ngb.classes_) == ["benign", "malignant"] + assert set(ngb.predict(x_test[:10])).issubset(set(ngb.classes_)) + assert ngb.predict_proba(x_test[:10]).shape == (10, 2) + assert len(ngb.staged_predict(x_test[:10])) == len(ngb.base_models) + assert cross_val_score(ngb, x_train, y_labels, scoring="roc_auc", cv=3).shape == ( + 3, + ) + + display = RocCurveDisplay.from_estimator( + ngb, + x_test, + y_test_labels, + pos_label="malignant", + ) + assert display.roc_auc >= 0.5 + + # TODO: This is non-deterministic in the model fitting def test_classification(breast_cancer_data): from sklearn.metrics import ( # pylint: disable=import-outside-toplevel From 332a41c5bd81a34fa15d567afe01a51504141d44 Mon Sep 17 00:00:00 2001 From: Sarp Tan Doven <119648987+sarptandoven@users.noreply.github.com> Date: Wed, 24 Jun 2026 02:45:23 -0700 Subject: [PATCH 2/2] fix: satisfy lint on ci --- ngboost/api.py | 1 + ngboost/helpers.py | 1 + 2 files changed, 2 insertions(+) diff --git a/ngboost/api.py b/ngboost/api.py index efde683..43ac3ac 100644 --- a/ngboost/api.py +++ b/ngboost/api.py @@ -123,6 +123,7 @@ def __setstate__(self, state_dict): super().__setstate__(state_dict) +# pylint: disable=duplicate-code class NGBClassifier(ClassifierMixin, NGBoost, BaseEstimator): """ Constructor for NGBoost classification models. diff --git a/ngboost/helpers.py b/ngboost/helpers.py index ea3750d..1a24ce2 100644 --- a/ngboost/helpers.py +++ b/ngboost/helpers.py @@ -14,6 +14,7 @@ _TREE_MODULE_SWAP_LOCK = threading.RLock() +# pylint: disable-next=c-extension-no-member class _CompatTree(_sklearn_tree.Tree): # pylint: disable=too-few-public-methods """Transient subclass of sklearn's Tree used only during loading.