ml_meta_perf.search

  1"""Choosing which terms enter the equation.
  2
  3This is the search. `ml_meta_perf.fit` is the arithmetic it scores candidates with, and the
  4split is the study's own: a form is chosen once, and weights are refitted per fold against a
  5form held still -- which is what `validate.cross_validate_fixed_form` reports. Until
  62026-09-09 both lived in a module called `fit`, whose entry point was a function called `fit`
  7that ran a beam search, and the one distinction the results rest on was invisible in the names.
  8
  9Four ideas do the work here.
 10
 11*Correlation-guided term design.* Pearson and Spearman disagree in a informative way.
 12A feature whose Pearson and Spearman correlations against MCC are similar is related to
 13the target linearly, and ``f`` on its own is the right term -- adding ``log(f)`` and
 14``1/f`` for it only inflates the search space. A feature whose Spearman correlation is
 15much the larger is monotonically but *non*-linearly related, and that is precisely the
 16case where a compressing or inverting transform earns its place. Screening on that gap
 17shrinks the library before search begins, which is what keeps the selected equations
 18short without a parsimony penalty bolted on afterwards.
 19
 20*Beam search over subsets, not greedy descent.* Greedy forward selection commits to its
 21first term forever, and on collinear libraries that first commitment is often wrong.
 22A beam keeps several partial equations alive, and because the beam records its best
 23subset at every size, one run yields the whole accuracy-versus-number-of-terms curve.
 24
 25*One term per feature combination.* An equation that describes ``a`` against ``b`` twice --
 26once as ``a/b`` and once as ``b/a`` -- cannot be read a term at a time, because neither
 27term means on its own what the table says it means. The selector therefore refuses a
 28candidate whose feature combination is already represented, which is a constraint on
 29*form* rather than on fit and is not negotiable against R2: an equation nobody can reason
 30about has failed at the only thing this study asks of it. It is measured, not assumed, to
 31cost nothing -- see `terms.Library.feature_groups`. Note that this is a different test
 32from ``COLLINEARITY_LIMIT``, which asks a numerical question; the pairs this removes were
 33well inside that limit.
 34
 35*Pruning is a decision about form, so it is here.* `prune` drops terms that do no work and
 36simplifies what remains -- both changes to which terms the equation contains -- and then asks
 37`fit` for the weights of what is left. An equation shortened without that refit would be the
 38best fit of terms it no longer has.
 39
 40**`Selector` holds a ridge fit and is still in this module.** Its `_evaluate` is arithmetic
 41that `fit` would recognise, but it exists to be called several hundred thousand times by the
 42beam, off a Gram matrix and a cache the beam's constructor builds. Splitting it would put the
 43two halves of one constructor in two files to make a naming point, which is worse than the
 44naming point is worth.
 45
 46Study chapter: [3. Term generation and selection][study-chapter] -- the rationale, in
 47prose, with the figures.
 48
 49[study-chapter]: https://github.com/mariolpantunes/ml-meta-perf/blob/main/assets/docs/03-term-selection.md
 50"""
 51
 52from __future__ import annotations
 53
 54from dataclasses import dataclass
 55
 56import numpy as np
 57
 58from ml_meta_perf.fit import RIDGE_DEFAULT, Standardizer, Subset, to_equation
 59from ml_meta_perf.model import Equation
 60from ml_meta_perf.stats import pearson, pearson_columns, rank_columns, rankdata, spearman
 61from ml_meta_perf.terms import Library, Term, is_trivial, simplify
 62
 63BEAM_WIDTH_DEFAULT = 6
 64CANDIDATE_POOL_DEFAULT = 24
 65COLLINEARITY_LIMIT = 0.95
 66REFINE_ROUNDS_DEFAULT = 2
 67
 68
 69def transform_gap(values: np.ndarray, target: np.ndarray) -> float:
 70    """How much more monotone than linear a feature's relation to the target is.
 71
 72    Near zero means a plain linear term is adequate. Clearly positive means a transform
 73    (log, inverse, power) or a ratio is likely to pay, because the association is there
 74    but a straight line cannot capture it.
 75    """
 76    return abs(spearman(values, target)) - abs(pearson(values, target))
 77
 78
 79def guided_screen(
 80    library: Library,
 81    target: np.ndarray,
 82    *,
 83    keep: int,
 84    linear_gap: float = 0.05,
 85) -> list[int]:
 86    """Reduce the library to the terms worth searching over.
 87
 88    Terms are ranked by the stronger of their two correlations, so a term that is
 89    strongly monotone but weakly linear survives -- it is exactly the term a linear
 90    equation needs a transform for. Terms that are near-duplicates of a stronger term
 91    already kept are dropped, since the beam would otherwise spend its width on
 92    variations of one idea.
 93    """
 94    matrix = library.matrix
 95    # Ranking is the expensive half of a Spearman correlation -- it sorts, then averages
 96    # ties -- and screening a 600-term pool inside 20 folds asks for it tens of thousands
 97    # of times. Nothing here is per-candidate: the target's ranks do not change from one
 98    # candidate to the next, the candidates' own ranks are one vectorised pass over the
 99    # matrix (`stats.rank_columns`), and both correlations are one matrix-vector product
