fix(placement): rank CART splits by impurity decrease before spacing

pretab/core/selectors.py

@@ -22,7 +22,7 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Literal
+from typing import Literal, cast
 
 import numpy as np
 from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
@@ -133,10 +133,11 @@ def select(
     def _ordered_candidates(self, x_valid: np.ndarray, y_valid: np.ndarray, task: Task) -> tuple[list[float], object]:
         """Fit a model and return candidate locations plus trimming context.
 
-        The candidates must be returned in the selector's preferred order (the
-        order :meth:`_enforce_spacing` should honour): location order for a
-        single tree, gain-descending order for a boosted ensemble. ``context`` is
-        an opaque object passed straight through to :meth:`_trim_over_max`.
+        The candidates must be returned in the selector's preferred order -- the
+        order :meth:`_enforce_spacing` honours, which keeps the earlier of two
+        candidates that are too close: impurity-decrease order for a single tree,
+        gain-descending order for a boosted ensemble. ``context`` is an opaque
+        object passed straight through to :meth:`_trim_over_max`.
         """
         raise NotImplementedError
 
@@ -148,9 +149,10 @@ def _trim_over_max(self, points: list[float], context: object, max_count: int) -
     def _enforce_spacing(self, split_points: list[float], x: np.ndarray) -> list[float]:
         """Drop locations closer than ``min_location_spacing`` of the range.
 
-        Compares each candidate against every already-kept location so the filter
-        is order-independent and honours both ascending (CART) and gain-descending
-        (LightGBM) ordering.
+        Compares each candidate against every already-kept location, so of two
+        candidates that are too close the one earlier in ``split_points`` -- the
+        more important one, given the importance-descending order both selectors
+        produce -- is kept.
         """
         if len(split_points) <= 1:
             return split_points
@@ -211,8 +213,9 @@ class CARTLocationSelector(BaseLocationSelector):
 
     A ``DecisionTreeRegressor`` or ``DecisionTreeClassifier`` is fitted to the
     feature against the target, and its split thresholds become the candidate
-    locations. Candidates are spaced out, and if there are too many they are
-    ranked by weighted impurity decrease so the most informative splits are kept.
+    locations. Candidates are ranked by weighted impurity decrease, so the most
+    informative splits are kept both when two candidates are too close to each
+    other and when there are too many.
 
     Parameters
     ----------
@@ -254,62 +257,50 @@ def _ordered_candidates(self, x_valid: np.ndarray, y_valid: np.ndarray, task: Ta
         )
 
         tree.fit(x_valid, y_valid)
-        split_points = self._extract_split_points(tree, x_valid)
-        return split_points, tree
+        importance = self._split_importance(tree, x_valid)
+        # Most informative split first, so that when two candidates are closer
+        # than ``min_location_spacing`` the spacing filter keeps the stronger one.
+        return self._rank(importance, importance), importance
 
     def _trim_over_max(self, points: list[float], context: object, max_count: int) -> list[float]:
-        return self._select_top_locations(points, context, max_count)
+        importance = cast(dict[float, float], context)
+        return sorted(self._rank(points, importance)[:max_count])
 
-    def _extract_split_points(self, tree, x: np.ndarray) -> list[float]:
-        """Collect in-range split thresholds from a fitted decision tree."""
-        tree_structure = tree.tree_
-        split_points = []
-
-        x_min, x_max = float(x.min()), float(x.max())
-
-        for node_id in range(tree_structure.node_count):
-            is_split = tree_structure.children_left[node_id] != tree_structure.children_right[node_id]
-            if not is_split:
-                continue
-            if tree_structure.feature[node_id] != 0:
-                continue
-            threshold = tree_structure.threshold[node_id]
-            if x_min < threshold < x_max:
-                split_points.append(threshold)
-
-        return sorted(set(split_points))
+    @staticmethod
+    def _rank(points, importance: dict[float, float]) -> list[float]:
+        """Order ``points`` by decreasing impurity decrease (ties by location)."""
+        return sorted(points, key=lambda point: (-importance[point], point))
 
-    def _select_top_locations(self, candidates: list[float], tree, max_count: int) -> list[float]:
-        """Keep the locations whose splits reduce impurity the most."""
-        if len(candidates) <= max_count:
-            return candidates
+    def _split_importance(self, tree, x: np.ndarray) -> dict[float, float]:
+        """Map each in-range split threshold to its weighted impurity decrease.
 
+        In a single-feature tree every threshold occurs at most once: after a split
+        at ``t`` no descendant holds samples on both sides of ``t``.
+        """
         tree_structure = tree.tree_
-        split_importance = {}
+        x_min, x_max = float(x.min()), float(x.max())
+        split_importance: dict[float, float] = {}
 
         for node_id in range(tree_structure.node_count):
-            is_split = tree_structure.children_left[node_id] != tree_structure.children_right[node_id]
-            if not is_split:
+            left_child = tree_structure.children_left[node_id]
+            right_child = tree_structure.children_right[node_id]
+            if left_child == right_child or tree_structure.feature[node_id] != 0:
                 continue
-
-            threshold = tree_structure.threshold[node_id]
-            if threshold not in candidates:
+            threshold = float(tree_structure.threshold[node_id])
+            if not x_min < threshold < x_max:
                 continue
 
             n_samples = tree_structure.n_node_samples[node_id]
             impurity = tree_structure.impurity[node_id]
-
-            left_child = tree_structure.children_left[node_id]
-            right_child = tree_structure.children_right[node_id]
             n_left = tree_structure.n_node_samples[left_child]
             n_right = tree_structure.n_node_samples[right_child]
             impurity_left = tree_structure.impurity[left_child]
             impurity_right = tree_structure.impurity[right_child]
+            split_importance[threshold] = float(
+                n_samples * impurity - (n_left * impurity_left + n_right * impurity_right)
+            )
 
-            split_importance[threshold] = n_samples * impurity - (n_left * impurity_left + n_right * impurity_right)
-
-        top = sorted(split_importance, key=lambda k: split_importance[k], reverse=True)[:max_count]
-        return sorted(top)
+        return split_importance
 
 
 class LightGBMLocationSelector(BaseLocationSelector):

tests/core/test_location_selectors.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 import numpy as np
 import pytest
 
@@ -180,6 +182,48 @@ def test_supplement_preserves_high_end_split_end_to_end():
     assert locations.max() > 5.0, f"expected a high-end location to survive supplementing, got {locations}"
 
 
+def _single_step(seed):
+    rng = np.random.default_rng(seed)
+    x = rng.uniform(0, 10, 1000)
+    step = rng.uniform(2, 8)
+    y = np.where(x > step, 1.0, 0.0) + 0.5 * rng.normal(size=1000)
+    return x.reshape(-1, 1), y
+
+
+@pytest.mark.parametrize("task", ["regression", None])
+def test_cart_spacing_keeps_the_dominant_split_over_a_weak_neighbour(task):
+    """Regression guard for issue #70: candidates were spaced in ascending location
+    order, so a weak split just below the root split evicted it."""
+    from sklearn.tree import DecisionTreeRegressor
+
+    X, y = _single_step(34)
+    root = DecisionTreeRegressor(max_depth=1).fit(X, y).tree_.threshold[0]
+    locations = CARTLocationSelector().select(X, y, task=task, min_count=6, max_count=6)
+    assert np.isclose(locations, root).any(), f"root split {root:.3f} missing from {locations}"
+
+
+@pytest.mark.parametrize("seed", range(10))
+def test_cart_keeps_the_root_split_of_a_single_step_target(seed):
+    from sklearn.tree import DecisionTreeRegressor
+
+    X, y = _single_step(seed)
+    root = DecisionTreeRegressor(max_depth=1).fit(X, y).tree_.threshold[0]
+    locations = CARTLocationSelector().select(X, y, task="regression", min_count=6, max_count=6)
+    assert np.isclose(locations, root).any()
+
+
+@pytest.mark.parametrize("task", ["regression", "classification"])
+def test_cart_candidates_are_ordered_by_impurity_decrease(data, task):
+    X, y = data
+    if task == "classification":
+        y = np.where(y > 0, "positive", "negative")
+    selector = CARTLocationSelector()
+    candidates, context = selector._ordered_candidates(X, y, task)
+    importance = cast("dict[float, float]", context)
+    gains = [importance[c] for c in candidates]
+    assert gains == sorted(gains, reverse=True)
+
+
 @pytest.mark.parametrize(
     "make_target",
     [