fix: keep selected locations when supplementing

pretab/core/selectors.py

@@ -26,7 +26,7 @@
 from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 
 from .exceptions import IncompatibleParamsError, OptionalDependencyError
-from .knots import quantile_knots
+from .knots import quantile_knots, select_knots
 
 Task = Literal["regression", "classification"]
 
@@ -148,13 +148,30 @@ def _enforce_spacing(self, split_points: list[float], x: np.ndarray) -> list[flo
         return spaced
 
     def _supplement(self, existing: list[float], x: np.ndarray, target_count: int) -> list[float]:
-        """Top up an under-filled location set with quantile locations."""
-        if target_count - len(existing) <= 0:
-            return existing
+        """Top up an under-filled location set with quantile locations.
+
+        The locations already selected are always kept; only the shortfall is
+        filled, and the quantile candidates used to fill it are down-sampled by
+        even spacing so they stay spread across the range.
+
+        Merging both sets and truncating with ``sorted(...)[:target_count]``
+        instead always discarded the *largest* values. Since
+        :func:`~pretab.core.knots.quantile_knots` already returns
+        ``target_count`` candidates, the merged set always overflowed by exactly
+        ``len(existing)``, so the truncation reliably threw away the data-driven
+        locations the selector had just found whenever they sat above the median.
+        """
+        missing = target_count - len(existing)
+        if missing <= 0:
+            return sorted(existing)
 
-        quantile_candidates = quantile_knots(x, target_count)
-        all_locations = set(existing) | set(quantile_candidates.tolist())
-        return sorted(all_locations)[:target_count]
+        current = set(existing)
+        candidates = np.asarray(
+            [c for c in quantile_knots(x, target_count).tolist() if c not in current]
+        )
+        if len(candidates) > missing:
+            candidates = select_knots(candidates, missing)
+        return sorted(current | set(candidates.tolist()))
 
 
 class CARTLocationSelector(BaseLocationSelector):

tests/test_location_selectors.py

@@ -79,6 +79,57 @@ def test_cart_handles_nan_rows(data):
     assert np.isfinite(locations).all()
 
 
+# --------------------------------------------------------------------------- #
+# Supplementing an under-filled set must not discard the selected locations.
+#
+# ``quantile_knots`` already returns ``target_count`` candidates, so merging and
+# truncating with ``sorted(...)[:target_count]`` always dropped the largest
+# entries -- i.e. the tree-selected locations whenever they sat above the median.
+# --------------------------------------------------------------------------- #
+def test_supplement_keeps_every_existing_location():
+    x = np.linspace(0, 10, 500).reshape(-1, 1)
+    selector = CARTLocationSelector()
+
+    supplemented = selector._supplement([9.5, 9.7], x, 5)
+
+    assert 9.5 in supplemented
+    assert 9.7 in supplemented
+    assert len(supplemented) == 5
+    assert supplemented == sorted(supplemented)
+
+
+def test_supplement_fills_only_the_shortfall():
+    x = np.linspace(0, 10, 500).reshape(-1, 1)
+    selector = CARTLocationSelector()
+    existing = [1.0, 2.0, 3.0]
+
+    supplemented = selector._supplement(list(existing), x, 6)
+
+    assert set(existing).issubset(supplemented)
+    assert len(supplemented) == 6
+
+
+def test_supplement_is_a_noop_when_already_full():
+    x = np.linspace(0, 10, 500).reshape(-1, 1)
+    selector = CARTLocationSelector()
+
+    assert selector._supplement([2.0, 4.0, 6.0], x, 3) == [2.0, 4.0, 6.0]
+
+
+def test_supplement_spreads_the_added_candidates():
+    # A single existing location must not cause the fill-ins to bunch at one end.
+    x = np.linspace(0, 10, 500).reshape(-1, 1)
+    selector = CARTLocationSelector()
+
+    supplemented = selector._supplement([5.0], x, 4)
+
+    assert 5.0 in supplemented
+    assert min(supplemented) < 4.0
+    assert max(supplemented) > 6.0
+
+
+
+
 def test_cart_matches_knot_adapter(data):
     X, y = data
     adapter = CARTKnotSelector(max_basis_functions=12, degree=3)