set classifier metadata

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,8 @@ def __setstate__(self, state_dict):
         super().__setstate__(state_dict)
 
 
-class NGBClassifier(NGBoost, BaseEstimator):
+# pylint: disable=duplicate-code
+class NGBClassifier(ClassifierMixin, NGBoost, BaseEstimator):
     """
     Constructor for NGBoost classification models.
 
@@ -154,6 +156,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 +181,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 +200,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

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.
 

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