1"""P1: point-in-time stratified universe -> data/manifests/universe_main.json (+ pilot subset)."""2from __future__ import annotations34import argparse5import json6import random7import sys8from pathlib import Path910import pandas as pd1112sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))13from chronos_x.data.client import FinDataClient # noqa: E40214from chronos_x.data.universe import UniverseConfig, build_candidates, stratified_sample, manifest, load_sic_map # noqa: E4021516ROOT = Path(__file__).resolve().parents[1]171819def main() -> None:20 ap = argparse.ArgumentParser()21 ap.add_argument("--n-target", type=int, default=300)22 ap.add_argument("--sec-tickers", default="/tmp/wp41_ext/sec_company_tickers.json")23 ap.add_argument("--sec-submissions", default="/tmp/wp41_ext/sec_submissions")24 ap.add_argument("--n-pilot", type=int, default=20)25 ap.add_argument("--n-feasibility", type=int, default=50)26 ap.add_argument("--min-coverage", type=float, default=0.0)27 ap.add_argument("--coverage-file", default=None)28 ap.add_argument("--out", default="universe_main.json")29 a = ap.parse_args()30 cfg = UniverseConfig(n_target=a.n_target, min_coverage=a.min_coverage, coverage_file=a.coverage_file)31 c = FinDataClient(log_path=ROOT / "results" / "logs" / "provenance_universe.jsonl")32 sic = load_sic_map(Path(a.sec_tickers), Path(a.sec_submissions))33 print(f"SIC map: {len(sic)} tickers")34 cand = build_candidates(c, cfg, sic)35 print(f"candidates: {cand.shape[0]} (stocks {int((cand.asset_class=='stock').sum())}, etf {int((cand.asset_class=='etf').sum())}, delisted {int(cand.delisted.sum())})")36 sample = stratified_sample(cand, cfg)37 ds_version = None38 try:39 ds_version = c.dataset("market_bars.stock.1m.raw").get("version")40 except Exception:41 pass42 # pilot and feasibility subsets drawn from the pooled (non-held-out) instruments, stratified by asset class43 rng = random.Random(cfg.seed + 7)44 pooled = sample[sample["role"] == "pooled"]45 pilot = sorted(rng.sample(sorted(pooled[pooled.asset_class == "stock"]["symbol"]), int(a.n_pilot * 0.75)) +46 rng.sample(sorted(pooled[pooled.asset_class == "etf"]["symbol"]), a.n_pilot - int(a.n_pilot * 0.75)))47 rest = [s for s in sorted(pooled["symbol"]) if s not in pilot]48 feas = sorted(pilot + rng.sample(rest, a.n_feasibility - len(pilot)))49 man = manifest(sample, cand, cfg, ds_version, {"pilot_symbols": pilot, "feasibility_symbols": feas})50 out = ROOT / "data" / "manifests"51 out.mkdir(parents=True, exist_ok=True)52 (out / a.out).write_text(json.dumps(man, indent=1))53 cand.drop(columns=["name"]).to_csv(out / a.out.replace("universe", "candidates").replace(".json", ".csv"), index=False)54 summary = sample.groupby(["asset_class", "adv_tercile", "vol_tercile"]).size().unstack(fill_value=0)55 print(summary)56 print(sample.groupby(["asset_class", "group"]).size())57 print(f"selected {len(sample)}, held out {int((sample.role=='heldout_instrument').sum())}, delisted {int(sample.delisted.sum())}")58 print("pilot:", pilot)596061if __name__ == "__main__":62 main()63