record boundary info in nodes

Charlie-XIAO · Charlie-XIAO · commit cbeb2724cd5e · 2024-07-21T02:10:41.000+08:00
diff --git a/sklearn/ensemble/_gb.py b/sklearn/ensemble/_gb.py
@@ -140,7 +140,6 @@ def _update_terminal_regions(
     sample_mask,
     learning_rate=0.1,
     k=0,
-    line_search=True,
 ):
     """Update the leaf values to be predicted by the tree and raw_prediction.
 
@@ -183,14 +182,11 @@ def _update_terminal_regions(
          ``learning_rate``.
     k : int, default=0
         The index of the estimator being updated.
-    line_search : bool, default=True
-        Whether line search must be performed. Line search must not be
-        performed under monotonic constraints.
     """
     # compute leaf for each sample in ``X``.
     terminal_regions = tree.apply(X)
 
-    if line_search and not isinstance(loss, HalfSquaredError):
+    if not isinstance(loss, HalfSquaredError):
         # mask all which are not in sample mask.
         masked_terminal_regions = terminal_regions.copy()
         masked_terminal_regions[~sample_mask] = -1
@@ -262,6 +258,11 @@ def compute_update(y_, indices, neg_gradient, raw_prediction, k):
             sw = None if sample_weight is None else sample_weight[indices]
             update = compute_update(y_, indices, neg_gradient, raw_prediction, k)
 
+            if update > tree.upper_bound[leaf]:
+                update = tree.upper_bound[leaf]
+            elif update < tree.lower_bound[leaf]:
+                update = tree.lower_bound[leaf]
+
             # TODO: Multiply here by learning rate instead of everywhere else.
             tree.value[leaf, 0, 0] = update
 
@@ -515,7 +516,6 @@ def _fit_stage(
                 sample_mask,
                 learning_rate=self.learning_rate,
                 k=k,
-                line_search=self.monotonic_cst is None,
             )
 
             # add tree to ensemble
diff --git a/sklearn/tree/_tree.pxd b/sklearn/tree/_tree.pxd
@@ -22,7 +22,8 @@ cdef struct Node:
     intp_t n_node_samples                # Number of samples at the node
     float64_t weighted_n_node_samples    # Weighted number of samples at the node
     uint8_t missing_go_to_left     # Whether features have missing values
-
+    float64_t lower_bound          # Lower bound of the node's impurity
+    float64_t upper_bound          # Upper bound of the node's impurity
 
 cdef struct ParentInfo:
     # Structure to store information about the parent of a node
@@ -58,7 +59,9 @@ cdef class Tree:
                           intp_t feature, float64_t threshold, float64_t impurity,
                           intp_t n_node_samples,
                           float64_t weighted_n_node_samples,
-                          uint8_t missing_go_to_left) except -1 nogil
+                          uint8_t missing_go_to_left,
+                          float64_t lower_bound,
+                          float64_t upper_bound) except -1 nogil
     cdef int _resize(self, intp_t capacity) except -1 nogil
     cdef int _resize_c(self, intp_t capacity=*) except -1 nogil
 
diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx
@@ -268,7 +268,9 @@ cdef class DepthFirstTreeBuilder(TreeBuilder):
                 node_id = tree._add_node(parent, is_left, is_leaf, split.feature,
                                          split.threshold, parent_record.impurity,
                                          n_node_samples, weighted_n_node_samples,
-                                         split.missing_go_to_left)
+                                         split.missing_go_to_left,
+                                         parent_record.lower_bound,
+                                         parent_record.upper_bound)
 
                 if node_id == INTPTR_MAX:
                     rc = -1
@@ -626,7 +628,8 @@ cdef class BestFirstTreeBuilder(TreeBuilder):
                                  is_left, is_leaf,
                                  split.feature, split.threshold, parent_record.impurity,
                                  n_node_samples, weighted_n_node_samples,
-                                 split.missing_go_to_left)
+                                 split.missing_go_to_left,
+                                 parent_record.lower_bound, parent_record.upper_bound)
         if node_id == INTPTR_MAX:
             return -1
 
@@ -774,6 +777,14 @@ cdef class Tree:
     def missing_go_to_left(self):
         return self._get_node_ndarray()['missing_go_to_left'][:self.node_count]
 
+    @property
+    def lower_bound(self):
+        return self._get_node_ndarray()['lower_bound'][:self.node_count]
+
+    @property
+    def upper_bound(self):
+        return self._get_node_ndarray()['upper_bound'][:self.node_count]
+
     @property
     def value(self):
         return self._get_value_ndarray()[:self.node_count]
@@ -910,7 +921,9 @@ cdef class Tree:
                           intp_t feature, float64_t threshold, float64_t impurity,
                           intp_t n_node_samples,
                           float64_t weighted_n_node_samples,
-                          uint8_t missing_go_to_left) except -1 nogil:
+                          uint8_t missing_go_to_left,
+                          float64_t lower_bound,
+                          float64_t upper_bound) except -1 nogil:
         """Add a node to the tree.
 
         The new node registers itself as the child of its parent.
@@ -927,6 +940,8 @@ cdef class Tree:
         node.impurity = impurity
         node.n_node_samples = n_node_samples
         node.weighted_n_node_samples = weighted_n_node_samples
+        node.lower_bound = lower_bound
+        node.upper_bound = upper_bound
 
         if parent != _TREE_UNDEFINED:
             if is_left:
@@ -1934,7 +1949,8 @@ cdef _build_pruned_tree(
             new_node_id = tree._add_node(
                 parent, is_left, is_leaf, node.feature, node.threshold,
                 node.impurity, node.n_node_samples,
-                node.weighted_n_node_samples, node.missing_go_to_left)
+                node.weighted_n_node_samples, node.missing_go_to_left,
+                node.lower_bound, node.upper_bound)
 
             if new_node_id == INTPTR_MAX:
                 rc = -1
