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 empty = [k for k, v in idx.items() if len(v) == 0 and k in ("train", a.eval_split)]57 if empty:58 raise SystemExit(f"no eligible origins for split(s) {empty}: check splits/data window in {a.config}")59 return cfg, exp, out, man, rows, dirs, idx, tag606162def stage_baselines(a):63 from chronos_x import experiment as E64 cfg, exp, out, man, rows, dirs, idx, tag = setup(a)65 ev, tr, horizons = cfg["eval"], cfg["train"], cfg["eval"]["horizons"]66 idx_eval = idx[a.eval_split]67 sizes = {k: len(v) for k, v in idx.items()}68 print(f"[{exp}:baselines] {len(dirs)} grids, origins {sizes}", flush=True)69 write_json({"config": cfg, "symbols": [r["symbol"] for r in rows], "origins": sizes, "eval_split": a.eval_split, "heldout": a.heldout},70 out / f"run_meta_{tag}.json")71 t0 = time.time()72 which_b = tuple(cfg.get("baselines", ["naive", "har", "garch", "trees"]))73 garch_params = E.fit_garch_params(idx["train"]) if "garch" in which_b else None74 tab = P.origin_table(idx_eval, horizons, garch_params, workers=tr["num_workers"])75 tab.to_csv(out / f"origins_{tag}.csv", index=False)76 print(f"[{exp}:baselines] eval origin frame {len(tab)} rows in {time.time() - t0:.0f}s", flush=True)77 t0 = time.time()78 tab_tr = P.origin_table(idx["train"], horizons, None, workers=tr["num_workers"])79 tab_tr.to_csv(out / "origins_train.csv", index=False)80 preds = E.fit_and_predict_baselines(tab_tr, tab, horizons, out, tr["seed"], which_b)81 if garch_params is not None:82 write_json({str(i): p for i, p in garch_params.items()}, out / "garch_params.json")83 for name, pH in preds.items():84 E.score_predictions(f"{name}_{tag}", tab, horizons, out, ev["n_samples"], tr["seed"], param_preds=pH, universe=man)85 print(f"[{exp}:baselines] done in {time.time() - t0:.0f}s: {list(preds)}", flush=True)868788def stage_neural(a):89 import resource90 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)91 resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 65536), hard)) # memory-mapped grids x workers92 import torch93 torch.multiprocessing.set_sharing_strategy("file_system")94 from chronos_x import experiment as E95 from chronos_x.inference.predict import predict_index96 from chronos_x.inference.serialization import save_predictions97 from chronos_x.train import checkpoint as ck98 from chronos_x.eval.statistics import cluster_bootstrap_mean99 cfg, exp, out, man, rows, dirs, idx, tag = setup(a)100 ev, tr, horizons = cfg["eval"], cfg["train"], cfg["eval"]["horizons"]101 idx_eval = idx[a.eval_split]; n_grids = len(dirs)102 f_tab = out / f"origins_{tag}.csv"103 if f_tab.exists():104 tab = pd.read_csv(f_tab)105 else:106 tab = P.origin_table(idx_eval, horizons, None, workers=tr["num_workers"]); tab.to_csv(f_tab, index=False)107 print(f"[{exp}:neural] {n_grids} grids, origins { {k: len(v) for k, v in idx.items()} }, device {E.pick_device()}", flush=True)108 only = set(a.only_models.split(",")) if a.only_models else None109 for spec in cfg.get("models", []):110 name = spec["name"]111 if only and name not in only:112 continue113 mdir = (Path(a.models_from) / name) if a.models_from else out / name114 mdir.mkdir(parents=True, exist_ok=True)115 t0 = time.time()116 if a.skip_train or a.heldout or a.models_from:117 if spec.get("kwargs", {}).get("ticker_embedding") and a.heldout:118 print(f"[{exp}:neural] {name}: ticker embedding cannot be used zero-shot, skipping", flush=True); continue119 model = E.build_model(spec, 0 if a.heldout else n_grids)120 ck.load(mdir / "best.pt", model)121 if getattr(model, "copula", None) is not None and (mdir / "copula" / "best.pt").exists():122 ck.load(mdir / "copula" / "best.pt", model)123 summ = json.loads((mdir / "train_summary.json").read_text()) if (mdir / "train_summary.json").exists() else {}124 else:125 model, summ = E.train_variant(spec, idx["train"], idx["val"], cfg, mdir, n_grids)126 ck.load(mdir / "best.pt", model)127 if getattr(model, "copula", None) is not None:128 ck.load(mdir / "copula" / "best.pt", model)129 write_json(summ, mdir / "train_summary.json")130 joint = getattr(model, "copula", None) is not None131 kw = dict(n_samples=ev["n_samples"], batch_size=tr.get("eval_batch_size", 128), feature_set=spec.get("feature_set", "full"),132 num_workers=tr["num_workers"], seed=tr["seed"], keep_paths=spec.get("keep_paths", False))133 res = predict_index(model, idx_eval, horizons, shuffle_minutes=spec.get("shuffle_minutes_eval", False), joint=joint, **kw)134 mtag = f"{name}_{tag}"135 save_predictions(res, out / f"pred_{mtag}.npz")136 E.score_predictions(mtag, tab, horizons, out, ev["n_samples"], tr["seed"], neural=res, universe=man)137 if joint:138 res_i = predict_index(model, idx_eval, horizons, joint=False, **kw)139 E.score_predictions(f"{name}_indep_{tag}", tab, horizons, out, ev["n_samples"], tr["seed"], neural=res_i, universe=man)140 if spec.get("keep_paths"):141 cl = pd.factorize(tab["session_date"])[0]; ps = {}142 for lab, r in (("joint", res), ("indep", res_i)):143 ps[lab] = {k: cluster_bootstrap_mean(r[k], cl, n_boot=500) for k in ("energy", "variogram", "mdd_crps", "rv_crps")}144 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")}145 write_json(ps, out / f"path_scores_{mtag}.json")146 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)147148149def stage_tables(a):150 from chronos_x import experiment as E151 cfg = load_config(a.config); exp = a.exp_name or cfg["name"]152 out = ROOT / "results" / "metrics" / exp153 tag = a.eval_split + ("_heldout" if a.heldout else "")154 horizons = cfg["eval"]["horizons"]; ev = cfg["eval"]155 tab = pd.read_csv(out / f"origins_{tag}.csv")156 scored = {}157 for f in sorted(out.glob(f"summary_*_{tag}.json")):158 name = f.name[len("summary_"):-len(f"_{tag}.json")]159 per = np.load(out / f"per_origin_{name}_{tag}.npz")160 scored[name] = {"summary": {int(H): v for H, v in json.loads(f.read_text()).items()},161 "per": {H: {k[:-len(f"_{H}")]: per[k] for k in per.files if k.endswith(f"_{H}")} for H in horizons}}162 lb = E.leaderboard(scored, horizons); lb.to_csv(out / f"leaderboard_{tag}.csv", index=False)163 ref = cfg.get("reference", "naive_rw")164 if ref in scored:165 for metric in ("crps", "wis", "qlike"):166 E.comparisons(scored, tab, horizons, ev["stride"], metric, ref).to_csv(out / f"compare_{metric}_vs_{ref}_{tag}.csv", index=False)167 cols = ["model", "H", "crps_scaled", "wis", "qlike", "cov_50", "cov_80", "cov_95", "pit_ks", "dir_acc"]168 print(lb[[c for c in cols if c in lb.columns]].to_string(index=False, float_format=lambda x: f"{x:.5g}"))169170171def main() -> None:172 a = parse()173 if a.stage == "baselines":174 stage_baselines(a)175 elif a.stage == "neural":176 stage_neural(a)177 elif a.stage == "tables":178 stage_tables(a)179 else:180 passthrough = [x for x in sys.argv[1:] if x not in ("--stage", "all")]181 if "--stage" in sys.argv:182 i = sys.argv.index("--stage"); passthrough = [x for j, x in enumerate(sys.argv[1:], 1) if j not in (i, i + 1)]183 if not a.skip_baselines:184 subprocess.run([sys.executable, __file__, "--stage", "baselines", *passthrough], check=True)185 subprocess.run([sys.executable, __file__, "--stage", "neural", *passthrough], check=True)186 stage_tables(a)187188189if __name__ == "__main__":190 main()191