diff --git a/pretab/core/selectors.py b/pretab/core/selectors.py index bf172e4..d828a61 100644 --- a/pretab/core/selectors.py +++ b/pretab/core/selectors.py @@ -133,16 +133,29 @@ def _trim_over_max(self, points: list[float], context: object, max_count: int) - raise NotImplementedError def _enforce_spacing(self, split_points: list[float], x: np.ndarray) -> list[float]: - """Drop locations closer than ``min_location_spacing`` of the range.""" + """Drop locations closer than ``min_location_spacing`` of the range. + + The distance test runs against *every* location kept so far rather than + only the most recent one, which makes the filter independent of the + order in which candidates arrive. That matters because subclasses + deliberately order their candidates differently -- location order for a + single tree, gain-descending order for a boosted ensemble -- and an + order-sensitive test silently dropped every candidate positioned below + the previously kept one, collapsing a gain-ranked set into a small + clustered subsequence. + + For already-ascending input (the single-tree path) this is equivalent to + comparing against the last kept location, so that behaviour is unchanged. + """ if len(split_points) <= 1: return split_points x_range = float(x.max() - x.min()) min_distance = self.min_location_spacing * x_range - spaced = [split_points[0]] + spaced: list[float] = [split_points[0]] for point in split_points[1:]: - if point - spaced[-1] >= min_distance: + if all(abs(point - kept) >= min_distance for kept in spaced): spaced.append(point) return spaced diff --git a/tests/test_location_selectors.py b/tests/test_location_selectors.py index 15fe2e3..1a3e96e 100644 --- a/tests/test_location_selectors.py +++ b/tests/test_location_selectors.py @@ -119,3 +119,76 @@ def test_lightgbm_matches_knot_adapter(data): X, y, task="regression", min_count=adapter.min_knots, max_count=adapter.max_knots ) np.testing.assert_array_equal(from_adapter, from_selector) + + +# --------------------------------------------------------------------------- # +# Spacing must not depend on the order candidates arrive in. +# +# ``_ordered_candidates`` returns location order for a single tree but +# gain-descending order for a boosted ensemble, so an order-sensitive spacing +# filter discarded every candidate sitting below the previously kept one. +# --------------------------------------------------------------------------- # +def test_enforce_spacing_is_order_independent(): + x = np.linspace(0, 10, 500).reshape(-1, 1) + selector = CARTLocationSelector(min_location_spacing=0.01) + ascending = [1.0, 2.0, 3.0, 8.0, 9.0] + + from_ascending = selector._enforce_spacing(list(ascending), x) + from_shuffled = selector._enforce_spacing([9.0, 1.0, 8.0, 3.0, 2.0], x) + + assert sorted(from_shuffled) == sorted(from_ascending) == ascending + + +def test_enforce_spacing_still_drops_close_neighbours(): + x = np.linspace(0, 10, 500).reshape(-1, 1) # min_distance = 0.1 + selector = CARTLocationSelector(min_location_spacing=0.01) + + assert selector._enforce_spacing([1.0, 1.01, 5.0], x) == [1.0, 5.0] + + +def test_enforce_spacing_matches_legacy_result_on_ascending_input(): + # The single-tree path always supplies ascending candidates; comparing + # against every kept location must be equivalent there. + rng = np.random.RandomState(0) + x = rng.uniform(0, 10, size=(500, 1)) + points = sorted(rng.uniform(0, 10, size=40).tolist()) + selector = CARTLocationSelector(min_location_spacing=0.05) + + min_distance = 0.05 * float(x.max() - x.min()) + legacy = [points[0]] + for point in points[1:]: + if point - legacy[-1] >= min_distance: + legacy.append(point) + + assert selector._enforce_spacing(points, x) == legacy + + +def test_lightgbm_locations_cover_the_feature_range(data): + pytest.importorskip("lightgbm") + X, y = data + x_min, x_max = float(X.min()), float(X.max()) + + lgbm = LightGBMLocationSelector(n_estimators=30).select( + X, y, task="regression", min_count=3, max_count=8 + ) + cart = CARTLocationSelector().select(X, y, task="regression", min_count=3, max_count=8) + + def span(locations): + return (locations.max() - locations.min()) / (x_max - x_min) + + # Previously the gain-ordered candidates collapsed into a narrow cluster; + # coverage should now be comparable to the single-tree selector. + assert span(lgbm) > 0.5 * span(cart) + + +def test_lightgbm_keeps_the_highest_gain_locations(data): + pytest.importorskip("lightgbm") + X, y = data + selector = LightGBMLocationSelector(n_estimators=30) + + candidates, _ = selector._ordered_candidates(X, y, "regression") + spaced = selector._enforce_spacing(candidates, X) + # Gain-descending order survives the spacing filter, so the trim keeps the + # top-ranked entries rather than whichever happen to be numerically lowest. + assert spaced == [c for c in candidates if c in set(spaced)] + assert selector._trim_over_max(spaced, None, 5) == spaced[:5]