Fit Honest leaves after building pruned tree

neurodata · Feb 7, 2025 · de2195c · de2195c
1 parent 2d31df5
commit de2195c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/treeple/tree/_honest_tree.py b/treeple/tree/_honest_tree.py
@@ -752,6 +752,12 @@ def _fit_leaves(self, X, y, sample_weight):
                 pruned_tree, self.tree_, pruner, X, y, sample_weight, missing_values_in_feature_mask
             )
             self.tree_ = pruned_tree
+
+            # Fit leaves using other subsample
+            honest_leaves = self.tree_.apply(X[self.honest_indices_])
+
+            # y-encoded ensures that y values match the indices of the classes
+            self._set_leaf_nodes(honest_leaves, y, sample_weight)
 
         if self.n_outputs_ == 1:
             self.n_classes_ = self.n_classes_[0]