diff --git a/pretab/core/selectors.py b/pretab/core/selectors.py index bf172e4..a4331d6 100644 --- a/pretab/core/selectors.py +++ b/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): diff --git a/tests/test_location_selectors.py b/tests/test_location_selectors.py index 15fe2e3..928f170 100644 --- a/tests/test_location_selectors.py +++ b/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)