spb/uqo-research-db

1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617"""Build data/metadata/universe.json = S&P 500 (cached) ∪ Russell 3000 (iShares IWV).1819Optionally split the companies into per-node shards for the cluster run:2021    python scripts/build_universe.py --shards M3U96a:28,M2U64:24,... --bhs64-share 0.52223writes data/metadata/shards.json plus data/metadata/shards/<node>.{self,bhs64}.txt24(CIK lists).  Companies already present locally (the S&P 500 corpus) are pinned25to --master and are not re-sharded.26"""2728from __future__ import annotations2930import argparse31import json32import sys33from pathlib import Path3435sys.path.insert(0, str(Path(__file__).resolve().parent.parent))3637from src.acquisition.cik_mapper import CIKMapper38from src.acquisition.sp500 import SP500Universe39from src.acquisition.universe import Russell3000Universe, build_universe, load_universe40from src.utils.constants import DATA_DIRS41from src.utils.logging import setup_logging4243logger = setup_logging("build_universe")444546def parse_args() -> argparse.Namespace:47    p = argparse.ArgumentParser(description="Build the company universe (S&P 500 ∪ Russell 3000)")48    p.add_argument("--refresh-sp500", action="store_true", help="Re-scrape the S&P 500 list (default: cached)")49    p.add_argument("--skip-fetch", action="store_true", help="Reuse cached Russell 3000 file")50    p.add_argument("--shards", type=str, default=None,51                   help="node:cores,... for the new companies (excludes --master)")52    p.add_argument("--master", type=str, default="M3U96b", help="node holding the existing S&P 500 corpus")53    p.add_argument("--bhs64-share", type=float, default=0.5,54                   help="fraction of each node's new CIKs downloaded through BHS64 (second IP)")55    return p.parse_args()565758def main() -> None:59    args = parse_args()60    sp = SP500Universe()61    sp500 = sp.get_constituents() if args.refresh_sp500 else sp.load_cached()6263    mapper = CIKMapper()64    mapper.build_mapping()65    ru = Russell3000Universe()66    russell = ru.load_cached() if args.skip_fetch else ru.get_constituents(mapper)6768    universe = build_universe(sp500, russell)6970    if not args.shards:71        return7273    existing = {str(c["cik"]).zfill(10) for c in sp500 if c.get("cik")}74    new_ciks = [e["cik"] for e in universe if e["cik"] not in existing]75    nodes = []76    for item in args.shards.split(","):77        name, cores = item.split(":")78        nodes.append((name.strip(), int(cores)))79    total = sum(c for _, c in nodes)8081    # weight each node by cores; deal CIKs round-robin weighted so sizes are proportional82    shares = {n: c / total for n, c in nodes}83    assigned: dict[str, list[str]] = {n: [] for n, _ in nodes}84    assigned[args.master] = []85    quota = {n: shares[n] * len(new_ciks) for n, _ in nodes}86    # largest-remainder allocation on a stable CIK order87    order = sorted(new_ciks)88    idx = 089    for n, _ in nodes:90        k = int(round(quota[n]))91        assigned[n].extend(order[idx:idx + k])92        idx += k93    for leftover in order[idx:]:94        assigned[nodes[0][0]].append(leftover)9596    shard_dir = DATA_DIRS["metadata"] / "shards"97    shard_dir.mkdir(parents=True, exist_ok=True)98    manifest = {"master": args.master, "bhs64_share": args.bhs64_share, "nodes": {}}99    for n, ciks in assigned.items():100        ciks = sorted(ciks)101        n_bhs = int(round(len(ciks) * args.bhs64_share))102        bhs = ciks[:n_bhs]103        own = ciks[n_bhs:]104        if n == args.master:105            own = sorted(existing)  # refresh of the S&P 500 corpus on the master106            bhs = []107        (shard_dir / f"{n}.self.txt").write_text("\n".join(own) + ("\n" if own else ""))108        (shard_dir / f"{n}.bhs64.txt").write_text("\n".join(bhs) + ("\n" if bhs else ""))109        manifest["nodes"][n] = {"self": own, "bhs64": bhs, "n_total": len(own) + len(bhs)}110        logger.info("shard %-8s self=%4d bhs64=%4d", n, len(own), len(bhs))111    with open(DATA_DIRS["metadata"] / "shards.json", "w") as f:112        json.dump(manifest, f, indent=2)113    logger.info("Wrote %s (%d new CIKs over %d nodes)", DATA_DIRS["metadata"] / "shards.json",114                len(new_ciks), len(nodes))115116117if __name__ == "__main__":118    main()119