Skip to content

Commit 547b4db

Browse files
fix(fast): commit missing feols.py + tests to repair CI/CD
src/statspai/fast/__init__.py was committed referencing `from .feols import feols, FeolsResult` but the feols.py module itself was never staged, so every push since aee6583 broke the CI/CD pytest collection step on all five OS/Python shards with: ModuleNotFoundError: No module named 'statspai.fast.feols' Ship the missing files: - src/statspai/fast/feols.py — native OLS HDFE estimator (closed-form WLS on FE-residualised X/y, vcov: iid/hc1/cr1, optional weights → WLS). Mirrors fepois result shape so downstream code (sp.fast.event_study, etable) treats them uniformly. - tests/test_fast_feols.py — 22-test suite (parity vs pyfixest, vcov matrix vs hand-rolled HC1/CR1, weights, formula DSL edge cases). - benchmarks/hdfe/run_feols_bench.py + feols_bench.json — wall-clock bench harness vs pyfixest / R fixest::feols, mirrors run_baseline.py. 22 + 5 (test_fast_bench, was failing on collection) = 27 tests now pass locally on Python 3.13.
1 parent ba8fb43 commit 547b4db

4 files changed

Lines changed: 1034 additions & 0 deletions

File tree

benchmarks/hdfe/feols_bench.json

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
{
2+
"datasets": {
3+
"medium": {
4+
"config": {
5+
"n": 1000000,
6+
"fe1_card": 100000,
7+
"fe2_card": 1000,
8+
"seed": 43
9+
},
10+
"n_rows": 1000000,
11+
"statspai": {
12+
"wall_min": 0.13518087519332767,
13+
"wall_mean": 0.1357012224228432,
14+
"n_runs": 3,
15+
"coef": {
16+
"x1": 0.29995277603470477,
17+
"x2": -0.2017701666154102
18+
},
19+
"se": {
20+
"x1": 0.0010533848798775225,
21+
"x2": 0.0010555318690953826
22+
}
23+
},
24+
"pyfixest": {
25+
"wall_min": 0.2095034159719944,
26+
"wall_mean": 0.21892262459732592,
27+
"n_runs": 3,
28+
"available": true,
29+
"coef": {
30+
"x1": 0.29995277603470355,
31+
"x2": -0.20177016661540814
32+
},
33+
"se": {
34+
"x1": 0.0010533854657401864,
35+
"x2": 0.0010555324561521368
36+
}
37+
},
38+
"fixest": {
39+
"wall_min": 0.106,
40+
"wall_mean": 0.120333333333333,
41+
"n_runs": 3,
42+
"coef": {
43+
"x1": 0.299952776034704,
44+
"x2": -0.201770166615406
45+
},
46+
"se": {
47+
"x1": 0.0010533854657402,
48+
"x2": 0.00105553245615213
49+
},
50+
"available": true
51+
},
52+
"coef_max_abs_diff_vs_fixest": 4.191091917959966e-15
53+
}
54+
}
55+
}

