ml_meta_perf.model

  1"""The fitted object: an equation of the form ``MCC = w0 + w1*t1 + w2*t2 + ...``.
  2
  3Selection runs on standardised terms, because greedy selection compares candidates by
  4correlation with the residual and raw scales would make that comparison meaningless.
  5What gets published, though, must be evaluable as written, so the weights are folded
  6back into raw units before they are stored. Both sets are kept: the raw weights are the
  7equation, the standardised weights are how the terms rank against each other.
  8
  9Study chapter: [2. The additive model][study-chapter] -- the rationale, in
 10prose, with the figures.
 11
 12[study-chapter]: https://github.com/mariolpantunes/ml-meta-perf/blob/main/assets/docs/02-additive-model.md
 13"""
 14
 15from __future__ import annotations
 16
 17import json
 18from dataclasses import dataclass
 19from pathlib import Path
 20
 21import numpy as np
 22
 23from ml_meta_perf.terms import Term
 24
 25MCC_LOWER = -1.0
 26MCC_UPPER = 1.0
 27
 28
 29@dataclass(frozen=True)
 30class Equation:
 31    """A sparse linear equation over named terms."""
 32
 33    intercept: float
 34    terms: tuple[Term, ...]
 35    weights: tuple[float, ...]
 36    standardized_weights: tuple[float, ...]
 37    name: str = "equation"
 38
 39    def __post_init__(self) -> None:
 40        if not (len(self.terms) == len(self.weights) == len(self.standardized_weights)):
 41            raise ValueError("terms and weights must have the same length")
 42
 43    @property
 44    def n_terms(self) -> int:
 45        return len(self.terms)
 46
 47    def evaluate(self, columns: dict[str, np.ndarray]) -> np.ndarray:
 48        """Apply the equation as written, without clipping."""
 49        n = len(next(iter(columns.values())))
 50        total = np.full(n, self.intercept, dtype=np.float64)
 51        with np.errstate(all="ignore"):
 52            for term, weight in zip(self.terms, self.weights, strict=True):
 53                total += weight * term.evaluate(columns)
 54        return total
 55
 56    def predict(self, columns: dict[str, np.ndarray]) -> np.ndarray:
 57        """Apply the equation and clip into the range MCC can actually take.
 58
 59        A linear form has no idea that MCC stops at 1.0, and the meta-dataset
 60        includes exact values at that bound. Clipping is the domain knowledge imposed on the
 61        output, and it is applied identically during cross-validation.
 62        """
 63        return np.clip(self.evaluate(columns), MCC_LOWER, MCC_UPPER)
 64
 65    def ranked_terms(self) -> list[tuple[Term, float, float]]:
 66        """Terms ordered by standardised weight magnitude: strongest contributor first."""
 67        triples = list(zip(self.terms, self.weights, self.standardized_weights, strict=True))
 68        return sorted(triples, key=lambda item: abs(item[2]), reverse=True)
 69
 70    def __str__(self) -> str:
 71        lines = [f"MCC = {self.intercept:+.6g}"]
 72        for term, weight, standardized in self.ranked_terms():
 73            # Long terms overflow the column rather than being truncated, so the padding
 74            # has to keep a separator of its own -- ljust alone glues the comment to a
 75            # name that is already past the stop.
 76            lines.append(f"      {weight:+.6g} * {term.name}".ljust(64) + f"  # beta={standardized:+.4f}")
 77        return "\n".join(lines)
 78
 79    def to_latex(self) -> str:
 80        """A LaTeX rendering of the equation, for dropping straight into the paper."""
 81        parts = [f"{self.intercept:+.4g}"]
 82        for term, weight, _ in self.ranked_terms():
 83            parts.append(f"{weight:+.4g} \\cdot \\mathrm{{{_escape(term.name)}}}")
 84        return "\\mathrm{MCC} = " + " ".join(parts)
 85
 86    def to_dict(self) -> dict[str, object]:
 87        return {
 88            "name": self.name,
 89            "intercept": self.intercept,
 90            "terms": [term.to_dict() for term in self.terms],
 91            "weights": list(self.weights),
 92            "standardized_weights": list(self.standardized_weights),
 93        }
 94
 95    @classmethod
 96    def from_dict(cls, payload: dict[str, object]) -> Equation:
 97        terms = payload["terms"]
 98        weights = payload["weights"]
 99        standardized = payload["standardized_weights"]
100        if not isinstance(terms, list) or not isinstance(weights, list) or not isinstance(standardized, list):
101            raise ValueError("malformed equation payload")
102        return cls(
103            intercept=float(payload["intercept"]),  # pyright: ignore[reportArgumentType]
104            terms=tuple(Term.from_dict(item) for item in terms),  # pyright: ignore[reportArgumentType]
105            weights=tuple(float(value) for value in weights),  # pyright: ignore[reportArgumentType]
106            standardized_weights=tuple(float(value) for value in standardized),  # pyright: ignore[reportArgumentType]
107            name=str(payload.get("name", "equation")),
108        )
109
110    def save(self, path: str | Path) -> None:
111        Path(path).write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8")
112
113    @classmethod
114    def load(cls, path: str | Path) -> Equation:
115        return cls.from_dict(json.loads(Path(path).read_text(encoding="utf-8")))
116
117
118def _escape(name: str) -> str:
119    return name.replace("_", r"\_").replace("^", r"\^{}")
120
121
122#: How a term's direction is written in every exported table.
123#:
124#: One spelling, in one place, because `attribution.term_effects` and `report` both write a
125#: ``direction`` column and used to disagree -- "increases MCC" in one CSV, "raises MCC" in
126#: the next -- which reads as two different quantities to anyone joining them.
127RAISES, LOWERS = "raises MCC", "lowers MCC"
128
129
130def direction(signed: float) -> str:
131    """The word for the sign of a standardised weight or effect."""
132    return RAISES if signed > 0.0 else LOWERS