add tag-based branching to thresholding scorer

amueller · amueller · commit 306d788c39af · 2015-03-19T16:36:46.000-04:00
diff --git a/sklearn/metrics/scorer.py b/sklearn/metrics/scorer.py
@@ -157,20 +157,23 @@ def __call__(self, clf, X, y, sample_weight=None):
         if y_type not in ("binary", "multilabel-indicator"):
             raise ValueError("{0} format is not supported".format(y_type))
 
-        try:
-            y_pred = clf.decision_function(X)
+        if getattr(clf, "estimator_type", None) == "regressor":
+            y_pred = clf.predict(X)
+        else:
+            try:
+                y_pred = clf.decision_function(X)
 
-            # For multi-output multi-class estimator
-            if isinstance(y_pred, list):
-                y_pred = np.vstack(p for p in y_pred).T
+                # For multi-output multi-class estimator
+                if isinstance(y_pred, list):
+                    y_pred = np.vstack(p for p in y_pred).T
 
-        except (NotImplementedError, AttributeError):
-            y_pred = clf.predict_proba(X)
+            except (NotImplementedError, AttributeError):
+                y_pred = clf.predict_proba(X)
 
-            if y_type == "binary":
-                y_pred = y_pred[:, 1]
-            elif isinstance(y_pred, list):
-                y_pred = np.vstack([p[:, -1] for p in y_pred]).T
+                if y_type == "binary":
+                    y_pred = y_pred[:, 1]
+                elif isinstance(y_pred, list):
+                    y_pred = np.vstack([p[:, -1] for p in y_pred]).T
 
         if sample_weight is not None:
             return self._sign * self._score_func(y, y_pred,
diff --git a/sklearn/metrics/tests/test_score_objects.py b/sklearn/metrics/tests/test_score_objects.py
@@ -22,7 +22,7 @@
 from sklearn.cluster import KMeans
 from sklearn.dummy import DummyRegressor
 from sklearn.linear_model import Ridge, LogisticRegression
-from sklearn.tree import DecisionTreeClassifier
+from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 from sklearn.datasets import make_blobs
 from sklearn.datasets import make_classification
 from sklearn.datasets import make_multilabel_classification
@@ -219,6 +219,13 @@ def test_thresholded_scorers():
     score2 = roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1])
     assert_almost_equal(score1, score2)
 
+    # test with a regressor (no decision_function)
+    reg = DecisionTreeRegressor()
+    reg.fit(X_train, y_train)
+    score1 = get_scorer('roc_auc')(reg, X_test, y_test)
+    score2 = roc_auc_score(y_test, reg.predict(X_test))
+    assert_almost_equal(score1, score2)
+
     # Test that an exception is raised on more than two classes
     X, y = make_blobs(random_state=0, centers=3)
     X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