benchmarks/hdfe/run_feols_bench.py

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
"""End-to-end OLS HDFE benchmark — ``sp.fast.feols`` vs pyfixest vs R fixest.
2+
3+
Mirror of ``run_baseline.py`` (Phase 0 fepois bench) for the linear path.
4+
Usage:
5+
6+
python3 benchmarks/hdfe/run_feols_bench.py [--small] [--medium]
7+
8+
Reports wall-clock + coef diff vs R ``fixest::feols`` (the reference);
9+
falls back to pyfixest as the cross-check when R is unavailable. Output
10+
is JSON to stdout and a copy to ``benchmarks/hdfe/feols_bench.json``.
11+
"""
12+
from __future__ import annotations
13+
14+
import argparse
15+
import json
16+
import shutil
17+
import subprocess
18+
import time
19+
from pathlib import Path
20+
from typing import Any, Dict, List
21+
22+
import numpy as np
23+
import pandas as pd
24+
25+
26+
# ---------------------------------------------------------------------------
27+
# DGP
28+
# ---------------------------------------------------------------------------
29+
30+
def _make_panel(n: int, fe1_card: int, fe2_card: int, seed: int) -> pd.DataFrame:
31+
rng = np.random.default_rng(seed)
32+
fe1 = rng.integers(0, fe1_card, size=n)
33+
fe2 = rng.integers(0, fe2_card, size=n)
34+
x1 = rng.normal(size=n)
35+
x2 = rng.normal(size=n)
36+
a1 = rng.normal(0, 0.5, size=fe1_card)[fe1]
37+
a2 = rng.normal(0, 0.3, size=fe2_card)[fe2]
38+
cl = rng.normal(0, 0.4, size=fe1_card)[fe1]
39+
eps = cl + rng.normal(size=n)
40+
y = 0.30 * x1 - 0.20 * x2 + a1 + a2 + eps
41+
return pd.DataFrame({
42+
"y": y, "x1": x1, "x2": x2,
43+
"fe1": fe1.astype(np.int32), "fe2": fe2.astype(np.int32),
44+
})
45+
46+
47+
# ---------------------------------------------------------------------------
48+
# Backends
49+
# ---------------------------------------------------------------------------
50+
51+
def _time(fn, n_warm: int = 1, n_runs: int = 3) -> Dict[str, Any]:
52+
"""Run ``fn()`` ``n_warm`` warmups + ``n_runs`` timed runs; return summary."""
53+
for _ in range(n_warm):
54+
fn()
55+
times: List[float] = []
56+
last_result = None
57+
for _ in range(n_runs):
58+
t0 = time.perf_counter()
59+
last_result = fn()
60+
times.append(time.perf_counter() - t0)
61+
return {
62+
"wall_min": min(times),
63+
"wall_mean": float(np.mean(times)),
64+
"n_runs": n_runs,
65+
"result": last_result,
66+
}
67+
68+
69+
def _bench_statspai(df: pd.DataFrame) -> Dict[str, Any]:
70+
import statspai as sp
71+
def _run():
72+
return sp.fast.feols("y ~ x1 + x2 | fe1 + fe2", df, vcov="iid")
73+
summary = _time(_run)
74+
fit = summary.pop("result")
75+
summary["coef"] = {k: float(fit.coef()[k]) for k in ("x1", "x2")}
76+
summary["se"] = {k: float(fit.se()[k]) for k in ("x1", "x2")}
77+
return summary
78+
79+
80+
def _bench_pyfixest(df: pd.DataFrame) -> Dict[str, Any]:
81+
try:
82+
import pyfixest as pf
83+
except ImportError:
84+
return {"available": False, "reason": "pyfixest not installed"}
85+
def _run():
86+
return pf.feols(
87+
fml="y ~ x1 + x2 | fe1 + fe2", data=df,
88+
vcov="iid", fixef_rm="singleton",
89+
)
90+
summary = _time(_run)
91+
fit = summary.pop("result")
92+
summary["available"] = True
93+
summary["coef"] = {k: float(fit.coef()[k]) for k in ("x1", "x2")}
94+
summary["se"] = {k: float(fit.se()[k]) for k in ("x1", "x2")}
95+
return summary
96+
97+
98+
def _bench_r_fixest(df: pd.DataFrame, tmp_csv: Path) -> Dict[str, Any]:
99+
if shutil.which("Rscript") is None:
100+
return {"available": False, "reason": "Rscript not on PATH"}
101+
df.to_csv(tmp_csv, index=False)
102+
r_script = (
103+
"suppressMessages({library(data.table); library(fixest); library(jsonlite)})\n"
104+
f"d <- fread('{tmp_csv}')\n"
105+
"# Warm + 3 timed runs to mirror the Python harness\n"
106+
"for (i in 1:1) { invisible(feols(y ~ x1 + x2 | fe1 + fe2, data=d)) }\n"
107+
"ts <- numeric(3)\n"
108+
"for (i in 1:3) {\n"
109+
" t0 <- proc.time()['elapsed']\n"
110+
" f <- feols(y ~ x1 + x2 | fe1 + fe2, data=d)\n"
111+
" ts[i] <- as.numeric(proc.time()['elapsed'] - t0)\n"
112+
"}\n"
113+
"out <- list(\n"
114+
" wall_min = min(ts), wall_mean = mean(ts), n_runs = 3,\n"
115+
" coef = as.list(coef(f)), se = as.list(se(f))\n"
116+
")\n"
117+
"cat(toJSON(out, auto_unbox=TRUE, digits=14))\n"
118+
)
119+
proc = subprocess.run(
120+
["Rscript", "-e", r_script], capture_output=True, text=True, timeout=300,
121+
)
122+
if proc.returncode != 0:
123+
return {"available": False, "reason": f"Rscript failed: {proc.stderr[:200]}"}
124+
out = json.loads(proc.stdout.strip().splitlines()[-1])
125+
out["available"] = True
126+
return out
127+
128+
129+
# ---------------------------------------------------------------------------
130+
# Driver
131+
# ---------------------------------------------------------------------------
132+
133+
DATASETS = {
134+
"small": dict(n=100_000, fe1_card=1_000, fe2_card=50, seed=42),
135+
"medium": dict(n=1_000_000, fe1_card=100_000, fe2_card=1_000, seed=43),
136+
}
137+
138+
139+
def main(argv=None):
140+
parser = argparse.ArgumentParser()
141+
parser.add_argument("--small", action="store_true",
142+
help="Run only the 100k-row dataset")
143+
parser.add_argument("--medium", action="store_true",
144+
help="Run only the 1M-row dataset")
145+
args = parser.parse_args(argv)
146+
147+
sizes = []
148+
if args.small:
149+
sizes = ["small"]
150+
elif args.medium:
151+
sizes = ["medium"]
152+
else:
153+
sizes = ["small", "medium"]
154+
155+
out_dir = Path(__file__).resolve().parent
156+
tmp_csv = out_dir / "_bench_panel.csv"
157+
158+
report: Dict[str, Any] = {"datasets": {}}
159+
for size in sizes:
160+
cfg = DATASETS[size]
161+
df = _make_panel(**cfg)
162+
ds_report: Dict[str, Any] = {"config": cfg, "n_rows": len(df)}
163+
print(f"=== {size} ({len(df):,} rows, fe1={cfg['fe1_card']:,}, "
164+
f"fe2={cfg['fe2_card']:,}) ===")
165+
ds_report["statspai"] = _bench_statspai(df)
166+
print(f" statspai : {ds_report['statspai']['wall_min']*1000:.0f} ms")
167+
ds_report["pyfixest"] = _bench_pyfixest(df)
168+
if ds_report["pyfixest"].get("available"):
169+
print(f" pyfixest : {ds_report['pyfixest']['wall_min']*1000:.0f} ms")
170+
ds_report["fixest"] = _bench_r_fixest(df, tmp_csv)
171+
if ds_report["fixest"].get("available"):
172+
print(f" R fixest : {ds_report['fixest']['wall_min']*1000:.0f} ms")
173+
174+
# Cross-backend coef diff: reference is fixest if available, else pyfixest
175+
if ds_report["fixest"].get("available"):
176+
ref = ds_report["fixest"]
177+
ref_name = "fixest"
178+
elif ds_report["pyfixest"].get("available"):
179+
ref = ds_report["pyfixest"]
180+
ref_name = "pyfixest"
181+
else:
182+
ref = None
183+
ref_name = "none"
184+
if ref is not None:
185+
diffs = {
186+
k: abs(ds_report["statspai"]["coef"][k] - float(ref["coef"][k]))
187+
for k in ("x1", "x2")
188+
}
189+
ds_report["coef_max_abs_diff_vs_" + ref_name] = max(diffs.values())
190+
print(f" coef diff vs {ref_name}: max {max(diffs.values()):.3e}")
191+
192+
report["datasets"][size] = ds_report
193+
194+
out_json = out_dir / "feols_bench.json"
195+
with out_json.open("w") as f:
196+
json.dump(report, f, indent=2, default=float)
197+
print(f"\nWrote {out_json}")
198+
if tmp_csv.exists():
199+
tmp_csv.unlink()
200+
return report
201+
202+
203+
if __name__ == "__main__":
204+
main()

0 commit comments

Comments
 (0)