Demonstrate pickling is not an issue anymore

adam2392 · adam2392 · commit 8a720edd1f36 · 2022-05-12T12:04:10.000-04:00
diff --git a/sklearn/tree/_oblique_splitter.pyx b/sklearn/tree/_oblique_splitter.pyx
@@ -20,7 +20,6 @@ from libc.stdio cimport printf
 from libcpp.vector cimport vector
 
 from cython.operator cimport dereference as deref
-from cython.parallel import prange
 
 from ._utils cimport log
 from ._utils cimport rand_int
@@ -55,7 +54,7 @@ cdef class ObliqueSplitter(Splitter):
 
     def __cinit__(self, Criterion criterion, SIZE_t max_features,
                   SIZE_t min_samples_leaf, double min_weight_leaf,
-                  double feature_combinations, object random_state):
+                  double feature_combinations, object random_state, *argv):
         """
         Parameters
         ----------
diff --git a/sklearn/tree/_oblique_tree.pxd b/sklearn/tree/_oblique_tree.pxd
@@ -31,4 +31,5 @@ cdef class ObliqueTree(Tree):
     cdef DTYPE_t _compute_feature(self, const DTYPE_t[:] X_ndarray, Node *node, SIZE_t node_id) nogil
     cdef int _resize_c(self, SIZE_t capacity=*) nogil except -1
 
+    cpdef DTYPE_t compute_feature_value(self, object X, SIZE_t node_id)
     cpdef np.ndarray get_projection_matrix(self)
diff --git a/sklearn/tree/_oblique_tree.pyx b/sklearn/tree/_oblique_tree.pyx
@@ -244,6 +244,12 @@ cdef class ObliqueTree(Tree):
         self.proj_vec_weights[node_id] = deref(deref(oblique_split_node).proj_vec_weights)
         self.proj_vec_indices[node_id] = deref(deref(oblique_split_node).proj_vec_indices)
         return 1
+        
+    cpdef DTYPE_t compute_feature_value(self, object X, SIZE_t node_id):
+        cdef const DTYPE_t[:] X_vector = X
+        cdef Node* node = &self.nodes[node_id]
+        feature_value = self._compute_feature(X_vector, node, node_id)
+        return feature_value
 
     cdef DTYPE_t _compute_feature(self, const DTYPE_t[:] X_ndarray, Node *node, SIZE_t node_id) nogil:
         """Compute feature from a given data matrix, X.
@@ -254,11 +260,19 @@ cdef class ObliqueTree(Tree):
         cdef vector[DTYPE_t] proj_vec_weights
         cdef vector[SIZE_t] proj_vec_indices
         cdef DTYPE_t proj_feat = 0.0
+        cdef DTYPE_t weight = 0.0
+        cdef SIZE_t j = 0
+        cdef SIZE_t n_projections = proj_vec_indices.size()
 
         # compute projection of the data based on trained tree
         proj_vec_weights = self.proj_vec_weights[node_id]
         proj_vec_indices = self.proj_vec_indices[node_id]
-        for j in range(proj_vec_indices.size()):
-            proj_feat += X_ndarray[proj_vec_indices[j]] * proj_vec_weights[j]
+        for j in range(n_projections):
+            weight = proj_vec_weights[j]
+            
+            # skip a multiplication step if there is nothing to be done
+            if weight == 0:
+                continue
+            proj_feat += X_ndarray[proj_vec_indices[j]] * weight
 
         return proj_feat
diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx
@@ -766,15 +766,15 @@ cdef class Tree:
 
         self.capacity = capacity
         return 0
-    
+
     cdef int _set_node_values(self, SplitRecord* split_node,
             Node *node) nogil except -1:
         """Set node data.
         """
         node.feature = split_node.feature
         node.threshold = split_node.threshold
         return 1
-    
+
     cdef DTYPE_t _compute_feature(self, const DTYPE_t[:] X_ndarray,
             Node *node, SIZE_t node_id) nogil:
         """Compute feature from a given data matrix, X.
diff --git a/sklearn/tree/tests/test_tree.py b/sklearn/tree/tests/test_tree.py
@@ -998,7 +998,7 @@ def test_pickle():
         else:
             X, y = diabetes.data, diabetes.target
 
-        est = TreeEstimator(random_state=0, max_depth=5)
+        est = TreeEstimator(random_state=0)
         est.fit(X, y)
         score = est.score(X, y)
 
@@ -1032,9 +1032,7 @@ def test_pickle():
             est2_proj_mat = est2.tree_.get_projection_matrix()
             assert_array_equal(est_proj_mat, est2_proj_mat)
 
-        # TODO: this works when `max_depth=5`, but not 6?
-        # Must be some machine rounding error occurring
-        # probably needs ``_compute_feature`` to do some rounding?
+        # score should match before/after pickling
         score2 = est2.score(X, y)
         assert (
             score == score2
