ml_meta_perf.fit

  1"""Fitting weights to a form that has already been chosen.
  2
  3The study's central distinction, in a module boundary. An equation is two things -- which
  4terms it contains and what they are multiplied by -- and they are decided by different means
  5and validated by different protocols. `ml_meta_perf.search` chooses the terms;
  6this module does the arithmetic that turns a chosen set of them into an equation, and nothing
  7here has an opinion about which set that should be.
  8
  9That line is not a tidiness preference. `validate.cross_validate_fixed_form` is the study's
 10reported protocol precisely because it holds the form still and refits only the weights per
 11fold, and until 2026-09-09 the function that performed a search was called `fit`, which
 12blurred the one distinction the results rest on. Chapter 3 calls the other half "term
 13selection" and always has.
 14
 15*Gram-matrix arithmetic.* Every refit is a k-by-k solve against a precomputed Gram matrix
 16rather than a least-squares call against the full design, which is what makes an exhaustive
 17sweep inside 20 cross-validation folds tractable. `ridge_solve` states that arithmetic in one
 18place; `search.Selector` indexes submatrices out of a Gram it holds, for speed, and
 19`tests/test_search.py` checks the two agree.
 20
 21Study chapter: [3. Term generation and selection][study-chapter] -- the rationale, in
 22prose, with the figures.
 23
 24[study-chapter]: https://github.com/mariolpantunes/ml-meta-perf/blob/main/assets/docs/03-term-selection.md
 25"""
 26
 27from __future__ import annotations
 28
 29from dataclasses import dataclass
 30
 31import numpy as np
 32
 33from ml_meta_perf.model import Equation
 34from ml_meta_perf.terms import Library
 35
 36#: The ridge penalty a caller gets if it does not choose one.
 37RIDGE_DEFAULT = 10.0
 38
 39
 40@dataclass(frozen=True)
 41class Standardizer:
 42    """Column means and scales, learned on training rows only."""
 43
 44    mean: np.ndarray
 45    scale: np.ndarray
 46
 47    @classmethod
 48    def fit(cls, matrix: np.ndarray) -> Standardizer:
 49        scale = matrix.std(axis=0)
 50        return cls(matrix.mean(axis=0), np.where(scale < 1e-12, 1.0, scale))
 51
 52    def apply(self, matrix: np.ndarray) -> np.ndarray:
 53        return (matrix - self.mean) / self.scale
 54
 55
 56@dataclass(frozen=True)
 57class Subset:
 58    """One candidate equation: which terms, their standardised weights, and its fit."""
 59
 60    indices: tuple[int, ...]
 61    weights: np.ndarray
 62    #: The **penalised** objective this subset reached, ``||y - Xw||^2 + lambda*||w||^2``,
 63    #: not a residual sum of squares. It is only ever compared -- the beam sorts on it and
 64    #: the refinement pass tests it against an incumbent -- so its scale does not matter,
 65    #: but two things follow from the name. Ridge shrinkage can drive it slightly negative
 66    #: on a near-perfect fit, which is harmless for ordering and would not be for anything
 67    #: taking a square root of it; nothing does, and nothing should without clamping at the
 68    #: point of use rather than here, since clamping would flatten the ordering the beam
 69    #: depends on among near-perfect subsets.
 70    rss: float
 71
 72
 73def ridge_solve(gram: np.ndarray, rhs: np.ndarray, penalty: float) -> np.ndarray:
 74    """Solve ``(G + penalty*I) w = rhs`` for a centred, standardised design.
 75
 76    Centring both sides removes the intercept from the system, so no penalty is ever
 77    applied to it -- shrinking an intercept toward zero would bias every prediction
 78    toward MCC = 0, which is not a value this data goes anywhere near.
 79
 80    This is the **reference form**, written to be read. `Selector._evaluate` does not call
 81    it: it adds the penalty to the full Gram diagonal once at construction and indexes the
 82    submatrix straight out of that, which is the same arithmetic without a copy and an
 83    index array per solve. ``test_fit`` asserts the two agree, so this function is what
 84    the fast path is checked against rather than documentation of something else.
 85    """
 86    system = gram.copy()
 87    if penalty:
 88        system.flat[:: system.shape[0] + 1] += penalty
 89    try:
 90        return np.linalg.solve(system, rhs)
 91    except np.linalg.LinAlgError:
 92        return np.linalg.lstsq(system, rhs, rcond=None)[0]
 93
 94
 95def to_equation(
 96    library: Library,
 97    subset: Subset,
 98    standardizer: Standardizer,
 99    offset: float,
100    name: str,
101) -> Equation:
102    """Fold the standardisation back into the weights so the equation reads in raw units."""
103    order = list(subset.indices)
104    raw = subset.weights / standardizer.scale[order]
105    intercept = offset - float(raw @ standardizer.mean[order])
106    return Equation(
107        intercept=intercept,
108        terms=tuple(library.terms[index] for index in order),
109        weights=tuple(float(value) for value in raw),
110        standardized_weights=tuple(float(value) for value in subset.weights),
111        name=name,
112    )