From e26ca1e17898fca50cfc5cc08f4721546aeea855 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Jul 2026 15:51:12 -0400 Subject: [PATCH 1/2] bench: add FE-absorption benchmark lane with BEFORE baselines + pyfixest yardstick Seven realistic workloads (scenarios 7-13 in docs/performance-scenarios.md) timing the MAP-demeaning hot path (demean_by_groups / within_transform): county policy event study, firm panel with churn, scanner store-week, 5M-order geo experiment, survey BRR replicates, correlated-FE stress, and a small-panel guard. Subprocess-isolated multi-run protocol with CV reporting (all committed baselines at CV 0.3-8.6%), ATT/SE identity capture with a --check-estimates gate (validated pass and fail directions), and an optional pyfixest yardstick lane guarded on import (coefficient parity < 1e-6 on the exact-estimand scenarios; 6e-13 to 2.7e-9 measured). Measurement-only: no library behavior change, no new package dependencies; benchmarks are not pytest-collected (verified). The committed BEFORE baseline records the tail_stress non-convergence warning as evidence of the max_iter=100 contract issue the follow-up optimization PR addresses. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 10 + benchmarks/speed_review/README.md | 35 ++ .../baselines/fe_absorption_before.json | 149 ++++++++ .../baselines/fe_absorption_pyfixest.json | 69 ++++ .../speed_review/bench_fe_absorption.py | 348 ++++++++++++++++++ .../bench_fe_absorption_pyfixest.py | 179 +++++++++ .../speed_review/fe_absorption_datagen.py | 235 ++++++++++++ .../speed_review/gen_findings_tables.py | 100 +++-- docs/performance-plan.md | 64 ++++ docs/performance-scenarios.md | 59 +++ 10 files changed, 1216 insertions(+), 32 deletions(-) create mode 100644 benchmarks/speed_review/baselines/fe_absorption_before.json create mode 100644 benchmarks/speed_review/baselines/fe_absorption_pyfixest.json create mode 100644 benchmarks/speed_review/bench_fe_absorption.py create mode 100644 benchmarks/speed_review/bench_fe_absorption_pyfixest.py create mode 100644 benchmarks/speed_review/fe_absorption_datagen.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1478d7f39..9bdc70dec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **FE-absorption benchmark suite** (`benchmarks/speed_review/bench_fe_absorption.py`, scenarios + 7-13 in `docs/performance-scenarios.md`): seven realistic workloads (county policy event study, + firm panel with churn, scanner store-week, 5M-order geo experiment, survey BRR replicates, + correlated-FE stress, small-panel guard) timing the MAP-demeaning hot path + (`demean_by_groups` / `within_transform`) with subprocess isolation, multi-run CV reporting, and + ATT/SE identity capture so optimization PRs can prove estimates are unchanged + (`--check-estimates`). Includes an optional pyfixest yardstick lane + (`bench_fe_absorption_pyfixest.py`, guarded on import - never a dependency) asserting < 1e-6 + coefficient parity on the exact-estimand scenarios, and committed BEFORE baselines + (`baselines/fe_absorption_before.json`). Measurement-only: no library behavior change. - **`HeterogeneousAdoptionDiD` cluster-robust SE on the continuous paths** (Phase 2a). `cluster=` is now threaded into `bias_corrected_local_linear` on the `continuous_at_zero` / `continuous_near_d_lower` designs, so the CCT-2014 robust variance becomes cluster-robust and the diff --git a/benchmarks/speed_review/README.md b/benchmarks/speed_review/README.md index 33a77c392..c622ca1fa 100644 --- a/benchmarks/speed_review/README.md +++ b/benchmarks/speed_review/README.md @@ -34,8 +34,16 @@ benchmarks/speed_review/ ├── bench_memory_scaling.py # peak-RSS sweep for the memory-scaling │ # work (B1 #561 / B2 #563 / C #567); │ # subprocess-isolated ru_maxrss, median +├── bench_fe_absorption.py # Scenarios 7-13: MAP-demeaning hot path +│ # (subprocess-isolated, multi-run CV, +│ # ATT/SE identity capture + gate) +├── bench_fe_absorption_pyfixest.py # optional external yardstick (guarded +│ # on `import pyfixest`; never a dep) +├── fe_absorption_datagen.py # seeded DGPs shared by both FE lanes └── baselines/ # this effort's output ├── memory_scaling_{before,after}.json # peak RSS pre-#561 vs current + ├── fe_absorption_{before,after}.json # FE-absorption timings + identity + ├── fe_absorption_pyfixest.json # yardstick timings + parity (optional) ├── _.json # phase-level wall-clock + peak RSS ├── mem_profile_brfss_large_.txt # tracemalloc top-N sites └── profiles/ # flame HTMLs (gitignored) @@ -76,6 +84,33 @@ Multi-scale scenarios write per-scale outputs (e.g. `dose_response_rust.json`). Full runtime for all scales × both backends is ~90 seconds on Apple Silicon M4. +### FE-absorption suite (scenarios 7-13) + +Standalone like `bench_memory_scaling.py` (not part of `run_all.py` - the 5M-row +scenarios are too heavy for the routine sweep): + +```bash +# Full suite (~20-40 min on M4; strictly sequential subprocesses by design) +python benchmarks/speed_review/bench_fe_absorption.py \ + --out benchmarks/speed_review/baselines/fe_absorption_before.json + +# Smoke test / one scenario +python benchmarks/speed_review/bench_fe_absorption.py --quick +python benchmarks/speed_review/bench_fe_absorption.py --only geo_experiment + +# After an optimization: regenerate + prove estimates did not move +python benchmarks/speed_review/bench_fe_absorption.py \ + --out benchmarks/speed_review/baselines/fe_absorption_after.json \ + --check-estimates benchmarks/speed_review/baselines/fe_absorption_before.json + +# Optional external yardstick (skipped cleanly if pyfixest is absent) +pip install pyfixest +python benchmarks/speed_review/bench_fe_absorption_pyfixest.py +``` + +Do not run anything else on the machine during a baseline run - the committed +JSONs carry a CV field per scenario and the driver flags CV > 10% as unusable. + ## Where to look for findings [`docs/performance-plan.md`](../../docs/performance-plan.md) - "Practitioner diff --git a/benchmarks/speed_review/baselines/fe_absorption_before.json b/benchmarks/speed_review/baselines/fe_absorption_before.json new file mode 100644 index 000000000..2502b72a2 --- /dev/null +++ b/benchmarks/speed_review/baselines/fe_absorption_before.json @@ -0,0 +1,149 @@ +{ + "suite": "fe_absorption", + "platform": "macOS-26.5.1-arm64-arm-64bit-Mach-O", + "python": "3.14.4", + "versions": { + "diff_diff": "3.6.1", + "numpy": "2.5.0", + "pandas": "3.0.3" + }, + "backend_requested": "python", + "repeats": 3, + "quick": false, + "tolerances": { + "att_atol": 1e-09, + "se_rtol": 1e-07, + "se_rtol_survey": 1e-06, + "gate_exempt": [ + "tail_stress" + ] + }, + "results": [ + { + "scenario": "county_policy", + "n_obs": 177289, + "checksum": 6315.071127371084, + "att": 0.28037854139903007, + "se": 0.031244864610216142, + "fit_median_s": 3.8655, + "fit_min_s": 3.7841, + "fit_max_s": 4.0366, + "fit_cv": 0.0253, + "noisy": false, + "n_timed_fits": 9, + "datagen_s": 0.005, + "peak_rss_mb": 3547.5, + "backend_resolved": "python", + "warnings": [] + }, + { + "scenario": "firm_churn", + "n_obs": 2400000, + "checksum": 240997.06223418942, + "att": 0.3144327184497805, + "se": 0.005837281701254032, + "fit_median_s": 92.9567, + "fit_min_s": 91.5167, + "fit_max_s": 93.0237, + "fit_cv": 0.0092, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.075, + "peak_rss_mb": 11550.3, + "backend_resolved": "python", + "warnings": [] + }, + { + "scenario": "scanner_twfe", + "n_obs": 3255000, + "checksum": 76832.22571739106, + "att": 0.2511641314956207, + "se": 0.0024149387421607484, + "fit_median_s": 1.5491, + "fit_min_s": 1.5475, + "fit_max_s": 1.5551, + "fit_cv": 0.0026, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.08, + "peak_rss_mb": 980.1, + "backend_resolved": "python", + "warnings": [] + }, + { + "scenario": "geo_experiment", + "n_obs": 5000000, + "checksum": 443696.3834000407, + "att": 0.2519437259991462, + "se": 0.0017886316514778143, + "fit_median_s": 2.6299, + "fit_min_s": 2.6156, + "fit_max_s": 2.6465, + "fit_cv": 0.0059, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.062, + "peak_rss_mb": 1631.5, + "backend_resolved": "python", + "warnings": [ + "Rank-deficient design matrix: dropping 1 of 4 columns (column 1). Coefficients for these columns are set to NA" + ] + }, + { + "scenario": "survey_absorb", + "n_obs": 500000, + "checksum": -176482.6014001448, + "att": 0.24553586050769946, + "se": 0.007662904200749773, + "fit_median_s": 7.3845, + "fit_min_s": 7.3301, + "fit_max_s": 8.5032, + "fit_cv": 0.0856, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.109, + "peak_rss_mb": 3621.1, + "backend_resolved": "python", + "warnings": [ + "Rank-deficient design matrix: dropping 1 of 4 columns (column 1). Coefficients for these columns are set to NA" + ] + }, + { + "scenario": "tail_stress", + "n_obs": 5000000, + "checksum": 376703.7248141762, + "att": 0.25020825600819235, + "se": 0.004392831033439421, + "fit_median_s": 26.2715, + "fit_min_s": 26.125, + "fit_max_s": 26.3224, + "fit_cv": 0.0039, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.072, + "peak_rss_mb": 1632.3, + "backend_resolved": "python", + "warnings": [ + "Rank-deficient design matrix: dropping 1 of 4 columns (column 1). Coefficients for these columns are set to NA", + "demean_by_groups alternating projection (variables: ['y', 'post', '_treat_time']) did not converge in 100 iter" + ] + }, + { + "scenario": "guard_small", + "n_obs": 20000, + "checksum": -571.6190664591859, + "att": 0.2515087011451766, + "se": 0.028726073333820948, + "fit_median_s": 0.0046, + "fit_min_s": 0.0044, + "fit_max_s": 0.0051, + "fit_cv": 0.0414, + "noisy": false, + "n_timed_fits": 21, + "datagen_s": 0.001, + "peak_rss_mb": 145.1, + "backend_resolved": "python", + "warnings": [] + } + ] +} \ No newline at end of file diff --git a/benchmarks/speed_review/baselines/fe_absorption_pyfixest.json b/benchmarks/speed_review/baselines/fe_absorption_pyfixest.json new file mode 100644 index 000000000..b03f65ffd --- /dev/null +++ b/benchmarks/speed_review/baselines/fe_absorption_pyfixest.json @@ -0,0 +1,69 @@ +{ + "suite": "fe_absorption_pyfixest", + "pyfixest": "0.60.0", + "platform": "macOS-26.5.1-arm64-arm-64bit-Mach-O", + "repeats": 3, + "quick": false, + "results": [ + { + "scenario": "county_policy", + "n_obs": 177289, + "checksum": 6315.071127371084, + "formula": "y ~ i(rel_time, ref=-1) | unit + time", + "coef": null, + "fit_median_s": 0.1904, + "fit_cv": 0.0401, + "parity_ok": null + }, + { + "scenario": "firm_churn", + "n_obs": 2400000, + "checksum": 240997.06223418942, + "formula": "y ~ i(rel_time, ref=-1) | unit + time", + "coef": null, + "fit_median_s": 1.9116, + "fit_cv": 0.2027, + "parity_ok": null + }, + { + "scenario": "scanner_twfe", + "n_obs": 3255000, + "checksum": 76832.22571739106, + "formula": "y ~ d | unit + post", + "coef": 0.2511641314950137, + "fit_median_s": 0.6441, + "fit_cv": 0.1229, + "parity_ok": true + }, + { + "scenario": "geo_experiment", + "n_obs": 5000000, + "checksum": 443696.3834000407, + "formula": "y ~ d | store + week", + "coef": 0.2519437259991667, + "fit_median_s": 0.6228, + "fit_cv": 0.078, + "parity_ok": true + }, + { + "scenario": "tail_stress", + "n_obs": 5000000, + "checksum": 376703.7248141762, + "formula": "y ~ d | store + week", + "coef": 0.2502082587524758, + "fit_median_s": 3.0748, + "fit_cv": 0.0149, + "parity_ok": true + }, + { + "scenario": "guard_small", + "n_obs": 20000, + "checksum": -571.6190664591859, + "formula": "y ~ d | unit + post", + "coef": 0.2515087011451808, + "fit_median_s": 0.0066, + "fit_cv": 0.0222, + "parity_ok": true + } + ] +} \ No newline at end of file diff --git a/benchmarks/speed_review/bench_fe_absorption.py b/benchmarks/speed_review/bench_fe_absorption.py new file mode 100644 index 000000000..316801d73 --- /dev/null +++ b/benchmarks/speed_review/bench_fe_absorption.py @@ -0,0 +1,348 @@ +"""FE-absorption (MAP demeaning) wall-clock + identity benchmark suite. + +Measures the fixed-effects absorption hot path across seven realistic +practitioner workloads (see ``fe_absorption_datagen.py`` and +docs/performance-scenarios.md for shape provenance). This is the before/after +surface for the demean_by_groups optimization work motivated by the 2026-07 +pyfixest gap analysis: the ``_before``/``_after`` baseline JSONs committed +under ``baselines/`` carry both the timings and the ATT/SE identity fields. + +Noise protocol +-------------- +Every (scenario, repeat) runs in a fresh subprocess, strictly sequential. +Small scenarios time a warmup fit then the median of several in-process fits; +large scenarios time one fit per subprocess. The driver pools all timed fits +across subprocesses and reports median / min / max / CV, flagging CV > 10% +as too noisy to trust. + +Identity gate +------------- +``--check-estimates baselines/fe_absorption_before.json`` compares each +scenario's ATT/SE against a previously committed baseline at the unified +tolerances (ATT atol=1e-9; SE rtol=1e-7; survey_absorb SE rtol=1e-6 because +the ~5.8e-11 accumulation-order drift stacks across 80 replicate re-demeans). +``tail_stress`` is excluded from the gate: its estimate is EXPECTED to move +when the MAP iteration cap is raised (it currently stops at max_iter without +converging) - the driver prints the delta instead. + +Usage:: + + python benchmarks/speed_review/bench_fe_absorption.py # full sweep + python benchmarks/speed_review/bench_fe_absorption.py --quick # smoke + python benchmarks/speed_review/bench_fe_absorption.py --only geo_experiment + python benchmarks/speed_review/bench_fe_absorption.py --repeats 3 + python benchmarks/speed_review/bench_fe_absorption.py \ + --out benchmarks/speed_review/baselines/fe_absorption_after.json \ + --check-estimates benchmarks/speed_review/baselines/fe_absorption_before.json +""" + +import argparse +import json +import os +import platform +import resource +import statistics +import subprocess +import sys +import time +import warnings +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import fe_absorption_datagen as datagen # noqa: E402 + +# Unified identity tolerances (single source for PR-A capture and PR-B gate). +ATT_ATOL = 1e-9 +SE_RTOL = 1e-7 +SE_RTOL_SURVEY = 1e-6 # drift stacks across ~80 replicate re-demeans +GATE_EXEMPT = ("tail_stress",) # estimate expected to move when the cap rises +CV_FLAG = 0.10 + +# scenario -> number of timed in-process fits (after one warmup fit). +# Sub-second scenarios repeat in-process so the subprocess median is stable. +INPROC_FITS = { + "county_policy": 3, + "firm_churn": 1, + "scanner_twfe": 1, + "geo_experiment": 1, + "survey_absorb": 1, + "tail_stress": 1, + "guard_small": 7, +} + + +def _peak_rss_mb(): + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # macOS ru_maxrss is bytes; Linux is KiB. + return peak / (1024 * 1024) if platform.system() == "Darwin" else peak / 1024 + + +# ------------------------------------------------------------------- fitting +def _fit(scenario, df): + """Run the scenario's estimator front door once; return (att, se, extras).""" + if scenario in ("county_policy", "firm_churn"): + from diff_diff import SunAbraham + + res = SunAbraham().fit(df, outcome="y", unit="unit", time="time", first_treat="first_treat") + elif scenario in ("scanner_twfe", "guard_small"): + from diff_diff import TwoWayFixedEffects + + res = TwoWayFixedEffects().fit( + df, outcome="y", treatment="treated", time="post", unit="unit" + ) + elif scenario in ("geo_experiment", "tail_stress"): + from diff_diff import DifferenceInDifferences + + res = DifferenceInDifferences().fit( + df, outcome="y", treatment="treated", time="post", absorb=["store", "week"] + ) + elif scenario == "survey_absorb": + from diff_diff import DifferenceInDifferences + from diff_diff.survey import SurveyDesign + + rep_cols = [c for c in df.columns if c.startswith("rw")] + design = SurveyDesign(weights="w", replicate_weights=rep_cols, replicate_method="BRR") + res = DifferenceInDifferences().fit( + df, + outcome="y", + treatment="treated", + time="post", + absorb=["state", "month"], + survey_design=design, + ) + else: + raise ValueError(f"unknown scenario {scenario!r}") + return float(res.att), float(res.se) + + +def _warmup(scenario): + """Tiny fit through the same front door (imports, caches, BLAS init).""" + df, _ = datagen.build(scenario, quick=True, seed=1) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _fit(scenario, df) + + +# ------------------------------------------------------------------ worker IO +def _worker(cfg_json): + cfg = json.loads(cfg_json) + scenario, quick = cfg["scenario"], cfg["quick"] + + t0 = time.perf_counter() + df, meta = datagen.build(scenario, quick=quick) + t_datagen = time.perf_counter() - t0 + + _warmup(scenario) + + fit_times, warn_msgs = [], set() + att = se = None + for _ in range(INPROC_FITS[scenario] if not quick else 1): + t1 = time.perf_counter() + with warnings.catch_warnings(record=True) as wl: + warnings.simplefilter("always") + att, se = _fit(scenario, df) + fit_times.append(time.perf_counter() - t1) + warn_msgs.update(str(w.message)[:110] for w in wl) + + try: + from diff_diff._backend import HAS_RUST_BACKEND + except Exception: + HAS_RUST_BACKEND = False + env_b = os.environ.get("DIFF_DIFF_BACKEND", "auto") + resolved = "python" if (env_b == "python" or not HAS_RUST_BACKEND) else "rust" + + print( + "FE_RESULT " + + json.dumps( + { + "fit_times": fit_times, + "t_datagen": t_datagen, + "att": att, + "se": se, + "n_obs": meta["n_obs"], + "checksum": meta["checksum"], + "peak_rss_mb": _peak_rss_mb(), + "backend_resolved": resolved, + "warnings": sorted(warn_msgs), + } + ) + ) + + +def _measure(scenario, repeats, backend, quick): + env = {**os.environ, "DIFF_DIFF_BACKEND": backend} + cfg = json.dumps({"scenario": scenario, "quick": quick}) + payloads = [] + for _ in range(repeats): + p = subprocess.run( + [sys.executable, __file__, "--worker", cfg], + capture_output=True, + text=True, + env=env, + ) + line = next((ln for ln in p.stdout.splitlines() if ln.startswith("FE_RESULT ")), None) + if line is None: + return None, p.stderr.strip().splitlines()[-3:] or ["(no output)"] + payloads.append(json.loads(line[len("FE_RESULT ") :])) + return payloads, None + + +# --------------------------------------------------------------- identity gate +def check_estimates(results, before_path): + """Compare ATT/SE per scenario against a prior baseline JSON. Returns the + number of failures (0 = pass). tail_stress reports its delta, ungated. + """ + before = {r["scenario"]: r for r in json.loads(Path(before_path).read_text())["results"]} + failures = 0 + for r in results: + scen = r["scenario"] + b = before.get(scen) + if b is None or r.get("att") is None or b.get("att") is None: + print(f" identity {scen:16s} SKIP (missing in one side)") + continue + d_att = abs(r["att"] - b["att"]) + rel_se = abs(r["se"] - b["se"]) / max(abs(b["se"]), 1e-300) + se_rtol = SE_RTOL_SURVEY if scen == "survey_absorb" else SE_RTOL + if scen in GATE_EXEMPT: + print( + f" identity {scen:16s} EXEMPT (expected shift): " + f"|d att|={d_att:.3e} rel d se={rel_se:.3e}" + ) + continue + ok = d_att <= ATT_ATOL and rel_se <= se_rtol + status = "ok" if ok else "FAIL" + print( + f" identity {scen:16s} {status}: |d att|={d_att:.3e} (atol {ATT_ATOL}) " + f"rel d se={rel_se:.3e} (rtol {se_rtol})" + ) + failures += 0 if ok else 1 + return failures + + +# ----------------------------------------------------------------------- main +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--worker") + ap.add_argument("--quick", action="store_true") + ap.add_argument("--repeats", type=int, default=3, help="subprocess repeats") + ap.add_argument("--only", choices=datagen.SCENARIO_IDS) + ap.add_argument( + "--backend", + choices=("python", "rust", "auto"), + default="python", + help="DIFF_DIFF_BACKEND for the fits (default python, matching the committed baselines)", + ) + ap.add_argument("--out", default="benchmarks/speed_review/baselines/fe_absorption.json") + ap.add_argument( + "--check-estimates", + metavar="BEFORE_JSON", + help="compare ATT/SE against a prior baseline at the unified tolerances", + ) + args = ap.parse_args() + + if args.worker: + _worker(args.worker) + return + + if args.repeats < 1: + ap.error("--repeats must be >= 1") + scenarios = [args.only] if args.only else list(datagen.SCENARIO_IDS) + + print( + f"FE-absorption suite | repeats={args.repeats} | backend={args.backend} | " + f"{platform.platform()} | py{sys.version.split()[0]}" + ) + results = [] + for scen in scenarios: + payloads, err = _measure(scen, args.repeats, args.backend, args.quick) + if payloads is None: + print(f" {scen:16s} -> ERROR: {err}") + results.append({"scenario": scen, "error": err}) + continue + pooled = [t for p in payloads for t in p["fit_times"]] + med = statistics.median(pooled) + cv = (statistics.stdev(pooled) / statistics.mean(pooled)) if len(pooled) > 1 else 0.0 + checksums = {p["checksum"] for p in payloads} + assert len(checksums) == 1, f"{scen}: data not deterministic across runs" + atts = {p["att"] for p in payloads} + assert len(atts) == 1, f"{scen}: estimate not deterministic across runs" + last = payloads[-1] + noisy = cv > CV_FLAG + results.append( + { + "scenario": scen, + "n_obs": last["n_obs"], + "checksum": last["checksum"], + "att": last["att"], + "se": last["se"], + "fit_median_s": round(med, 4), + "fit_min_s": round(min(pooled), 4), + "fit_max_s": round(max(pooled), 4), + "fit_cv": round(cv, 4), + "noisy": noisy, + "n_timed_fits": len(pooled), + "datagen_s": round(last["t_datagen"], 3), + "peak_rss_mb": round(max(p["peak_rss_mb"] for p in payloads), 1), + "backend_resolved": last["backend_resolved"], + "warnings": last["warnings"], + } + ) + flag = " ** CV>10%, rerun **" if noisy else "" + print( + f" {scen:16s} n={last['n_obs']:>10,} median={med:8.3f}s " + f"[{min(pooled):.3f}, {max(pooled):.3f}] cv={cv:.1%} " + f"rss={last['peak_rss_mb']:.0f}MB{flag}" + ) + if last["warnings"]: + for w in last["warnings"]: + print(f" warn: {w}") + + exit_code = 0 + if args.check_estimates: + print("\nidentity gate vs", args.check_estimates) + failures = check_estimates([r for r in results if "error" not in r], args.check_estimates) + if failures: + print(f"IDENTITY GATE FAILED: {failures} scenario(s) moved") + exit_code = 1 + + try: + import diff_diff + + versions = { + "diff_diff": diff_diff.__version__, + "numpy": __import__("numpy").__version__, + "pandas": __import__("pandas").__version__, + } + except Exception: + versions = {} + payload = { + "suite": "fe_absorption", + "platform": platform.platform(), + "python": sys.version.split()[0], + "versions": versions, + "backend_requested": args.backend, + "repeats": args.repeats, + "quick": args.quick, + "tolerances": { + "att_atol": ATT_ATOL, + "se_rtol": SE_RTOL, + "se_rtol_survey": SE_RTOL_SURVEY, + "gate_exempt": list(GATE_EXEMPT), + }, + "results": results, + } + try: + with open(args.out, "w") as f: + json.dump(payload, f, indent=2) + print(f"\nwrote {args.out}") + except OSError as e: + print(f"\n(could not write {args.out}: {e})") + + if any("error" in r for r in results): + exit_code = 1 + sys.exit(exit_code) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/speed_review/bench_fe_absorption_pyfixest.py b/benchmarks/speed_review/bench_fe_absorption_pyfixest.py new file mode 100644 index 000000000..4210b3630 --- /dev/null +++ b/benchmarks/speed_review/bench_fe_absorption_pyfixest.py @@ -0,0 +1,179 @@ +"""Optional pyfixest yardstick for the FE-absorption suite. + +Runs the same seeded scenarios as ``bench_fe_absorption.py`` through +``pyfixest.feols`` so the committed diff-diff baselines can be read against +the state of the art, not only against our own past. pyfixest is NOT a +dependency of diff-diff or of the test suite - this script exits cleanly if +it is not installed (mirror of the guarded R lanes in ``benchmarks/R/``). + +Honesty checks per scenario: + +- ``scanner_twfe`` / ``geo_experiment`` / ``tail_stress`` / ``guard_small`` + estimate the identical regression (``y ~ d | fe1 + fe2``); the yardstick + asserts coefficient agreement < 1e-6 against the diff-diff baseline JSON. +- ``county_policy`` / ``firm_churn`` are Sun-Abraham event studies; pyfixest + 0.60 has no ``sunab()``, so the yardstick runs the saturated + ``i(rel_time)`` event-study regression - the same demeaning load (outcome + + many interaction columns through unit+time FE) but a different estimand + weighting, so it is TIMING-ONLY (no coefficient assertion). +- ``survey_absorb`` is skipped: pyfixest has no BRR replicate-weight + refit path, so there is no comparable computation. + +Usage:: + + pip install pyfixest # in any environment alongside diff-diff + python benchmarks/speed_review/bench_fe_absorption_pyfixest.py \ + [--quick] [--repeats 3] [--only geo_experiment] \ + [--baseline benchmarks/speed_review/baselines/fe_absorption_before.json] +""" + +import argparse +import json +import platform +import statistics +import sys +import time +import warnings +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import fe_absorption_datagen as datagen # noqa: E402 + +try: + import pyfixest as pf +except ImportError: + print( + "pyfixest is not installed - the yardstick lane is optional; " + "run `pip install pyfixest` to enable it." + ) + sys.exit(0) + +COEF_PARITY_TOL = 1e-6 +PARITY_SCENARIOS = ("scanner_twfe", "geo_experiment", "tail_stress", "guard_small") +TIMING_ONLY = ("county_policy", "firm_churn") +SKIPPED = ("survey_absorb",) + + +def _formula_and_frame(scenario, df): + """Return (formula, frame, coef_name) for the pyfixest equivalent.""" + if scenario in ("scanner_twfe", "guard_small"): + return "y ~ d | unit + post", df, "d" + if scenario in ("geo_experiment", "tail_stress"): + return "y ~ d | store + week", df, "d" + if scenario in TIMING_ONLY: + # Saturated event study: same demeaning load as Sun-Abraham + # (outcome + interaction block through unit+time FE), timing-only. + frame = df.copy() + rel = frame["time"] - frame["first_treat"] + # never-treated (first_treat == 0) pinned to the reference bin + frame["rel_time"] = np.where(frame["first_treat"] > 0, rel, -1) + return "y ~ i(rel_time, ref=-1) | unit + time", frame, None + raise ValueError(f"no pyfixest mapping for {scenario!r}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--quick", action="store_true") + ap.add_argument("--repeats", type=int, default=3) + ap.add_argument("--only", choices=datagen.SCENARIO_IDS) + ap.add_argument( + "--baseline", + default="benchmarks/speed_review/baselines/fe_absorption_before.json", + help="diff-diff baseline JSON for the coefficient-parity check", + ) + ap.add_argument( + "--out", default="benchmarks/speed_review/baselines/fe_absorption_pyfixest.json" + ) + args = ap.parse_args() + + baseline = {} + bp = Path(args.baseline) + if bp.exists(): + baseline = {r["scenario"]: r for r in json.loads(bp.read_text())["results"] if "att" in r} + else: + print(f"(no diff-diff baseline at {args.baseline}; parity checks skipped)") + + scenarios = [args.only] if args.only else [s for s in datagen.SCENARIO_IDS if s not in SKIPPED] + print( + f"pyfixest yardstick {pf.__version__} | repeats={args.repeats} | " f"{platform.platform()}" + ) + + results, failures = [], 0 + for scen in scenarios: + if scen in SKIPPED: + print(f" {scen:16s} SKIP (no BRR replicate path in pyfixest)") + continue + df, meta = datagen.build(scen, quick=args.quick) + fml, frame, coef_name = _formula_and_frame(scen, df) + + # JIT/compile warmup on the quick-size frame through the same formula + wdf, _ = datagen.build(scen, quick=True, seed=1) + wfml, wframe, _ = _formula_and_frame(scen, wdf) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pf.feols(wfml, data=wframe, vcov="hetero") + + times, coef = [], None + for _ in range(args.repeats): + t0 = time.perf_counter() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + m = pf.feols(fml, data=frame, vcov="hetero") + times.append(time.perf_counter() - t0) + if coef_name is not None: + coef = float(m.coef()[coef_name]) + med = statistics.median(times) + cv = (statistics.stdev(times) / statistics.mean(times)) if len(times) > 1 else 0.0 + + parity = None + if coef_name is not None and scen in baseline: + delta = abs(coef - baseline[scen]["att"]) + parity = delta <= COEF_PARITY_TOL + if scen in PARITY_SCENARIOS and not parity: + failures += 1 + print( + f" {scen:16s} n={meta['n_obs']:>10,} median={med:8.3f}s cv={cv:.1%} " + f"|d coef vs diff-diff|={delta:.2e} {'ok' if parity else 'FAIL'}" + ) + else: + print( + f" {scen:16s} n={meta['n_obs']:>10,} median={med:8.3f}s cv={cv:.1%} " + f"(timing-only)" + ) + results.append( + { + "scenario": scen, + "n_obs": meta["n_obs"], + "checksum": meta["checksum"], + "formula": fml, + "coef": coef, + "fit_median_s": round(med, 4), + "fit_cv": round(cv, 4), + "parity_ok": parity, + } + ) + + payload = { + "suite": "fe_absorption_pyfixest", + "pyfixest": pf.__version__, + "platform": platform.platform(), + "repeats": args.repeats, + "quick": args.quick, + "results": results, + } + try: + with open(args.out, "w") as f: + json.dump(payload, f, indent=2) + print(f"\nwrote {args.out}") + except OSError as e: + print(f"\n(could not write {args.out}: {e})") + if failures: + print(f"COEFFICIENT PARITY FAILED for {failures} scenario(s)") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/speed_review/fe_absorption_datagen.py b/benchmarks/speed_review/fe_absorption_datagen.py new file mode 100644 index 000000000..db6f593ce --- /dev/null +++ b/benchmarks/speed_review/fe_absorption_datagen.py @@ -0,0 +1,235 @@ +"""Deterministic DGPs for the FE-absorption benchmark suite. + +Shared by ``bench_fe_absorption.py`` (diff-diff lane) and +``bench_fe_absorption_pyfixest.py`` (optional external yardstick). Both lanes +regenerate data locally from the same seeded PCG64 stream, so no data files +cross process boundaries; each scenario records a checksum of the outcome +column so the driver can prove both lanes saw identical matrices. + +Scenario shapes are calibrated to practitioner workloads, not to flatter the +library (see docs/performance-scenarios.md, FE-absorption suite section, for +the provenance of each): + +- ``county_policy`` 3,109 counties (the US county count) x 60 months, + staggered adoption, ~5% attrition - the everyday + applied-micro event study. +- ``firm_churn`` 100k firms x 40 quarters with contiguous entry/exit + lifetimes (~60% of the window) - Compustat/LBD-style + churn; correlated FE incidence is the slow regime for + alternating projections. +- ``scanner_twfe`` 35k stores x 156 weeks with entry/exit - Nielsen-scanner + shape for the TWFE workhorse. +- ``geo_experiment`` 5M order-level rows, 100k stores + 1k weeks - the + static-geo experiment shape from the Instacart pyfixest + writeup (tech.instacart.com, 2025). +- ``survey_absorb`` 500k microdata rows, state+month FE, 80 BRR replicate + weights - the replicate-refit multiplier path (each + replicate re-runs the weighted demean). +- ``tail_stress`` 5M rows with contiguous ~20% store lifetimes - the + correlated-FE stress regime (~250+ MAP iterations). + Reported separately; NOT a headline scenario. +- ``guard_small`` balanced 1k x 20 panel - regression guard so the small- + data path never regresses. +""" + +import numpy as np +import pandas as pd + +TAU = 0.25 # true treatment effect in every DGP + + +def checksum(df, col="y"): + return float(np.asarray(df[col], dtype=np.float64).sum()) + + +# --------------------------------------------------------------------- panels +def make_panel(n_units, n_periods, seed=42, drop_frac=0.10, lifetime_span=None): + """2x2 block-adoption panel: half the units treated after the midpoint. + + drop_frac > 0 drops rows at random (mild unbalance). lifetime_span in + (0, 1] instead gives each unit a contiguous observation window covering + that fraction of the periods (entry/exit churn - correlated FE incidence). + """ + rng = np.random.default_rng(seed) + unit = np.repeat(np.arange(n_units, dtype=np.int64), n_periods) + time = np.tile(np.arange(n_periods, dtype=np.int64), n_units) + alpha = rng.normal(0.0, 1.0, n_units) + gamma = rng.normal(0.0, 0.5, n_periods) + treated = (unit < n_units // 2).astype(np.int8) + post = (time >= n_periods // 2).astype(np.int8) + d = (treated * post).astype(np.int8) + y = alpha[unit] + gamma[time] + TAU * d + rng.normal(0.0, 1.0, unit.size) + df = pd.DataFrame( + {"y": y, "treated": treated, "post": post, "d": d, "unit": unit, "time": time} + ) + if lifetime_span is not None: + span = max(2, int(n_periods * lifetime_span)) + entry = rng.integers(0, n_periods - span + 1, n_units) + keep = (time >= entry[unit]) & (time < entry[unit] + span) + df = df.loc[keep].reset_index(drop=True) + elif drop_frac: + keep = rng.random(len(df)) >= drop_frac + df = df.loc[keep].reset_index(drop=True) + return df + + +def make_staggered_panel(n_units, n_periods, cohorts, seed=42, drop_frac=0.0, lifetime_span=None): + """Staggered-adoption panel for SunAbraham: cohorts = adoption periods + (0 entries mean never-treated). Equal cohort shares including never. + Dynamic effect ramps with event time (heterogeneous-by-cohort trends are + not needed to exercise the demeaning machinery). + """ + rng = np.random.default_rng(seed) + unit = np.repeat(np.arange(n_units, dtype=np.int64), n_periods) + time = np.tile(np.arange(n_periods, dtype=np.int64), n_units) + choices = np.asarray(list(cohorts) + [0], dtype=np.int64) + g_unit = rng.choice(choices, n_units) + first = g_unit[unit] + rel = time - first + d = ((first > 0) & (rel >= 0)).astype(np.int8) + alpha = rng.normal(0.0, 1.0, n_units) + gamma = rng.normal(0.0, 0.5, n_periods) + effect = np.where(d == 1, TAU * (1.0 + 0.05 * np.clip(rel, 0, 10)), 0.0) + y = alpha[unit] + gamma[time] + effect + rng.normal(0.0, 1.0, unit.size) + df = pd.DataFrame({"y": y, "unit": unit, "time": time, "first_treat": first, "d": d}) + if lifetime_span is not None: + span = max(2, int(n_periods * lifetime_span)) + entry = rng.integers(0, n_periods - span + 1, n_units) + t = df["time"].values + u = df["unit"].values + keep = (t >= entry[u]) & (t < entry[u] + span) + df = df.loc[keep].reset_index(drop=True) + elif drop_frac: + keep = rng.random(len(df)) >= drop_frac + df = df.loc[keep].reset_index(drop=True) + return df + + +def make_orders(n_rows, n_stores, n_weeks, seed=42, lifetime_span=None): + """Order-level static-geo experiment: store + week FE, store-level + treatment switched on for the post half of the weeks. lifetime_span + gives each store a contiguous active week window (the correlated-FE + regime); None gives uniform random incidence (the experiment regime). + """ + rng = np.random.default_rng(seed) + store = rng.integers(0, n_stores, n_rows).astype(np.int64) + if lifetime_span is not None: + span = max(2, int(n_weeks * lifetime_span)) + entry = rng.integers(0, n_weeks - span + 1, n_stores) + week = (entry[store] + rng.integers(0, span, n_rows)).astype(np.int64) + else: + week = rng.integers(0, n_weeks, n_rows).astype(np.int64) + a = rng.normal(0.0, 1.0, n_stores) + b = rng.normal(0.0, 0.5, n_weeks) + treated = (rng.random(n_stores) < 0.5).astype(np.int8)[store] + post = (week >= n_weeks // 2).astype(np.int8) + d = (treated * post).astype(np.int8) + y = a[store] + b[week] + TAU * d + rng.normal(0.0, 1.0, n_rows) + return pd.DataFrame( + {"y": y, "treated": treated, "post": post, "d": d, "store": store, "week": week} + ) + + +def make_survey_absorb(n_rows, n_states, n_months, n_replicates, seed=42): + """State-policy microdata with sampling weights and BRR replicate weights. + + Half the states adopt at the month midpoint. Replicate weights follow the + standard BRR half-sample convention (w * 2 on the selected half, w * 0 on + the other), supplied as pre-computed columns rw1..rwK the way BRFSS-style + public-use files ship them. + """ + rng = np.random.default_rng(seed) + state = rng.integers(0, n_states, n_rows).astype(np.int64) + month = rng.integers(0, n_months, n_rows).astype(np.int64) + a = rng.normal(0.0, 1.0, n_states) + b = rng.normal(0.0, 0.3, n_months) + treated = (state < n_states // 2).astype(np.int8) + post = (month >= n_months // 2).astype(np.int8) + d = (treated * post).astype(np.int8) + y = a[state] + b[month] + TAU * d + rng.normal(0.0, 1.0, n_rows) + w = rng.lognormal(0.0, 0.5, n_rows) + cols = { + "y": y, + "treated": treated, + "post": post, + "d": d, + "state": state, + "month": month, + "w": w, + } + # BRR half-samples drawn at the state level (states act as the PSUs). + for r in range(1, n_replicates + 1): + half = rng.random(n_states) < 0.5 + cols[f"rw{r}"] = w * np.where(half[state], 2.0, 0.0) + return pd.DataFrame(cols) + + +# ----------------------------------------------------------- scenario registry +def build(scenario, quick=False, seed=42): + """Return (DataFrame, meta) for a scenario id. quick=True shrinks every + dimension ~100x for smoke runs; the committed baselines use quick=False. + """ + q = quick + if scenario == "county_policy": + df = make_staggered_panel( + n_units=300 if q else 3_109, + n_periods=12 if q else 60, + cohorts=(6, 8) if q else (24, 30, 36, 42), + seed=seed, + drop_frac=0.05, + ) + elif scenario == "firm_churn": + df = make_staggered_panel( + n_units=1_000 if q else 100_000, + n_periods=12 if q else 40, + cohorts=(6, 8) if q else (16, 20, 24, 28), + seed=seed, + lifetime_span=0.6, + ) + elif scenario == "scanner_twfe": + df = make_panel( + n_units=500 if q else 35_000, + n_periods=12 if q else 156, + seed=seed, + lifetime_span=0.6, + ) + elif scenario == "geo_experiment": + df = make_orders( + n_rows=50_000 if q else 5_000_000, + n_stores=1_000 if q else 100_000, + n_weeks=100 if q else 1_000, + seed=seed, + ) + elif scenario == "survey_absorb": + df = make_survey_absorb( + n_rows=5_000 if q else 500_000, + n_states=50, + n_months=12 if q else 60, + n_replicates=8 if q else 80, + seed=seed, + ) + elif scenario == "tail_stress": + df = make_orders( + n_rows=50_000 if q else 5_000_000, + n_stores=1_000 if q else 100_000, + n_weeks=100 if q else 1_000, + seed=seed, + lifetime_span=0.2, + ) + elif scenario == "guard_small": + df = make_panel(n_units=1_000, n_periods=20, seed=seed, drop_frac=0.0) + else: + raise ValueError(f"unknown scenario {scenario!r}") + meta = {"n_obs": int(len(df)), "checksum": checksum(df)} + return df, meta + + +SCENARIO_IDS = ( + "county_policy", + "firm_churn", + "scanner_twfe", + "geo_experiment", + "survey_absorb", + "tail_stress", + "guard_small", +) diff --git a/benchmarks/speed_review/gen_findings_tables.py b/benchmarks/speed_review/gen_findings_tables.py index e1258094f..cb3455407 100644 --- a/benchmarks/speed_review/gen_findings_tables.py +++ b/benchmarks/speed_review/gen_findings_tables.py @@ -43,12 +43,12 @@ SINGLE_SCALE = ("reversible_dcdh", "dose_response") SCENARIO_DISPLAY = { - "campaign_staggered": "1. Staggered campaign", + "campaign_staggered": "1. Staggered campaign", "brand_awareness_survey": "2. Brand awareness survey", - "brfss_panel": "3. BRFSS microdata -> CS panel", - "geo_few_markets": "4. SDiD few markets", - "reversible_dcdh": "5. Reversible dCDH", - "dose_response": "6. Pricing dose-response", + "brfss_panel": "3. BRFSS microdata -> CS panel", + "geo_few_markets": "4. SDiD few markets", + "reversible_dcdh": "5. Reversible dCDH", + "dose_response": "6. Pricing dose-response", } @@ -84,15 +84,12 @@ def render_scale_sweep_totals(): py_t = py["total_seconds"] if py else None rs_t = rs["total_seconds"] if rs else None ratio = ( - f"{py_t/rs_t:.1f}x" - if (py_t is not None and rs_t is not None and rs_t > 0) - else "-" + f"{py_t/rs_t:.1f}x" if (py_t is not None and rs_t is not None and rs_t > 0) else "-" ) name_col = display if first else "" first = False rows.append( - f"| {name_col} | {scale} | " - f"{fmt_secs(py_t)} | {fmt_secs(rs_t)} | {ratio} |" + f"| {name_col} | {scale} | " f"{fmt_secs(py_t)} | {fmt_secs(rs_t)} | {ratio} |" ) for scen in SINGLE_SCALE: display = SCENARIO_DISPLAY[scen] @@ -100,15 +97,8 @@ def render_scale_sweep_totals(): rs = load(scen, None, "rust") py_t = py["total_seconds"] if py else None rs_t = rs["total_seconds"] if rs else None - ratio = ( - f"{py_t/rs_t:.1f}x" - if (py_t is not None and rs_t is not None and rs_t > 0) - else "-" - ) - rows.append( - f"| {display} | single | " - f"{fmt_secs(py_t)} | {fmt_secs(rs_t)} | {ratio} |" - ) + ratio = f"{py_t/rs_t:.1f}x" if (py_t is not None and rs_t is not None and rs_t > 0) else "-" + rows.append(f"| {display} | single | " f"{fmt_secs(py_t)} | {fmt_secs(rs_t)} | {ratio} |") return "\n".join(rows) @@ -160,8 +150,7 @@ def render_top_phases_by_scenario(): dropping the row entirely. """ rows = [ - "| Scenario | Scale | Backend | Top phase (%) " - "| 2nd phase (%) | 3rd phase (%) |", + "| Scenario | Scale | Backend | Top phase (%) " "| 2nd phase (%) | 3rd phase (%) |", "|---|---|---|---|---|---|", ] @@ -196,10 +185,7 @@ def largest_available(scen, backend): top = phase_rank(rec) if not top: continue - rows.append( - f"| {display} | {scale} | {backend} | " - f"{top[0]} | {top[1]} | {top[2]} |" - ) + rows.append(f"| {display} | {scale} | {backend} | " f"{top[0]} | {top[1]} | {top[2]} |") for scen in SINGLE_SCALE: display = SCENARIO_DISPLAY[scen] for backend in ("python", "rust"): @@ -207,10 +193,61 @@ def largest_available(scen, backend): top = phase_rank(rec) if not top: continue - rows.append( - f"| {display} | single | {backend} | " - f"{top[0]} | {top[1]} | {top[2]} |" - ) + rows.append(f"| {display} | single | {backend} | " f"{top[0]} | {top[1]} | {top[2]} |") + return "\n".join(rows) + + +FE_SCENARIO_DISPLAY = { + "county_policy": "7. County policy event study (SunAbraham)", + "firm_churn": "8. Firm panel with churn (SunAbraham)", + "scanner_twfe": "9. Scanner store-week (TWFE)", + "geo_experiment": "10. Geo experiment 5M orders (DiD absorb)", + "survey_absorb": "11. Survey BRR replicates (DiD absorb)", + "tail_stress": "12. Correlated-FE stress (DiD absorb)", + "guard_small": "13. Small-panel guard (TWFE)", +} + + +def render_fe_absorption(): + """Before/after/yardstick wall-clock for the FE-absorption suite. + + Renders whichever of fe_absorption_before.json / fe_absorption_after.json + / fe_absorption_pyfixest.json exist; absent files render as "-" columns so + the table is valid from PR-A (before only) onward. + """ + + def load_suite(name): + path = BASELINES / f"{name}.json" + if not path.exists(): + return {} + return { + r["scenario"]: r for r in json.loads(path.read_text())["results"] if "error" not in r + } + + before = load_suite("fe_absorption_before") + after = load_suite("fe_absorption_after") + yard = load_suite("fe_absorption_pyfixest") + + def cell(rec): + if not rec or rec.get("fit_median_s") is None: + return "-" + flag = " (noisy)" if rec.get("noisy") else "" + return f"{rec['fit_median_s']:.3f} (cv {rec['fit_cv']:.1%}){flag}" + + rows = [ + "| Scenario | n rows | Before (s) | After (s) | Speedup | pyfixest (s) |", + "|---|---:|---:|---:|---:|---:|", + ] + for scen, display in FE_SCENARIO_DISPLAY.items(): + b, a, y = before.get(scen), after.get(scen), yard.get(scen) + n_obs = (b or a or y or {}).get("n_obs") + n_col = f"{n_obs:,}" if n_obs else "-" + speedup = ( + f"{b['fit_median_s'] / a['fit_median_s']:.1f}x" + if (b and a and a.get("fit_median_s")) + else "-" + ) + rows.append(f"| {display} | {n_col} | {cell(b)} | {cell(a)} | {speedup} | {cell(y)} |") return "\n".join(rows) @@ -218,6 +255,7 @@ def largest_available(scen, backend): "scale_sweep_totals": render_scale_sweep_totals, "memory_by_scenario": render_memory_by_scenario, "top_phases_by_scenario": render_top_phases_by_scenario, + "fe_absorption": render_fe_absorption, } @@ -240,9 +278,7 @@ def update_markdown(path): f" to the document first." ) if n > 1: - raise RuntimeError( - f"Multiple marker pairs for '{table_id}' in {path}." - ) + raise RuntimeError(f"Multiple marker pairs for '{table_id}' in {path}.") text = new_text path.write_text(text) diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 3e91b30f2..5697e2a9a 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -4,6 +4,70 @@ This document outlines the strategy for improving diff-diff's performance on lar --- +## FE-absorption baseline: the MAP demeaning hot path (v3.6.x, July 2026) + +A measured head-to-head against pyfixest 0.60 (Rust demeaner core), prompted by +Instacart's high-cardinality marketplace-modeling writeup, established that the +method-of-alternating-projections demeaning engine (`demean_by_groups`, +`diff_diff/utils.py`) is the dominant cost of every TWFE-family fit - and that +the gap is machinery, not architecture: + +- **The engine is mainline, not edge.** `within_transform` is a thin two-way + wrapper over `demean_by_groups`, so every fit of TwoWayFixedEffects, + SunAbraham, BaconDecomposition, and WooldridgeDiD runs the MAP loop; + DiD/MultiPeriodDiD hit it via `absorb=`; the survey replicate path re-runs + the full weighted demean once per replicate (~80-200x per fit). +- **Attribution.** cProfile on SunAbraham's canonical staggered entry/exit + panel (30k units x 30 periods, 540k rows): 13.5s fit with **82% inside + `demean_by_groups`** - 2,948 `groupby().transform("mean")` calls, of which + 3.9s is pandas rebuilding the group index (`result_index_and_ids`) on every + call. The demeaning share of fit time across the estimator sweep was 64-87%. +- **Machinery gap.** On identical 5M-row arrays at tol=1e-8, the pandas MAP + loop is ~15x slower than pyfixest's compiled demeaner (2.39s vs 0.15s + two-way; 3.01s vs 0.20s three-way). Estimator-level gaps are 2-7x on + realistic workloads (shared fit costs dilute the demeaner share). Peak RSS + is at parity or better - memory is a non-issue. +- **Iteration regimes** (two-way, tol=1e-8, measured): balanced panel 2 + iterations; 10% random gaps 6; entry/exit lifetimes at 60% span 16-17; + order-level random incidence 8-9; contiguous 20% lifetimes 238-279 - the + last crosses the `max_iter=100` cap, so the fit warns and returns + slightly-off residuals where fixest/reghdfe/pyfixest (cap 10,000) converge. +- **Measured fix candidate.** A factorize-once + `np.bincount` rewrite of the + MAP loop (pure numpy, identical convergence semantics) measured **3.4x** on + the 5M-row two-way case and reaches full convergence on the stress case in + 19.2s where the current loop burns 25.1s hitting the cap without + converging. The same pattern already exists in-tree + (`diff_diff/survey.py` `_PsuScaffolding`). + +The measurement surface is the seven-scenario FE-absorption suite +(`benchmarks/speed_review/bench_fe_absorption.py`; shapes documented in +`docs/performance-scenarios.md` scenarios 7-13; committed baseline +`baselines/fe_absorption_before.json`, with `_after.json` to be added by the +optimization PR). The optional pyfixest yardstick +(`bench_fe_absorption_pyfixest.py`, guarded on import) asserts coefficient +parity < 1e-6 on the four exact-estimand scenarios so the comparison stays +honest over time. + +### FE-absorption suite results + + +| Scenario | n rows | Before (s) | After (s) | Speedup | pyfixest (s) | +|---|---:|---:|---:|---:|---:| +| 7. County policy event study (SunAbraham) | 177,289 | 3.865 (cv 2.5%) | - | - | 0.190 (cv 4.0%) | +| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | - | - | 1.912 (cv 20.3%) | +| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | - | - | 0.644 (cv 12.3%) | +| 10. Geo experiment 5M orders (DiD absorb) | 5,000,000 | 2.630 (cv 0.6%) | - | - | 0.623 (cv 7.8%) | +| 11. Survey BRR replicates (DiD absorb) | 500,000 | 7.385 (cv 8.6%) | - | - | - | +| 12. Correlated-FE stress (DiD absorb) | 5,000,000 | 26.271 (cv 0.4%) | - | - | 3.075 (cv 1.5%) | +| 13. Small-panel guard (TWFE) | 20,000 | 0.005 (cv 4.1%) | - | - | 0.007 (cv 2.2%) | + + +`tail_stress` is a deliberately adversarial correlated-FE shape and is +reported separately from the headline scenarios; the BEFORE baseline records +its non-convergence warning as evidence of the `max_iter=100` contract issue. + +--- + ## Memory scaling at the millions-of-units tail (v3.x, June 2026) At the scale where the dense working arrays - not compute - are the binding constraint diff --git a/docs/performance-scenarios.md b/docs/performance-scenarios.md index b0a42ea92..2dc3f2763 100644 --- a/docs/performance-scenarios.md +++ b/docs/performance-scenarios.md @@ -320,6 +320,65 @@ serves a different purpose: R-parity accuracy). They complement it. Callaway, Goodman-Bacon & Sant'Anna (2024), `docs/methodology/REGISTRY.md` ContinuousDiD section. +### FE-absorption suite (scenarios 7-13) - isolated-fit machinery benchmarks + +Unlike scenarios 1-6, which time full practitioner operation chains, scenarios +7-13 time **isolated `fit()` calls** on the fixed-effects absorption hot path +(`demean_by_groups` / `within_transform`, the method-of-alternating-projections +demeaning engine). They exist because the 2026-07 pyfixest gap analysis +measured 64-87% of fit time inside that engine for the TWFE-family estimators, +so a chain benchmark would only dilute the signal. Scripts: +`benchmarks/speed_review/bench_fe_absorption.py` (diff-diff lane, committed +`fe_absorption_before/after.json` baselines) and +`bench_fe_absorption_pyfixest.py` (optional external yardstick, guarded on +`import pyfixest`). Data shapes live in +`benchmarks/speed_review/fe_absorption_datagen.py`; every scenario is +seeded-deterministic and records an outcome checksum plus ATT/SE identity +fields so optimization PRs can prove estimates did not move. + +- **7. `county_policy`** - staggered state/county policy event study. + 3,109 counties (the US county count) x 60 months, four adoption cohorts + plus never-treated, ~5% attrition; `SunAbraham()` fit. The everyday + applied-micro shape (minimum-wage / opioid-policy literatures). Anchor: + Sun & Abraham (2021) applications; REGISTRY SunAbraham section. +- **8. `firm_churn`** - corporate panel with entry/exit. 100k firms x 40 + quarters, contiguous lifetimes covering ~60% of the window (Compustat/LBD + convention that firms enter and exit the sample); `SunAbraham()` fit. + Correlated unit-time incidence is the slow-convergence regime for MAP + (measured 16-17 iterations vs 2 balanced). +- **9. `scanner_twfe`** - retail scanner panel. 35k stores x 156 weeks with + entry/exit, binary post treatment; `TwoWayFixedEffects()` fit (~3.3M rows). + Nielsen-scanner shape; TWFE is the workhorse estimator. +- **10. `geo_experiment`** - order-level static-geo experiment. 5M orders + across 100k stores + 1k weeks, store-level treatment on for the post + window; `DifferenceInDifferences(absorb=["store", "week"])`. The + high-cardinality marketplace shape from the Instacart pyfixest writeup + (tech.instacart.com, 2025) - the workload PR #586's multi-absorb MAP + exists for. +- **11. `survey_absorb`** - state-policy microdata with replicate weights. + 500k rows, state + month FE, sampling weights plus 80 BRR replicate-weight + columns (BRFSS public-use convention); + `DifferenceInDifferences(absorb=..., survey_design=...)`. Times the + replicate-refit multiplier path: each replicate re-runs the full weighted + demean (`estimators.py` `_refit_did_absorb`). +- **12. `tail_stress`** - correlated-FE stress case, NOT a headline number. + 5M order-level rows where stores live for a contiguous ~20% of the weeks; + ~250+ MAP iterations. Characterizes the iteration tail and the `max_iter` + cap (the committed BEFORE baseline records the non-convergence warning). + Reported separately from scenarios 7-11 by design: it is a deliberately + adversarial shape, kept to make the suite honest about the worst case + rather than to flatter average-case wins. +- **13. `guard_small`** - regression guard. Balanced 1k x 20 panel through + `TwoWayFixedEffects()`; exists so optimization PRs prove the small-data + path did not regress (factorization/setup overhead would show here first). + +Noise protocol: every (scenario, repeat) runs in a fresh subprocess, strictly +sequential; sub-second scenarios take the median of several in-process fits +per subprocess; the driver pools all timed fits and reports median / min / +max / CV, flagging CV > 10% as unusable. Identity tolerances (ATT +atol=1e-9, SE rtol=1e-7, survey replicate SE rtol=1e-6) are recorded in the +baseline JSON itself. + ## Backend and environment notes All scenarios run under both backends where available: From ebe0fcad1aeb907ef24e635f95c5fc32b552cf2a Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Jul 2026 15:57:49 -0400 Subject: [PATCH 2/2] bench: address AI review - flag noisy/proxy cells in findings table, add repeats guard - gen_findings_tables.py cell() now derives the noisy flag from fit_cv > 10% for records without a precomputed `noisy` field (pyfixest lane), and marks timing-only yardstick cells (coef=None) as `proxy`; rendered footnote explains both qualifiers. Regenerated the fe_absorption table: firm_churn pyfixest now reads (cv 20.3%, noisy, proxy), scanner_twfe (cv 12.3%, noisy). - bench_fe_absorption_pyfixest.py: add the same --repeats >= 1 argparse guard as the main lane (verified: --repeats 0 now exits 2 with a clear error). Co-Authored-By: Claude Fable 5 --- .../bench_fe_absorption_pyfixest.py | 3 +++ .../speed_review/gen_findings_tables.py | 22 +++++++++++++++++-- docs/performance-plan.md | 8 ++++--- 3 files changed, 28 insertions(+), 5 deletions(-) diff --git a/benchmarks/speed_review/bench_fe_absorption_pyfixest.py b/benchmarks/speed_review/bench_fe_absorption_pyfixest.py index 4210b3630..fcd673bb4 100644 --- a/benchmarks/speed_review/bench_fe_absorption_pyfixest.py +++ b/benchmarks/speed_review/bench_fe_absorption_pyfixest.py @@ -89,6 +89,9 @@ def main(): ) args = ap.parse_args() + if args.repeats < 1: + ap.error("--repeats must be >= 1") + baseline = {} bp = Path(args.baseline) if bp.exists(): diff --git a/benchmarks/speed_review/gen_findings_tables.py b/benchmarks/speed_review/gen_findings_tables.py index cb3455407..2643b2fa6 100644 --- a/benchmarks/speed_review/gen_findings_tables.py +++ b/benchmarks/speed_review/gen_findings_tables.py @@ -231,8 +231,18 @@ def load_suite(name): def cell(rec): if not rec or rec.get("fit_median_s") is None: return "-" - flag = " (noisy)" if rec.get("noisy") else "" - return f"{rec['fit_median_s']:.3f} (cv {rec['fit_cv']:.1%}){flag}" + quals = [] + # The 10% threshold mirrors CV_FLAG in bench_fe_absorption.py. The + # diff-diff lane pre-computes `noisy`; the pyfixest lane records only + # `fit_cv`, so derive the flag here for both. + if rec.get("noisy") or (rec.get("fit_cv") or 0) > 0.10: + quals.append("noisy") + # Yardstick records with no coefficient are timing-only proxies + # (saturated i(rel_time) stand-in for Sun-Abraham; see the footnote). + if "coef" in rec and rec.get("coef") is None: + quals.append("proxy") + qual_txt = f", {', '.join(quals)}" if quals else "" + return f"{rec['fit_median_s']:.3f} (cv {rec['fit_cv']:.1%}{qual_txt})" rows = [ "| Scenario | n rows | Before (s) | After (s) | Speedup | pyfixest (s) |", @@ -248,6 +258,14 @@ def cell(rec): else "-" ) rows.append(f"| {display} | {n_col} | {cell(b)} | {cell(a)} | {speedup} | {cell(y)} |") + rows.append("") + rows.append( + "*noisy* = CV above the 10% unusable threshold from the noise protocol. " + "*proxy* = timing-only pyfixest stand-in: the Sun-Abraham scenarios run a " + "saturated `i(rel_time)` event study there (pyfixest 0.60 has no `sunab()`) - " + "comparable demeaning load, different estimand, so those cells are not " + "exact-estimand comparisons (see `bench_fe_absorption_pyfixest.py`)." + ) return "\n".join(rows) diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 5697e2a9a..e3aaf3c14 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -53,13 +53,15 @@ honest over time. | Scenario | n rows | Before (s) | After (s) | Speedup | pyfixest (s) | |---|---:|---:|---:|---:|---:| -| 7. County policy event study (SunAbraham) | 177,289 | 3.865 (cv 2.5%) | - | - | 0.190 (cv 4.0%) | -| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | - | - | 1.912 (cv 20.3%) | -| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | - | - | 0.644 (cv 12.3%) | +| 7. County policy event study (SunAbraham) | 177,289 | 3.865 (cv 2.5%) | - | - | 0.190 (cv 4.0%, proxy) | +| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | - | - | 1.912 (cv 20.3%, noisy, proxy) | +| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | - | - | 0.644 (cv 12.3%, noisy) | | 10. Geo experiment 5M orders (DiD absorb) | 5,000,000 | 2.630 (cv 0.6%) | - | - | 0.623 (cv 7.8%) | | 11. Survey BRR replicates (DiD absorb) | 500,000 | 7.385 (cv 8.6%) | - | - | - | | 12. Correlated-FE stress (DiD absorb) | 5,000,000 | 26.271 (cv 0.4%) | - | - | 3.075 (cv 1.5%) | | 13. Small-panel guard (TWFE) | 20,000 | 0.005 (cv 4.1%) | - | - | 0.007 (cv 2.2%) | + +*noisy* = CV above the 10% unusable threshold from the noise protocol. *proxy* = timing-only pyfixest stand-in: the Sun-Abraham scenarios run a saturated `i(rel_time)` event study there (pyfixest 0.60 has no `sunab()`) - comparable demeaning load, different estimand, so those cells are not exact-estimand comparisons (see `bench_fe_absorption_pyfixest.py`). `tail_stress` is a deliberately adversarial correlated-FE shape and is