ml_meta_perf.attribution

  1"""What each term is worth, and whether dataset or model features carry the equation.
  2
  3An equation is only useful as guidance if its terms can be compared, and raw weights
  4cannot be: they are in the units of whatever the term happens to compute. Two comparable
  5quantities are derived here instead.
  6
  7*Effect* is the swing in predicted MCC produced by moving a term from its 10th to its
  890th percentile, holding everything else fixed. It is stated in MCC units, which is what
  9makes one term's importance comparable with another's and with the target itself.
 10
 11*Share* attributes the variance of the equation's output to groups of terms. Terms are
 12labelled by the features they use -- dataset-only, model-only, or mixed -- so the split
 13between "how hard is this data" and "how capable is this model" can be read off the
 14fitted equation rather than assumed.
 15
 16Study chapter: [6. Best practices against the equation][study-chapter] -- the rationale, in
 17prose, with the figures.
 18
 19[study-chapter]: https://github.com/mariolpantunes/ml-meta-perf/blob/main/assets/docs/06-practices.md
 20"""
 21
 22from __future__ import annotations
 23
 24import numpy as np
 25import polars as pl
 26
 27from ml_meta_perf.model import Equation, direction
 28
 29DATASET_ONLY = "dataset"
 30MODEL_ONLY = "model"
 31MIXED = "mixed"
 32
 33
 34def classify(
 35    features: tuple[str, ...],
 36    dataset_features: tuple[str, ...],
 37    model_features: tuple[str, ...],
 38) -> str:
 39    """Label a term by which feature groups it draws on."""
 40    uses_dataset = any(feature in dataset_features for feature in features)
 41    uses_model = any(feature in model_features for feature in features)
 42    if uses_dataset and uses_model:
 43        return MIXED
 44    return MODEL_ONLY if uses_model else DATASET_ONLY
 45
 46
 47def contributions(equation: Equation, columns: dict[str, np.ndarray]) -> np.ndarray:
 48    """The per-row contribution of each term, as a ``(n_rows, n_terms)`` matrix."""
 49    if not equation.terms:
 50        return np.zeros((len(next(iter(columns.values()))), 0))
 51    with np.errstate(all="ignore"):
 52        return np.column_stack(
 53            [weight * term.evaluate(columns) for term, weight in zip(equation.terms, equation.weights, strict=True)]
 54        )
 55
 56
 57def term_effects(
 58    equation: Equation,
 59    columns: dict[str, np.ndarray],
 60    dataset_features: tuple[str, ...],
 61    model_features: tuple[str, ...],
 62) -> pl.DataFrame:
 63    """Per-term effect sizes, in MCC units, strongest first.
 64
 65    ``effect`` is the 10th-to-90th-percentile swing of the term's contribution: how much
 66    predicted MCC moves across the bulk of the observed range of that term. Percentiles
 67    rather than min-to-max so that a single extreme row cannot define the headline number.
 68    """
 69    matrix = contributions(equation, columns)
 70    rows: list[dict[str, object]] = []
 71    for index, term in enumerate(equation.terms):
 72        column = matrix[:, index]
 73        low, high = np.percentile(column, [10.0, 90.0])
 74        rows.append(
 75            {
 76                "term": term.name,
 77                "group": classify(term.features, dataset_features, model_features),
 78                "weight": equation.weights[index],
 79                "beta": equation.standardized_weights[index],
 80                "effect": float(high - low),
 81                "direction": direction(equation.standardized_weights[index]),
 82            }
 83        )
 84    return pl.DataFrame(rows).sort(pl.col("effect").abs(), descending=True)
 85
 86
 87def group_shares(
 88    equation: Equation,
 89    columns: dict[str, np.ndarray],
 90    dataset_features: tuple[str, ...],
 91    model_features: tuple[str, ...],
 92) -> pl.DataFrame:
 93    """How much of the equation's output variance each feature group drives.
 94
 95    Shares are computed as ``cov(group total, prediction) / var(prediction)``. That
 96    decomposition sums to exactly 1 even when the groups are correlated, which a naive
 97    ``var(group) / var(total)`` does not -- and the groups here are correlated, because
 98    three of the model features are functions of dataset size.
 99    """
100    matrix = contributions(equation, columns)
101    total = matrix.sum(axis=1)
102    variance = float(np.var(total))
103    labels = [classify(term.features, dataset_features, model_features) for term in equation.terms]
104
105    rows: list[dict[str, object]] = []
106    for group in (DATASET_ONLY, MODEL_ONLY, MIXED):
107        selected = [index for index, label in enumerate(labels) if label == group]
108        if not selected:
109            rows.append({"group": group, "n_terms": 0, "share": 0.0, "effect_sum": 0.0})
110            continue
111        subtotal = matrix[:, selected].sum(axis=1)
112        # Population covariance (ddof=0) to match np.var. np.cov defaults to ddof=1, and
113        # mixing the two inflates every share by n/(n-1) so the decomposition stops
114        # summing to 1.
115        covariance = float(((subtotal - subtotal.mean()) * (total - total.mean())).mean())
116        share = covariance / variance if variance > 1e-15 else 0.0
117        low, high = np.percentile(subtotal, [10.0, 90.0])
118        rows.append(
119            {
120                "group": group,
121                "n_terms": len(selected),
122                "share": share,
123                "effect_sum": float(high - low),
124            }
125        )
126    return pl.DataFrame(rows)
127
128
129def variance_decomposition(
130    target: np.ndarray,
131    datasets: np.ndarray,
132    models: np.ndarray,
133) -> pl.DataFrame:
134    """How much of MCC is explained by dataset identity alone, and by model identity alone.
135
136    This is measured on the target itself, independent of any equation, and is the honest
137    answer to "does the choice of model matter more than the choice of data". Each row is
138    the variance explained by knowing *only* that group label -- an upper bound on what
139    any equation over that group's features could reach.
140    """
141    total = float(((target - target.mean()) ** 2).sum())
142    rows: list[dict[str, object]] = []
143    for name, labels in (("dataset identity", datasets), ("model identity", models)):
144        prediction = np.zeros_like(target)
145        for label in np.unique(labels):
146            mask = labels == label
147            prediction[mask] = target[mask].mean()
148        residual = float(((target - prediction) ** 2).sum())
149        rows.append(
150            {
151                "knowing only": name,
152                "n_groups": int(np.unique(labels).shape[0]),
153                "variance_explained": 1.0 - residual / total if total > 1e-15 else 0.0,
154            }
155        )
156    return pl.DataFrame(rows)