FIX revert to sample weight method

neurodata · Aug 22, 2023 · 6620f37 · 6620f37
1 parent 4a8f9c8
commit 6620f37
Showing 1 changed file with 6 additions and 10 deletions.
diff --git a/sktree/tree/_honest_tree.py b/sktree/tree/_honest_tree.py
@@ -420,13 +420,11 @@ def partial_fit(self, X, y, sample_weight=None, check_input=True, classes=None):
         )
         self.honest_indices_ = np.setdiff1d(nonzero_indices, self.structure_indices_)
 
-        _X = X[self.structure_indices_]
-        _y = y[self.structure_indices_]
-        _sample_weight = _sample_weight[self.structure_indices_]
+        _sample_weight[self.honest_indices_] = 0
 
         self.estimator_.partial_fit(
-            _X,
-            _y,
+            X,
+            y,
             sample_weight=_sample_weight,
             check_input=check_input,
             classes=classes if classes else np.unique(y),
@@ -535,9 +533,7 @@ def _fit(
         )
         self.honest_indices_ = np.setdiff1d(nonzero_indices, self.structure_indices_)
 
-        _X = X[self.structure_indices_]
-        _y = y[self.structure_indices_]
-        _sample_weight = _sample_weight[self.structure_indices_]
+        _sample_weight[self.honest_indices_] = 0
 
         if not self.tree_estimator:
             self.estimator_ = DecisionTreeClassifier(
@@ -562,8 +558,8 @@ def _fit(
 
         # Learn structure on subsample
         self.estimator_._fit(
-            _X,
-            _y,
+            X,
+            y,
             sample_weight=_sample_weight,
             check_input=check_input,
             missing_values_in_feature_mask=missing_values_in_feature_mask,