Merge pull request #352 from neurodata/honestoblique_yuxin

YuxinB · web-flow · commit 75c2cf919939 · 2025-02-25T11:46:17.000-05:00
FIX attempt to correct pruning
diff --git a/treeple/tree/_honest_tree.py b/treeple/tree/_honest_tree.py
@@ -761,6 +761,7 @@ def _fit_leaves(self, X, y, sample_weight):
             )
             self.tree_ = pruned_tree
 
+            y = y_encoded
             # Fit leaves using other subsample
             honest_leaves = self.tree_.apply(X[self.honest_indices_])
 
@@ -883,6 +884,7 @@ class in a leaf.
 
         if self.n_outputs_ == 1:
             proba = proba[:, : self._tree_n_classes_]
+
             if not self.kernel_method:
                 normalizer = proba.sum(axis=1)[:, np.newaxis]
                 normalizer[normalizer == 0.0] = 1.0
@@ -896,11 +898,13 @@ class in a leaf.
 
             for k in range(self.n_outputs_):
                 proba_k = proba[:, k, : self._tree_n_classes_[k]]
+
                 if not self.kernel_method:
                     normalizer = proba_k.sum(axis=1)[:, np.newaxis]
                     normalizer[normalizer == 0.0] = 1.0
                     proba_k /= normalizer
                     proba_k = self._empty_leaf_correction(proba_k, k)
+
                 all_proba.append(proba_k)
 
             return all_proba
diff --git a/treeple/tree/honesty/_honest_prune.pyx b/treeple/tree/honesty/_honest_prune.pyx
@@ -110,7 +110,7 @@ cdef class HonestPruner(Splitter):
             The original tree to be pruned.
         """
         self.tree = orig_tree
-        self.capacity = 0
+        self.capacity = 2047
 
     cdef int init(
         self,
@@ -133,7 +133,6 @@ cdef class HonestPruner(Splitter):
         right-most end of `samples`, that is `samples[end_non_missing:end]`.
         """
         cdef float64_t threshold = self.tree.nodes[node_idx].threshold
-        cdef intp_t feature = self.tree.nodes[node_idx].feature
         cdef intp_t n_missing = 0
         cdef intp_t pos = self.start
         cdef intp_t p
@@ -147,10 +146,9 @@ cdef class HonestPruner(Splitter):
             sample_idx = self.samples[p]
 
             # missing-values are always placed at the right-most end
-            if isnan(X_ndarray[sample_idx, feature]):
+            if isnan(self.tree._compute_feature(X_ndarray, sample_idx, &self.tree.nodes[node_idx])):
                 self.samples[p], self.samples[current_end] = \
                     self.samples[current_end], self.samples[p]
-
                 n_missing += 1
                 current_end -= 1
 
