1"""Run one experiment config end to end: indices -> baselines -> neural variants -> predictions -> scores -> tables.2LightGBM (baselines) and torch (neural) cannot share a process on macOS (duplicate libomp: segfault / Cholesky3deadlock, verified 2026-11-02), so the stages run as separate subprocesses and communicate through files.4Usage: run_experiment.py --config configs/pilot.yaml [--stage all|baselines|neural|tables] [--symbols-from pilot|feasibility|all]5 [--eval-split test|final] [--only-models m1,m2] [--skip-baselines] [--skip-train] [--exp-name NAME] [--heldout] [--models-from DIR]"""6from __future__ import annotations78import sys910if "baselines" in sys.argv: # LightGBM must be the first native library loaded in the baselines process11 import lightgbm # noqa: F4011213import argparse14import json15import subprocess16import time17from pathlib import Path1819import numpy as np20import pandas as pd2122ROOT = Path(__file__).resolve().parents[1]23sys.path.insert(0, str(ROOT / "src"))24from chronos_x.config import load_config # noqa: E40225from chronos_x import pipeline as P # noqa: E40226from chronos_x.data.windows import WindowConfig # noqa: E40227from chronos_x.inference.serialization import write_json # noqa: E402282930def parse():31 ap = argparse.ArgumentParser()32 ap.add_argument("--config", default="configs/pilot.yaml"); ap.add_argument("--stage", default="all")33 ap.add_argument("--symbols-from", default=None); ap.add_argument("--eval-split", default="test")34 ap.add_argument("--only-models", default=None); ap.add_argument("--skip-baselines", action="store_true")35 ap.add_argument("--skip-train", action="store_true", help="reuse saved checkpoints"); ap.add_argument("--exp-name", default=None)36 ap.add_argument("--heldout", action="store_true", help="evaluate on held-out instruments (zero-shot) with saved models")37 ap.add_argument("--models-from", default=None, help="experiment directory holding trained models to reuse")38 return ap.parse_args()394041def setup(a):42 from chronos_x import experiment as E43 cfg = load_config(a.config)44 exp = a.exp_name or cfg["name"]45 out = ROOT / "results" / "metrics" / exp; out.mkdir(parents=True, exist_ok=True)46 man = json.loads((ROOT / cfg["universe"]["manifest"]).read_text())47 which = a.symbols_from or cfg["universe"].get("symbols_from", "pilot")48 roles = ("heldout_instrument",) if a.heldout else ("pooled",)49 rows = E.symbols_for(man, which, roles)50 d = cfg["data"]51 dirs = P.grid_dirs(Path(cfg["paths"]["cache"]), d["start"], d["end"], [r["symbol"] for r in rows])52 wcfg = WindowConfig(**cfg["windows"]); ev = cfg["eval"]; tr = cfg["train"]53 idx = P.build_indices(dirs, cfg["splits"], wcfg, d["warmup_sessions"], tr["origins_per_session"], ev["stride"], tr["seed"],54 embargo_sessions=cfg["splits"].get("embargo_sessions", 1))55 tag = a.eval_split + ("_heldout" if a.heldout else "")56 return cfg, exp, out, man, rows, dirs, idx, tag575859def stage_baselines(a):60 from chronos_x import experiment as E61 cfg, exp, out, man, rows, dirs, idx, tag = setup(a)62 ev, tr, horizons = cfg["eval"], cfg["train"], cfg["eval"]["horizons"]63 idx_eval = idx[a.eval_split]64 sizes = {k: len(v) for k, v in idx.items()}65 print(f"[{exp}:baselines] {len(dirs)} grids, origins {sizes}", flush=True)66 write_json({"config": cfg, "symbols": [r["symbol"] for r in rows], "origins": sizes, "eval_split": a.eval_split, "heldout": a.heldout},67 out / f"run_meta_{tag}.json")68 t0 = time.time()69 which_b = tuple(cfg.get("baselines", ["naive", "har", "garch", "trees"]))70 garch_params = E.fit_garch_params(idx["train"]) if "garch" in which_b else None71 tab = P.origin_table(idx_eval, horizons, garch_params, workers=tr["num_workers"])72 tab.to_csv(out / f"origins_{tag}.csv", index=False)73 print(f"[{exp}:baselines] eval origin frame {len(tab)} rows in {time.time() - t0:.0f}s", flush=True)74 t0 = time.time()75 tab_tr = P.origin_table(idx["train"], horizons, None, workers=tr["num_workers"])76 tab_tr.to_csv(out / "origins_train.csv", index=False)77 preds = E.fit_and_predict_baselines(tab_tr, tab, horizons, out, tr["seed"], which_b)78 if garch_params is not None:79 write_json({str(i): p for i, p in garch_params.items()}, out / "garch_params.json")80 for name, pH in preds.items():81 E.score_predictions(f"{name}_{tag}", tab, horizons, out, ev["n_samples"], tr["seed"], param_preds=pH, universe=man)82 print(f"[{exp}:baselines] done in {time.time() - t0:.0f}s: {list(preds)}", flush=True)838485def stage_neural(a):86 import resource87 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)88 resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 65536), hard)) # memory-mapped grids x workers89 import torch90 torch.multiprocessing.set_sharing_strategy("file_system")91 from chronos_x import experiment as E92 from chronos_x.inference.predict import predict_index93 from chronos_x.inference.serialization import save_predictions94 from chronos_x.train import checkpoint as ck95 from chronos_x.eval.statistics import cluster_bootstrap_mean96 cfg, exp, out, man, rows, dirs, idx, tag = setup(a)97 ev, tr, horizons = cfg["eval"], cfg["train"], cfg["eval"]["horizons"]98 idx_eval = idx[a.eval_split]; n_grids = len(dirs)99 f_tab = out / f"origins_{tag}.csv"100 if f_tab.exists():101 tab = pd.read_csv(f_tab)102 else:103 tab = P.origin_table(idx_eval, horizons, None, workers=tr["num_workers"]); tab.to_csv(f_tab, index=False)104 print(f"[{exp}:neural] {n_grids} grids, origins { {k: len(v) for k, v in idx.items()} }, device {E.pick_device()}", flush=True)105 only = set(a.only_models.split(",")) if a.only_models else None106 for spec in cfg.get("models", []):107 name = spec["name"]108 if only and name not in only:109 continue110 mdir = (Path(a.models_from) / name) if a.models_from else out / name111 mdir.mkdir(parents=True, exist_ok=True)112 t0 = time.time()113 if a.skip_train or a.heldout or a.models_from:114 if spec.get("kwargs", {}).get("ticker_embedding") and a.heldout:115 print(f"[{exp}:neural] {name}: ticker embedding cannot be used zero-shot, skipping", flush=True); continue116 model = E.build_model(spec, 0 if a.heldout else n_grids)117 ck.load(mdir / "best.pt", model)118 if getattr(model, "copula", None) is not None and (mdir / "copula" / "best.pt").exists():119 ck.load(mdir / "copula" / "best.pt", model)120 summ = json.loads((mdir / "train_summary.json").read_text()) if (mdir / "train_summary.json").exists() else {}121 else:122 model, summ = E.train_variant(spec, idx["train"], idx["val"], cfg, mdir, n_grids)123 ck.load(mdir / "best.pt", model)124 if getattr(model, "copula", None) is not None:125 ck.load(mdir / "copula" / "best.pt", model)126 write_json(summ, mdir / "train_summary.json")127 joint = getattr(model, "copula", None) is not None128 kw = dict(n_samples=ev["n_samples"], batch_size=tr.get("eval_batch_size", 128), feature_set=spec.get("feature_set", "full"),129 num_workers=tr["num_workers"], seed=tr["seed"], keep_paths=spec.get("keep_paths", False))130 res = predict_index(model, idx_eval, horizons, shuffle_minutes=spec.get("shuffle_minutes_eval", False), joint=joint, **kw)131 mtag = f"{name}_{tag}"132 save_predictions(res, out / f"pred_{mtag}.npz")133 E.score_predictions(mtag, tab, horizons, out, ev["n_samples"], tr["seed"], neural=res, universe=man)134 if joint:135 res_i = predict_index(model, idx_eval, horizons, joint=False, **kw)136 E.score_predictions(f"{name}_indep_{tag}", tab, horizons, out, ev["n_samples"], tr["seed"], neural=res_i, universe=man)137 if spec.get("keep_paths"):138 cl = pd.factorize(tab["session_date"])[0]; ps = {}139 for lab, r in (("joint", res), ("indep", res_i)):140 ps[lab] = {k: cluster_bootstrap_mean(r[k], cl, n_boot=500) for k in ("energy", "variogram", "mdd_crps", "rv_crps")}141 ps["joint_minus_indep"] = {k: cluster_bootstrap_mean(res[k] - res_i[k], cl, n_boot=1000) for k in ("energy", "variogram", "mdd_crps", "rv_crps")}142 write_json(ps, out / f"path_scores_{mtag}.json")143 print(f"[{exp}:neural] {name}: {time.time() - t0:.0f}s, params={model.n_params()}, train={summ.get('steps')} steps, best_val={summ.get('best_val')}", flush=True)144145146def stage_tables(a):147 from chronos_x import experiment as E148 cfg = load_config(a.config); exp = a.exp_name or cfg["name"]149 out = ROOT / "results" / "metrics" / exp150 tag = a.eval_split + ("_heldout" if a.heldout else "")151 horizons = cfg["eval"]["horizons"]; ev = cfg["eval"]152 tab = pd.read_csv(out / f"origins_{tag}.csv")153 scored = {}154 for f in sorted(out.glob(f"summary_*_{tag}.json")):155 name = f.name[len("summary_"):-len(f"_{tag}.json")]156 per = np.load(out / f"per_origin_{name}_{tag}.npz")157 scored[name] = {"summary": {int(H): v for H, v in json.loads(f.read_text()).items()},158 "per": {H: {k[:-len(f"_{H}")]: per[k] for k in per.files if k.endswith(f"_{H}")} for H in horizons}}159 lb = E.leaderboard(scored, horizons); lb.to_csv(out / f"leaderboard_{tag}.csv", index=False)160 ref = cfg.get("reference", "naive_rw")161 if ref in scored:162 for metric in ("crps", "wis", "qlike"):163 E.comparisons(scored, tab, horizons, ev["stride"], metric, ref).to_csv(out / f"compare_{metric}_vs_{ref}_{tag}.csv", index=False)164 cols = ["model", "H", "crps_scaled", "wis", "qlike", "cov_50", "cov_80", "cov_95", "pit_ks", "dir_acc"]165 print(lb[[c for c in cols if c in lb.columns]].to_string(index=False, float_format=lambda x: f"{x:.5g}"))166167168def main() -> None:169 a = parse()170 if a.stage == "baselines":171 stage_baselines(a)172 elif a.stage == "neural":173 stage_neural(a)174 elif a.stage == "tables":175 stage_tables(a)176 else:177 passthrough = [x for x in sys.argv[1:] if x not in ("--stage", "all")]178 if "--stage" in sys.argv:179 i = sys.argv.index("--stage"); passthrough = [x for j, x in enumerate(sys.argv[1:], 1) if j not in (i, i + 1)]180 if not a.skip_baselines:181 subprocess.run([sys.executable, __file__, "--stage", "baselines", *passthrough], check=True)182 subprocess.run([sys.executable, __file__, "--stage", "neural", *passthrough], check=True)183 stage_tables(a)184185186if __name__ == "__main__":187 main()188