ml_meta_perf.selection

  1"""How long an equation should be: one rule for E1, E2 and E3-Valid, and a bound for E3-MAX.
  2
  3Every equation is searched once up to a horizon and cross-validated at every length, so the
  4choice of length is a reading of a complete curve. The curve read is the **floor**: the worst
  5R2 over in-sample (IS), leave-one-dataset-out (LODO), leave-one-model-out (LOMO) and doubly
  6held-out (DHO) validation, so a length is only as good as the protocol it does worst on.
  7
  8Adjacent lengths on that curve differ by 0.03 to 0.05 for reasons that belong to one fold, not
  9to the length, so any argmax over it -- or any threshold on its raw steps -- chooses noise.
 10`plateau_knee` therefore reads a smoothed curve and asks where the returns stop:
 11
 121. **smooth** the floor with a running median, so no single length can choose or block;
 132. **propose** lengths with multi-Kneedle (the `kneeliverse` library) on the smoothed curve's
 14   Pareto front -- the lengths where the rate of improvement bends -- plus the smoothed
 15   maximum itself, after which by definition nothing improves;
 163. **choose** the first proposed length after which the smoothed floor gains at most ``delta``
 17   over the next ``window`` lengths.
 18
 19A knee alone lands where the *rate* of improvement first bends, five to nine terms in, well
 20before the curve levels off; a threshold alone reads noise. Together they choose the start of
 21the first sustained plateau, which is the shortest equation that has stopped improving. When
 22every knee is still followed by real gains, that start is the smoothed maximum.
 23
 24E3-MAX is not put forward as an equation to read: it is the same configuration at the configured
 25capability arity, at the raw maximum of its floor (`floor_argmax`), bounding what the additive form can
 26reach.
 27
 28Study chapter: [3. Term generation and selection][study-chapter] -- the rationale, in prose.
 29
 30[study-chapter]: https://github.com/mariolpantunes/ml-meta-perf/blob/main/assets/docs/03-term-selection.md
 31"""
 32
 33from __future__ import annotations
 34
 35import kneeliverse.kneedle as kneedle
 36import kneeliverse.multi_knee as multi_knee
 37import numpy as np
 38import polars as pl
 39
 40#: Every protocol a length is judged on; `floor_curve` takes the minimum over those present.
 41JUDGED_PROTOCOLS = ("r2_in_sample", "r2_loo_dataset", "r2_loo_model", "r2_loo_cell")
 42
 43#: multi-Kneedle's recursion settings: stop splitting a segment once a straight line explains
 44#: it to this coefficient of determination ...
 45KNEE_FIT_THRESHOLD = 0.001
 46#: ... or once it has fewer points than this.
 47KNEE_MIN_POINTS = 3
 48
 49
 50def complexity(arity: int, n_terms: int) -> int:
 51    """Feature slots an equation spends: a term of arity ``a`` names ``a`` raw features.
 52
 53    **Not a count of fitted coefficients** -- those number ``n_terms + 1``. This charges for how
 54    much of the *grammar* an equation uses, which is what tells the same term count at arity 2
 55    apart from arity 3.
 56    """
 57    return arity * n_terms
 58
 59
 60def floor_curve(curve: pl.DataFrame) -> np.ndarray:
 61    """The **worst** of every protocol present, per length. This is what a length is judged on.
 62
 63    The minimum rather than a median, because the weakest protocol is almost always DHO -- the
 64    cell where neither the dataset nor the model has been seen, which is what the study's
 65    headline is about -- and a rule that reports the strictest protocol but selects on a looser
 66    summary selects on a different quantity from the one it publishes. The minimum is always
 67    some protocol's own number.
 68    """
 69    columns = [name for name in JUDGED_PROTOCOLS if name in curve.columns]
 70    if not columns:
 71        raise ValueError("curve carries none of the protocol columns")
 72    return np.column_stack([curve[name].to_numpy() for name in columns]).min(axis=1)
 73
 74
 75def smoothed(values: np.ndarray, width: int) -> np.ndarray:
 76    """A centred running median of ``width`` points, shrinking at the ends.
 77
 78    A median rather than a mean because the noise here is spiky: one held-out dataset outside
 79    the convex hull of the other nineteen craters a single length, and a mean would spread that
 80    crater onto its neighbours where a median ignores it. The same holds for a single lucky
 81    length, which is what a raw argmax picks.
 82    """
 83    if width < 1:
 84        raise ValueError("the smoothing width must be positive")
 85    half = width // 2
 86    return np.array([np.median(values[max(index - half, 0) : index + half + 1]) for index in range(len(values))])
 87
 88
 89def knee_lengths(curve: pl.DataFrame, smoothing: int) -> list[int]:
 90    """The lengths multi-Kneedle proposes on the smoothed floor's Pareto front, ascending.
 91
 92    The front keeps only lengths that improve on every shorter one, so a knee is always a
 93    length worth its terms; the knees are where the front's rate of improvement bends.
 94    """
 95    lengths = curve["n_terms"].to_numpy()
 96    floor = smoothed(floor_curve(curve), smoothing)
 97    front = [index for index in range(len(floor)) if index == 0 or floor[index] > floor[:index].max()]
 98    if len(front) < KNEE_MIN_POINTS + 1:
 99        return []
