Merge pull request scikit-learn#4488 from glouppe/tree-apply

[MRG + 1] Public `apply` method for decision trees

Author: glouppe

doc/whats_new.rst

@@ -25,6 +25,10 @@ Bug fixes
 API changes summary
 -------------------
 
+    - :class:`tree.DecisionTreeClassifier` now exposes an ``apply`` method
+      for retrieving the leaf indices samples are predicted as. By
+      `Daniel Galvez`_ and `Gilles Louppe`_.
+
 .. _changes_0_16:
 
 0.16

sklearn/tree/tests/test_tree.py

@@ -29,6 +29,7 @@
 from sklearn.utils.testing import assert_true
 from sklearn.utils.testing import raises
 from sklearn.utils.validation import check_random_state
+from sklearn.utils.validation import NotFittedError
 
 from sklearn.tree import DecisionTreeClassifier
 from sklearn.tree import DecisionTreeRegressor
@@ -494,7 +495,7 @@ def test_error():
     for name, TreeEstimator in CLF_TREES.items():
         # predict before fit
         est = TreeEstimator()
-        assert_raises(Exception, est.predict_proba, X)
+        assert_raises(NotFittedError, est.predict_proba, X)
 
         est.fit(X, y)
         X2 = [-2, -1, 1]  # wrong feature shape for sample
@@ -527,7 +528,7 @@ def test_error():
 
         # predict before fitting
         est = TreeEstimator()
-        assert_raises(Exception, est.predict, T)
+        assert_raises(NotFittedError, est.predict, T)
 
         # predict on vector with different dims
         est.fit(X, y)
@@ -545,6 +546,10 @@ def test_error():
         clf.fit(X, y)
         assert_raises(ValueError, clf.predict, Xt)
 
+        # apply before fitting
+        est = TreeEstimator()
+        assert_raises(NotFittedError, est.apply, T)
+
 
 def test_min_samples_leaf():
     # Test if leaves contain more than leaf_count training examples
@@ -1208,6 +1213,8 @@ def check_explicit_sparse_zeros(tree, max_depth=3,
     Xs = (X_test, X_sparse_test)
     for X1, X2 in product(Xs, Xs):
         assert_array_almost_equal(s.tree_.apply(X1), d.tree_.apply(X2))
+        assert_array_almost_equal(s.apply(X1), d.apply(X2))
+        assert_array_almost_equal(s.apply(X1), s.tree_.apply(X1))
         assert_array_almost_equal(s.predict(X1), d.predict(X2))
 
         if tree in CLF_TREES:
@@ -1266,3 +1273,29 @@ def check_min_weight_leaf_split_level(name):
 def test_min_weight_leaf_split_level():
     for name in ALL_TREES:
         yield check_min_weight_leaf_split_level, name
+
+
+def check_public_apply(name):
+    X_small32 = X_small.astype(tree._tree.DTYPE)
+
+    est = ALL_TREES[name]()
+    est.fit(X_small, y_small)
+    assert_array_equal(est.apply(X_small),
+                       est.tree_.apply(X_small32))
+
+
+def check_public_apply_sparse(name):
+    X_small32 = csr_matrix(X_small.astype(tree._tree.DTYPE))
+
+    est = ALL_TREES[name]()
+    est.fit(X_small, y_small)
+    assert_array_equal(est.apply(X_small),
+                       est.tree_.apply(X_small32))
+
+
+def test_public_apply():
+    for name in ALL_TREES:
+        yield (check_public_apply, name)
+
+    for name in SPARSE_TREES:
+        yield (check_public_apply_sparse, name)

sklearn/tree/tree.py

@@ -370,6 +370,33 @@ def predict(self, X):
             else:
                 return proba[:, :, 0]
 
+    def apply(self, X):
+        """
+        Returns the index of the leaf that each sample is predicted as.
+
+        Parameters
+        ----------
+        X : array_like or sparse matrix, shape = [n_samples, n_features]
+            The input samples. Internally, it will be converted to
+            ``dtype=np.float32`` and if a sparse matrix is provided
+            to a sparse ``csr_matrix``.
+
+        Returns
+        -------
+        X_leaves : array_like, shape = [n_samples,]
+            For each datapoint x in X, return the index of the leaf x
+            ends up in. Leaves are numbered within
+            ``[0; self.tree_.node_count)``, possibly with gaps in the
+            numbering.
+        """
+        if self.tree_ is None:
+            raise NotFittedError("Estimator not fitted, "
+                                 "call `fit` before `apply`.")
+
+        X = check_array(X, dtype=DTYPE, accept_sparse="csr")
+
+        return self.tree_.apply(X)
+
     @property
     def feature_importances_(self):
         """Return the feature importances.