100    # over the whole pool (`stats.pearson_columns`) rather than a `pearson` call per column.
101    ranked_target = rankdata(target)
102    ranked = rank_columns(matrix) if matrix.shape[0] > 1 else matrix
103    linear = np.abs(pearson_columns(matrix, target))
104    monotone = np.abs(pearson_columns(ranked, ranked_target)) if matrix.shape[0] > 1 else np.zeros(matrix.shape[1])
105    # A term that is only monotone is still useful, but it is preferred a little
106    # less than an equally strong linear one: linear terms read more simply.
107    strength = np.maximum(linear, monotone) - np.where(monotone - linear > linear_gap, 0.02, 0.0)
108    # Descending by strength, ties broken by ascending library index, for the reason
109    # `_descending` sets out: the survivors of the duplicate pass below must be a property
110    # of the library rather than of the order a caller assembled it in.
111    scored = _descending(strength)
112
113    standardized = Standardizer.fit(matrix).apply(matrix)
114    rows = float(standardized.shape[0])
115
116    # The duplicate check is one matrix-vector product against the columns kept so far,
117    # not a Python loop of dot products over them. At 1599 candidates keeping 1200 the
118    # loop form ran 17 million times and accounted for over half the total runtime of a
119    # cross-validated sweep; this is the same arithmetic in one BLAS call per candidate.
120    accepted = np.empty((standardized.shape[0], min(keep, standardized.shape[1])))
121    kept: list[int] = []
122    for position in scored:
123        if len(kept) >= keep:
124            break
125        index = int(position)
126        column = standardized[:, index]
127        if kept:
128            correlations = np.abs(column @ accepted[:, : len(kept)]) / rows
129            if float(correlations.max()) > 0.995:
130                continue
131        accepted[:, len(kept)] = column
132        kept.append(index)
133    return kept
134
135
136def _descending(scores: np.ndarray) -> np.ndarray:
137    """Indices of ``scores`` from largest to smallest, ties broken by ascending index.
138
139    ``np.argsort(scores)[::-1]`` looks equivalent and is not, twice over: the default sort is
140    quicksort, which is unstable, and reversing even a stable ascending sort *inverts* the tie
141    order rather than preserving it. Either way ties were resolved by where a term happened to
142    land in the library, and `terms.build_library` emits terms in feature declaration order --
143    so handing `run_equation` the same columns in a different order produced a different
144    equation, and leave-one-dataset-out (LODO) R² moved by as much as 0.286 across seven orderings of a
145    single ten-feature set.
146
147    Sorting the negated scores with an explicitly stable kind breaks ties by ascending library
148    index instead, which is a property of the library rather than of the caller's argument
149    order.
150    """
151    return np.argsort(-scores, kind="stable")
152
153
154class Selector:
155    """Beam search over term subsets against a fixed, standardised design."""
156
157    def __init__(
158        self,
159        design: np.ndarray,
160        target: np.ndarray,
161        penalty: float,
162        groups: np.ndarray | None = None,
163    ) -> None:
164        self.design = design
165        self.centered = target - target.mean()
166        self.offset = float(target.mean())
167        self.penalty = penalty
168        self.gram = design.T @ design
169        # The ridge penalty added to the *full* Gram diagonal, once. A submatrix's
170        # diagonal is drawn from the full diagonal, so ``penalized[ix, ix]`` already
171        # equals ``gram[ix, ix] + penalty*I`` -- which removes a copy, an arange and a
172        # fancy-index assignment from every one of the million solves a study runs.
173        # ``gram`` itself stays unpenalised: the residual scoring and the collinearity
174        # guard both need it that way.
175        self.penalized = self.gram.copy()
176        if penalty:
177            self.penalized.flat[:: self.penalized.shape[0] + 1] += penalty
178        self.projection = design.T @ self.centered
179        self.total = float(self.centered @ self.centered)
180        self.normalizer = float(design.shape[0])
181        # The collinearity test as a boolean matrix, taken once. ``_blocked`` runs it for
182        # every parent of every beam step -- tens of thousands of times per fit -- and the
183        # arithmetic does not depend on the parent, only the columns selected out of it.
184        # Deciding it here turns each call into a bool gather over an eighth of the bytes.
185        self._collinear = np.abs(self.gram) / self.normalizer > COLLINEARITY_LIMIT
186        self._cache: dict[tuple[int, ...], Subset] = {}
187        # `terms.Library.feature_groups`, expanded once into a membership matrix so that
188        # `_blocked` is a row gather and an `any`, not an `np.isin` over the pool on each
189        # of several hundred thousand calls. Row ``g`` marks every term in group ``g``.
190        self.groups = groups
191        self._members: np.ndarray | None = None
192        if groups is not None and groups.size:
193            count = int(groups.max()) + 1
194            if count > 0:
195                members = np.zeros((count, groups.shape[0]), dtype=bool)
196                present = groups >= 0
197                members[groups[present], np.flatnonzero(present)] = True
198                self._members = members
199
200    def _evaluate(self, indices: tuple[int, ...]) -> Subset:
201        """Fit one subset. Memoised, because the refinement pass revisits subsets.
202
203        The residual sum of squares uses an identity rather than a second quadratic form.
204        Since ``(G + lambda I) w = b``, we have ``w'Gw = w'b - lambda w'w``, so
205
206            rss = total - 2 w'b + w'Gw = total - w'b - lambda w'w
207
208        which removes a matrix-vector product from the inner loop. The two are equal to
209        floating-point tolerance whenever the solve succeeded, which ``test_fit`` checks
210        directly against the literal definition.
211        """
212        cached = self._cache.get(indices)
213        if cached is not None:
214            return cached
215
216        order = np.array(indices, dtype=np.intp)
217        # ``penalized[order[:, None], order]`` rather than ``np.ix_(order, order)``.
218        # Identical arithmetic: ix_ exists to build exactly this pair of broadcast index
219        # arrays, and it spends most of its time on dtype checks already known to hold
220        # here. At a million calls that check cost 24 of the study's 137 seconds, and
221        # doing the reshape directly is 1.8x faster.
222        block = self.penalized[order[:, None], order]
223        rhs = self.projection[order]
224        try:
225            weights = np.linalg.solve(block, rhs)
226        except np.linalg.LinAlgError:
227            weights = np.linalg.lstsq(block, rhs, rcond=None)[0]
228        rss = self.total - float(weights @ rhs) - self.penalty * float(weights @ weights)
229        subset = Subset(indices, weights, rss)
230        self._cache[indices] = subset
231        return subset
232
233    def _evaluate_many(self, batch: list[tuple[int, ...]]) -> list[Subset]:
234        """Fit several subsets **of the same size** in one stacked solve.
235
236        A beam step and a refinement position both produce a few dozen candidate subsets
237        that differ by one index, and every one of them is a ``k``-by-``k`` solve. At
238        that size numpy's per-call wrapper -- dtype promotion, array coercion, the
239        errstate context manager -- costs several times the LAPACK call it guards, and
240        the study makes over a million of them. Stacking the blocks into one ``(n, k, k)``
241        array moves that loop into C: same routine per slice, bitwise identical weights,
242        two to ten times faster depending on ``k``.
243
244        Order is preserved and the cache is shared with `_evaluate`, so callers see
245        exactly what the one-at-a-time path gave them.
246        """
247        pending: list[tuple[int, ...]] = []
248        queued: set[tuple[int, ...]] = set()
249        for indices in batch:
250            if indices not in self._cache and indices not in queued:
251                queued.add(indices)
252                pending.append(indices)
253
254        if pending:
255            order = np.array(pending, dtype=np.intp)
256            blocks = self.penalized[order[:, :, None], order[:, None, :]]
257            rhs = self.projection[order]
258            try:
259                weights = np.linalg.solve(blocks, rhs[..., None])[..., 0]
260            except np.linalg.LinAlgError:
261                # One singular block fails the whole stack, so fall back to the
262                # per-subset path, which has its own least-squares rescue.
263                for indices in pending:
264                    self._evaluate(indices)
265            else:
266                scores = self.total - np.einsum("ij,ij->i", weights, rhs)
267                if self.penalty:
268                    scores -= self.penalty * np.einsum("ij,ij->i", weights, weights)
269                for indices, row, rss in zip(pending, weights, scores, strict=True):
270                    self._cache[indices] = Subset(indices, row, float(rss))
271
272        return [self._cache[indices] for indices in batch]
273
274    def _residual_scores(self, subset: Subset) -> np.ndarray:
275        order = np.array(subset.indices, dtype=np.intp)
276        return np.abs(self.projection - self.gram[:, order] @ subset.weights)
277
278    def _blocked(self, indices: tuple[int, ...]) -> np.ndarray | None:
279        """Which candidates this parent may not take, as one boolean mask.
280
281        Computed per parent rather than per candidate. The per-candidate form rebuilt an
282        index array and took a max on every one of several hundred thousand calls; this
283        does the same work as a single vectorised reduction over the Gram matrix.
284
285        Two independent refusals, and they answer different questions. The correlation
286        test asks whether a candidate is numerically redundant against what is already
287        chosen. The group test asks whether it would state an *already-stated relationship
288        a second time* -- see `terms.Library.feature_groups`. Neither implies the other:
289        the mirrored pairs the group test removes sat at 0.891 and 0.786, comfortably
290        inside ``COLLINEARITY_LIMIT``, and two terms may correlate above the limit while
291        describing different feature pairs entirely.
292        """
293        if not indices:
294            return None
295        order = np.array(indices, dtype=np.intp)
296        blocked = self._collinear[:, order].any(axis=1)
297        if self._members is None or self.groups is None:
298            return blocked
299        chosen = self.groups[order]
300        chosen = chosen[chosen >= 0]
301        if chosen.size == 0:
302            return blocked
303        # No `np.unique` on the way in: a repeated group id gathers the same membership row
304        # twice and an `any` over it reduces to the same mask, so the sort it costs bought
305        # nothing across several hundred thousand calls.
306        return blocked | self._members[chosen].any(axis=0)
307
308    def search(
309        self,
310        pool: list[int],
311        max_terms: int,
312        *,
313        beam_width: int = BEAM_WIDTH_DEFAULT,
314        candidates: int = CANDIDATE_POOL_DEFAULT,
315        refine_rounds: int = REFINE_ROUNDS_DEFAULT,
316    ) -> dict[int, Subset]:
317        """Return the best subset found at every size from 1 to ``max_terms``."""
318        available = np.array(pool)
319        beam: list[Subset] = [Subset((), np.zeros(0), self.total)]
320        best: dict[int, Subset] = {}
321
322        for size in range(1, max_terms + 1):
323            seen: set[tuple[int, ...]] = set()
324            children: list[tuple[int, ...]] = []
325            for parent in beam:
326                scores = self._residual_scores(parent)
327                ranked = available[_descending(scores[available])]
328                blocked = self._blocked(parent.indices)
329                taken = 0
330                for candidate in ranked:
331                    if taken >= candidates:
332                        break
333                    index = int(candidate)
334                    if index in parent.indices or (blocked is not None and blocked[index]):
335                        continue
336                    child = tuple(sorted((*parent.indices, index)))
337                    if child in seen:
338                        continue
339                    seen.add(child)
340                    children.append(child)
341                    taken += 1
342            # Every child of this step has the same size, which is what lets them go
343            # through the solver as one stack rather than one at a time.
344            generated = self._evaluate_many(children)
345            if not generated:
346                break
347            generated.sort(key=lambda item: item.rss)
348            beam = generated[:beam_width]
349            best[size] = self._refine(beam[0], available, refine_rounds, candidates)
350            if best[size].rss < beam[0].rss:
351                beam[0] = best[size]
352        return best
353
354    def _refine(
355        self,
356        subset: Subset,
357        available: np.ndarray,
358        rounds: int = REFINE_ROUNDS_DEFAULT,
359        candidates: int = CANDIDATE_POOL_DEFAULT,
360    ) -> Subset:
361        """Local search: try replacing each chosen term with a better one.
362
363        Beam search still fixes early choices under a narrow width. Swapping one term at
364        a time recovers the cases where a term that looked good at step two is dominated
365        once the rest of the equation is in place.
366        """
367        current = subset
368        for _ in range(rounds):
369            improved = False
370            warmed: tuple[int, ...] | None = None
371            for position in range(len(current.indices)):
372                # The probe for every position of one subset is a leave-one-out of it, so all
373                # of them are the same size and price in a single stacked solve. Priced up
374                # front and read out of the cache position by position; `_evaluate_many` shares
375                # the cache with `_evaluate`, so this only changes when the solves happen.
376                if warmed != current.indices:
377                    warmed = current.indices
378                    if len(warmed) > 1:
379                        self._evaluate_many(
380                            [tuple(x for i, x in enumerate(warmed) if i != p) for p in range(len(warmed))]
381                        )
382                remaining = tuple(x for i, x in enumerate(current.indices) if i != position)
383                probe = self._evaluate(remaining) if remaining else Subset((), np.zeros(0), self.total)
384                scores = self._residual_scores(probe) if remaining else np.abs(self.projection)
385                ranked = available[_descending(scores[available])][:candidates]
386                blocked = self._blocked(remaining)
387                trials = [
388                    tuple(sorted((*remaining, int(candidate))))
389                    for candidate in ranked
390                    if int(candidate) not in remaining and (blocked is None or not blocked[int(candidate)])
391                ]
392                # Scanned in ranked order for the *first* improvement, exactly as the
393                # one-at-a-time loop did. Evaluating the rest of the batch is not waste:
394                # a refinement round that improves nothing has to price every candidate
395                # anyway, and anything priced here is cached for the rounds after it.
396                for trial in self._evaluate_many(trials):
397                    if trial.rss < current.rss - 1e-12:
398                        current = trial
399                        improved = True
400                        break
401            if not improved:
402                break
403        return current
404
405
406@dataclass(frozen=True)
407class SearchResult:
408    """Everything one search produces: the equation per size, and the pool it drew from.
409
410    One search yields the whole curve rather than one equation, because the beam records its
411    best subset at every size as it grows. The selection rules read that complete curve, which
412    is why the selected length is not a parameter here.
413    """
414
415    equations: dict[int, Equation]
416    pool_size: int
417
418    def best(self, max_terms: int | None = None) -> Equation:
419        sizes = [size for size in self.equations if max_terms is None or size <= max_terms]
420        return self.equations[max(sizes)]
421
422
423def search(
424    library: Library,
425    target: np.ndarray,
426    *,
427    max_terms: int = 12,
428    penalty: float = RIDGE_DEFAULT,
429    pool_size: int = 250,
430    beam_width: int = BEAM_WIDTH_DEFAULT,
431    candidates: int = CANDIDATE_POOL_DEFAULT,
432    refine_rounds: int = REFINE_ROUNDS_DEFAULT,
433    name: str = "equation",
434) -> SearchResult:
435    """Search for the best subset of every size up to ``max_terms`` over the given library.
436
437    Screen, then beam, then fold the standardisation back out. The weights each returned
438    equation carries come from `fit`; what this function decides is which terms carry them.
439    """
440    standardizer = Standardizer.fit(library.matrix)
441    design = standardizer.apply(library.matrix)
442    pool = guided_screen(library, target, keep=pool_size)
443    selector = Selector(design, target, penalty, library.feature_groups)
444    subsets = selector.search(
445        pool,
446        max_terms,
447        beam_width=beam_width,
448        candidates=candidates,
449        refine_rounds=refine_rounds,
450    )
451    equations = {
452        size: to_equation(library, subset, standardizer, selector.offset, f"{name}_k{size}")
453        for size, subset in subsets.items()
454    }
455    return SearchResult(equations=equations, pool_size=len(pool))
456
457
458#: A term whose contribution never moves predicted MCC by this much across the data is
459#: not doing work worth printing. The default is deliberately well below the resolution
460#: anyone reads MCC at.
461MIN_CONTRIBUTION = 0.002
462
463
464def prune(
465    equation: Equation,
466    columns: dict[str, np.ndarray],
467    target: np.ndarray,
468    *,
469    penalty: float = RIDGE_DEFAULT,
470    min_contribution: float = MIN_CONTRIBUTION,
471) -> Equation:
472    """Drop terms that do no work, simplify what remains, and refit the weights.
473
474    Two things make a published equation longer than it needs to be.
475
476    A term can survive selection and then contribute nothing: beam search adds terms while
477    residual sum of squares falls, and the last few can fall by amounts invisible in the
478    output. One dendrogram-cut equation carried a term weighted -3.7e-15 -- a full
479    expression over six features, contributing zero.
480
481    A term can also be algebraically redundant, computing ``a`` while printing
482    ``([a] / [b]) * [b]``. ``terms.simplify`` handles that, and terms that reduce to a
483    constant are removed outright since they duplicate the intercept.
484
485    Weights are refitted after pruning rather than carried over. Dropping a term changes
486    what the others should be, and keeping stale weights would leave an equation that is
487    shorter but no longer the best fit of its own remaining terms.
488    """
489    if not equation.terms:
490        return equation
491
492    contributions = np.column_stack(
493        [weight * term.evaluate(columns) for term, weight in zip(equation.terms, equation.weights, strict=True)]
494    )
495    keep: list[Term] = []
496    seen: set[str] = set()
497    for index, term in enumerate(equation.terms):
498        column = contributions[:, index]
499        low, high = np.percentile(column, [1.0, 99.0])
500        if float(high - low) < min_contribution:
501            continue
502        reduced = simplify(term)
503        if is_trivial(reduced) or reduced.name in seen:
504            continue
505        seen.add(reduced.name)
506        keep.append(reduced)
507
508    if not keep:
509        return Equation(
510            intercept=float(target.mean()), terms=(), weights=(), standardized_weights=(), name=equation.name
511        )
512
513    library = Library(keep, columns)
514    standardizer = Standardizer.fit(library.matrix)
515    selector = Selector(standardizer.apply(library.matrix), target, penalty, library.feature_groups)
516    subset = selector._evaluate(tuple(range(len(library))))
517    return to_equation(library, subset, standardizer, selector.offset, equation.name)
518
519
520def selected_terms(equation: Equation) -> list[Term]:
521    """The terms of an equation, strongest standardised contribution first."""
522    return [term for term, _, _ in equation.ranked_terms()]