100    points = np.column_stack([lengths[front], floor[front]]).astype(np.float64)
101    found = multi_knee.multi_knee(kneedle.knee, points, t1=KNEE_FIT_THRESHOLD, t2=KNEE_MIN_POINTS)
102    return sorted({int(lengths[front[int(index)]]) for index in found})
103
104
105def plateau_start(curve: pl.DataFrame, *, delta: float, window: int, smoothing: int) -> tuple[int, str]:
106    """The equation length and how it was reached: ``"knee"`` or ``"maximum"``.
107
108    Candidates are the lengths `knee_lengths` proposes and the smoothed maximum. The chosen one
109    is the first, in increasing length, whose smoothed floor gains at most ``delta`` over the
110    next ``window`` lengths. The maximum always qualifies -- nothing after it gains -- so the
111    rule always returns; a maximum closer to the horizon than ``window`` is still returned,
112    and `experiment.selection_summary` flags a length that sits at the horizon.
113    """
114    if delta < 0.0 or window < 1:
115        raise ValueError("delta must be non-negative and the window positive")
116    lengths = curve["n_terms"].to_numpy()
117    floor = smoothed(floor_curve(curve), smoothing)
118    peak = int(lengths[int(np.argmax(floor))])
119    position = {int(length): index for index, length in enumerate(lengths)}
120    for candidate in sorted({*knee_lengths(curve, smoothing), peak}):
121        if candidate == peak:
122            return peak, "maximum"
123        index = position[candidate]
124        ahead = floor[index + 1 : index + 1 + window]
125        if len(ahead) == window and float(ahead.max() - floor[index]) <= delta:
126            return candidate, "knee"
127    return peak, "maximum"  # unreachable: the peak is always a candidate
128
129
130def plateau_knee(curve: pl.DataFrame, *, delta: float, window: int, smoothing: int) -> int:
131    """The equation length: the start of the first sustained plateau. **The length rule.**
132
133    See `plateau_start`, which also says whether the length is a knee or the smoothed maximum.
134    """
135    return plateau_start(curve, delta=delta, window=window, smoothing=smoothing)[0]
136
137
138def protocol_spread(curve: pl.DataFrame) -> np.ndarray:
139    """How far a length falls from its fit to its worst protocol: ``IS - floor``.
140
141    **Reported, not selected on.** An equation that fits well and transfers badly is a
142    different object from one that does both moderately, and the floor alone cannot tell them
143    apart. Folding the spread into the selection would need a weight between level and spread,
144    which is a free parameter; the two are reported side by side instead.
145    """
146    if "r2_in_sample" not in curve.columns:
147        raise ValueError("a spread needs the IS column to measure the drop from")
148    return curve["r2_in_sample"].to_numpy() - floor_curve(curve)
149
150
151def floor_argmax_index(curve: pl.DataFrame) -> int:
152    """The *row position* of `floor_argmax`'s length, for callers indexing the curve."""
153    return int(np.argmax(floor_curve(curve)))
154
155
156def floor_argmax(curve: pl.DataFrame) -> int:
157    """The length maximising the raw floor: **E3-MAX's rule**, a capability bound.
158
159    No smoothing and no complexity penalty, because this equation is not put forward to be
160    read -- it bounds what the additive form can reach, and a bound should not be discounted
161    for being long. Ties go to the shorter equation. A result at the horizon means the bound
162    is the horizon's, not the data's; `experiment` reports that.
163    """
164    return int(curve["n_terms"].to_numpy()[floor_argmax_index(curve)])