ml_meta_perf.sweep

  1"""Choose the shared configuration: ``ml-meta-perf-search``.
  2
  3The study fits E1, E2 and E3 under **one** hyperparameter configuration (``config/study.json``),
  4chosen by E3-Valid's own criterion. This sweep makes that choice reproducible:
  5
  61. every point of the grid in the file's ``sweep`` section -- stability cap x ridge penalty x
  7   arity, with the pool, beam and horizon of its ``search`` section -- is fitted once with
  8   `experiment.fit_path` (E3, every feature) and cross-validated at every length;
  92. each point's length is chosen by the study's own rule, `selection.plateau_start`;
 103. points at or below ``readable_terms`` whose smoothed worst-protocol R2 is within ``band`` of
 11   the best such point are **indistinguishable on R2**, because adjacent configurations differ
 12   by less than adjacent lengths do;
 134. among those, the point whose doubly-held-out (DHO) predictions rank the models best --
 14   mean average precision over the twenty datasets -- is proposed; ties go to fewer terms,
 15   lower arity and then the heavier penalty.
 16
 17The DHO task metrics break the tie because they are what the equation is *for* -- ranking
 18models for a dataset nobody has run -- and they separate configurations that R2 cannot: at the
 19same floor, ranking AP ranges from 0.70 to 0.85.
 20
 21Nothing is changed in place. The run writes every curve, one row per configuration, and a
 22``proposed_study.json``; adopting it is a reviewed edit to ``config/study.json``.
 23
 24The search is small enough for a workstation: one fit per configuration, no re-selection
 25inside folds, one BLAS thread per worker.
 26"""
 27
 28from __future__ import annotations
 29
 30import dataclasses
 31import itertools
 32import json
 33import os
 34import sys
 35import time
 36from concurrent.futures import ProcessPoolExecutor
 37from pathlib import Path
 38from typing import Any
 39
 40import numpy as np
 41import polars as pl
 42from jsonargparse import ActionConfigFile, ArgumentParser  # pyright: ignore[reportPrivateImportUsage]
 43
 44from ml_meta_perf.config import (
 45    DEFAULT_CONFIG_PATH,
 46    Configuration,
 47    OpaqueConfig,
 48    SelectionConfig,
 49    StudyConfig,
 50    SweepConfig,
 51)
 52from ml_meta_perf.data import DATASET_COLUMN, DATASET_FEATURES, MODEL_FEATURES, groups, load, target
 53from ml_meta_perf.selection import floor_curve, plateau_start, smoothed
 54from ml_meta_perf.validate import decision_report, ranking_report
 55
 56#: The go/no-go threshold the decision metrics are reported at, as in chapter 5.
 57DECISION_THRESHOLD = 0.7
 58THREAD_VARIABLES = ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS")
 59
 60
 61def grid(study: StudyConfig) -> list[Configuration]:
 62    """Every configuration the sweep fits: the ``sweep`` grid over the ``search`` base."""
 63    return [
 64        dataclasses.replace(study.search, max_abs_zscore=zscore, penalty=penalty, max_arity=arity)
 65        for zscore, penalty, arity in itertools.product(study.sweep.zscores, study.sweep.penalties, study.sweep.arities)
 66    ]
 67
 68
 69def evaluate(
 70    config: Configuration, selection: SelectionConfig, data: str | None
 71) -> tuple[dict[str, Any], pl.DataFrame]:
 72    """One configuration: its curve, its chosen length, and that length's DHO task metrics."""
 73    from ml_meta_perf.experiment import fit_path
 74
 75    frame = load(data)
 76    truth, datasets = target(frame), groups(frame, DATASET_COLUMN)
 77    fitted = fit_path(frame, DATASET_FEATURES, MODEL_FEATURES, config, "E3")
 78    curve = fitted.curve
 79    size, how = plateau_start(curve, delta=selection.delta, window=selection.window, smoothing=selection.smoothing)
 80    position = list(curve["n_terms"]).index(size)
 81    at = curve.row(position, named=True)
 82    predictions = fitted.paths["loo_cell"][size].predictions
 83    ranking = ranking_report(truth, predictions, datasets)
 84    decision = decision_report(truth, predictions, datasets, thresholds=(DECISION_THRESHOLD,)).row(0, named=True)
 85    row = {
 86        "max_abs_zscore": config.max_abs_zscore,
 87        "penalty": config.penalty,
 88        "max_arity": config.max_arity,
 89        "n_terms": size,
 90        "reached_by": how,
 91        "smoothed_floor": float(smoothed(floor_curve(curve), selection.smoothing)[position]),
 92        "floor": float(floor_curve(curve)[position]),
 93        **{name: float(at[name]) for name in ("r2_in_sample", "r2_loo_dataset", "r2_loo_model", "r2_loo_cell")},
 94        "dho_ap": float(np.mean(ranking["ap"].to_numpy())),
 95        "dho_mrr": float(np.mean(ranking["mrr"].to_numpy())),
 96        "dho_hit_at_1": float(np.mean(ranking["hit_at_1"].to_numpy())),
 97        "dho_regret": float(np.mean(ranking["regret"].to_numpy())),
 98        "dho_accuracy": float(decision["accuracy"]),
 99        "dho_mcc": float(decision["mcc"]),
100        "terms": "; ".join(term.name for term in fitted.equations[size].terms),
101    }
102    labelled = curve.with_columns(
103        pl.lit(config.max_abs_zscore).alias("max_abs_zscore"),
104        pl.lit(config.penalty).alias("penalty"),
105        pl.lit(config.max_arity).alias("max_arity"),
106    )
107    return row, labelled
108
109
110def choose(candidates: pl.DataFrame, sweep: SweepConfig) -> dict[str, Any]:
111    """The proposed configuration: best DHO ranking AP within the readable R2 band."""
112    readable = candidates.filter(pl.col("n_terms") <= sweep.readable_terms)
113    if readable.height == 0:
114        raise ValueError(f"no configuration chose an equation of at most {sweep.readable_terms} terms")
115    best = float(readable["smoothed_floor"].max())  # pyright: ignore[reportArgumentType]
116    band = readable.filter(pl.col("smoothed_floor") >= best - sweep.band)
117    ranked = band.sort(["dho_ap", "n_terms", "max_arity", "penalty"], descending=[True, False, False, True])
118    return ranked.row(0, named=True)
119
120
121def _worker_initialiser() -> None:
122    for variable in THREAD_VARIABLES:
123        os.environ[variable] = "1"
124
125
126def build_parser() -> ArgumentParser:
127    parser = ArgumentParser(
128        prog="ml-meta-perf-search",
129        description="Sweep the shared hyperparameters and propose config/study.json.",
130        default_config_files=[str(DEFAULT_CONFIG_PATH)],
131    )
132    parser.add_argument("--config", action=ActionConfigFile, help="study configuration JSON")
133    parser.add_argument("--search", type=Configuration, help="the base configuration the grid varies")
134    parser.add_argument("--selection", type=SelectionConfig, help="the length rule applied to every point")
135    parser.add_argument("--opaque", type=OpaqueConfig, help="carried into the proposal unchanged")
136    parser.add_argument("--sweep", type=SweepConfig, help="the grid and the choice among close candidates")
137    parser.add_argument("--data", type=str | None, default=None, help="meta-dataset CSV")
138    parser.add_argument("--output", type=str, default="results/sweep", help="directory for the sweep outputs")
139    parser.add_argument("--jobs", type=int, default=0, help="parallel workers; 0 uses every CPU")
140    return parser
141
142
143def main(argv: list[str] | None = None) -> int:
144    parser = build_parser()
145    arguments = parser.instantiate(parser.parse_args(argv))
146    study = StudyConfig(
147        search=arguments.search, selection=arguments.selection, opaque=arguments.opaque, sweep=arguments.sweep
148    )
149    points = grid(study)
150    output = Path(arguments.output)
151    output.mkdir(parents=True, exist_ok=True)
152    jobs = arguments.jobs or os.cpu_count() or 1
153    print(f"{len(points)} configurations on {jobs} workers", flush=True)
154    started = time.perf_counter()
155    with ProcessPoolExecutor(max_workers=jobs, initializer=_worker_initialiser) as pool:
156        results = list(
157            pool.map(evaluate, points, itertools.repeat(study.selection), itertools.repeat(arguments.data), chunksize=1)
158        )
159    candidates = pl.DataFrame([row for row, _ in results]).sort(["max_abs_zscore", "penalty", "max_arity"])
160    pl.concat([curve for _, curve in results]).write_csv(output / "curves.csv")
161    candidates.write_csv(output / "candidates.csv")
162    chosen = choose(candidates, study.sweep)
163    proposed = dataclasses.replace(
164        study,
165        search=dataclasses.replace(
166            study.search,
167            max_abs_zscore=float(chosen["max_abs_zscore"]),
168            penalty=float(chosen["penalty"]),
169            max_arity=int(chosen["max_arity"]),
170        ),
171    )
172    (output / "proposed_study.json").write_text(json.dumps(dataclasses.asdict(proposed), indent=2) + "\n")
173    elapsed = time.perf_counter() - started
174    with pl.Config(tbl_rows=20, tbl_cols=12, float_precision=3):
175        print(candidates.sort("dho_ap", descending=True).drop("terms").head(10))
176    print(
177        f"\nproposed: z-cap {chosen['max_abs_zscore']}, penalty {chosen['penalty']}, arity {chosen['max_arity']} "
178        f"-> {chosen['n_terms']} terms ({chosen['reached_by']}), smoothed floor {chosen['smoothed_floor']:.4f}, "
179        f"DHO AP {chosen['dho_ap']:.4f}\nwritten to {output} in {elapsed:.0f} s"
180    )
181    return 0
182
183
184if __name__ == "__main__":
185    sys.exit(main())