SPB Git

spb/wp40_uqo

Public MIT
5.2 KB · 41 lines python
1#!/usr/bin/env python32"""E6 baselines: the same Parquet exports processed with pandas, polars and duckdb (read → regular session filter →3log returns within session → per-session RV, BPV, RSV±, RQ). Wall-clock per stage and peak RSS. Usage: e6_baselines.py e0_universe.json out.jsonl [n_files] [reps]"""4import sys, json, time, resource, os, math5import numpy as np6e0 = json.load(open(sys.argv[1])); out = open(sys.argv[2], "a"); nf = int(sys.argv[3]) if len(sys.argv) > 3 else 4; reps = int(sys.argv[4]) if len(sys.argv) > 4 else 37def rss(): return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss8import pandas as pd, polars as pl, duckdb, pyarrow, pyarrow.parquet as pq9versions = {"pandas": pd.__version__, "polars": pl.__version__, "duckdb": duckdb.__version__, "pyarrow": pyarrow.__version__, "numpy": np.__version__}10MU1 = math.sqrt(2 / math.pi)11def measures_np(r):12    M = r.size; return dict(rv=float(np.sum(r * r)), bpv=float((M / (M - 1)) / MU1 ** 2 * np.sum(np.abs(r[1:]) * np.abs(r[:-1]))) if M > 1 else float("nan"), rsv_neg=float(np.sum(r[r < 0] ** 2)), rsv_pos=float(np.sum(r[r > 0] ** 2)), rq=float(M / 3 * np.sum(r ** 4)))13for ex in e0["exports"][:nf]:14    path = ex["path"]15    for rep in range(reps):16        # ---- pandas17        t0 = time.perf_counter(); df = pq.read_table(path, columns=["symbol_as_reported", "timestamp_utc", "close"]).to_pandas(); t1 = time.perf_counter()18        et = df["timestamp_utc"].dt.tz_convert("America/New_York"); df["session"] = et.dt.date; mins = et.dt.hour * 60 + et.dt.minute19        early = et.dt.strftime("%m-%d").isin(["11-29", "07-03", "12-24"]); close_min = np.where(early, 13 * 60, 16 * 60)   # approximation of the early-close rule (documented): exact calendar in RAID20        reg = df[(mins >= 570) & (mins <= close_min)].sort_values(["symbol_as_reported", "timestamp_utc"]); t2 = time.perf_counter()21        reg["lr"] = reg.groupby(["symbol_as_reported", "session"])["close"].transform(lambda s: np.log(s).diff()); reg = reg.dropna(subset=["lr"]); t3 = time.perf_counter()22        res = reg.groupby(["symbol_as_reported", "session"])["lr"].apply(lambda s: measures_np(s.to_numpy())); t4 = time.perf_counter()23        out.write(json.dumps({"experiment": "E6", "engine": "pandas", "file": path, "rep": rep, "bars": int(len(df)), "sessions": int(len(res)), "read_seconds": t1 - t0, "filter_seconds": t2 - t1, "returns_seconds": t3 - t2, "estimate_seconds": t4 - t3, "total_seconds": t4 - t0, "peak_rss": rss(), "versions": versions}) + "\n")24        # ---- polars25        t0 = time.perf_counter(); lf = pl.scan_parquet(path).select(["symbol_as_reported", "timestamp_utc", "close"]); dfp = lf.collect(); t1 = time.perf_counter()26        dfp = dfp.with_columns(et=pl.col("timestamp_utc").dt.convert_time_zone("America/New_York")).with_columns(session=pl.col("et").dt.date(), mins=pl.col("et").dt.hour().cast(pl.Int32) * 60 + pl.col("et").dt.minute().cast(pl.Int32), md=pl.col("et").dt.strftime("%m-%d"))27        dfp = dfp.with_columns(close_min=pl.when(pl.col("md").is_in(["11-29", "07-03", "12-24"])).then(13 * 60).otherwise(16 * 60)).filter((pl.col("mins") >= 570) & (pl.col("mins") <= pl.col("close_min"))).sort(["symbol_as_reported", "timestamp_utc"]); t2 = time.perf_counter()28        dfp = dfp.with_columns(lr=pl.col("close").log().diff().over(["symbol_as_reported", "session"])).drop_nulls("lr"); t3 = time.perf_counter()29        resp = dfp.group_by(["symbol_as_reported", "session"]).agg([(pl.col("lr") ** 2).sum().alias("rv"), (pl.col("lr").abs() * pl.col("lr").abs().shift(1)).sum().alias("adj"), pl.len().alias("M"), (pl.col("lr") ** 4).sum().alias("s4"), pl.when(pl.col("lr") < 0).then(pl.col("lr") ** 2).otherwise(0.0).sum().alias("rsv_neg")]); t4 = time.perf_counter()30        out.write(json.dumps({"experiment": "E6", "engine": "polars", "file": path, "rep": rep, "bars": int(dfp.height), "sessions": int(resp.height), "read_seconds": t1 - t0, "filter_seconds": t2 - t1, "returns_seconds": t3 - t2, "estimate_seconds": t4 - t3, "total_seconds": t4 - t0, "peak_rss": rss(), "versions": versions}) + "\n")31        # ---- duckdb32        t0 = time.perf_counter(); con = duckdb.connect()33        q = f"""WITH b AS (SELECT symbol_as_reported s, timestamp_utc AT TIME ZONE 'America/New_York' et, close FROM read_parquet('{path}')),34        r AS (SELECT s, CAST(et AS DATE) session, et, close, hour(et)*60+minute(et) mins, strftime(et, '%m-%d') md FROM b),35        f AS (SELECT * FROM r WHERE mins >= 570 AND mins <= CASE WHEN md IN ('11-29','07-03','12-24') THEN 780 ELSE 960 END),36        l AS (SELECT s, session, ln(close) - lag(ln(close)) OVER (PARTITION BY s, session ORDER BY et) lr FROM f)37        SELECT s, session, sum(lr*lr) rv, count(*) M, sum(lr*lr*lr*lr) s4, sum(CASE WHEN lr<0 THEN lr*lr ELSE 0 END) rsv_neg FROM l WHERE lr IS NOT NULL GROUP BY s, session"""38        resd = con.sql(q).to_arrow_table(); t4 = time.perf_counter(); con.close()39        out.write(json.dumps({"experiment": "E6", "engine": "duckdb", "file": path, "rep": rep, "sessions": int(resd.num_rows), "total_seconds": t4 - t0, "peak_rss": rss(), "versions": versions}) + "\n")40        out.flush(); print(f"{os.path.basename(path)} rep {rep}: pandas {res.shape[0]} sessions, polars {resp.height}, duckdb {resd.num_rows}", flush=True)41
Files · ae5a4e6ff6df81cef716eacbb5d51143a008d45b
Menu

Theme