From 71bde7742060b2be321ed29a13ddd368df12eb29 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Jul 2026 22:39:19 -0400 Subject: [PATCH] perf(rust): demean_map kernel - compiled MAP sweeps, rayon across columns, GIL released Optional Rust acceleration for FE-absorption demeaning, mirroring the canonical numpy engine exactly (python-canonical policy): - rust/src/demean.rs: per-column independent convergence loops, same sweep order, row-order scatter-add accumulation (np.bincount parity), division by per-group sums, zero-total-weight inert guard, max|x - x_old| < tol with NaN-poisoning semantics. Owned copies first, then Python::detach releases the GIL around the rayon compute; Fortran-order result so per-column views are contiguous. Codes validated against n_groups while copying. Cargo unit tests included. - diff_diff/utils.py: numpy MAP loops extracted into _demean_map_numpy (the reference engine, directly testable) + _demean_map_rust wrapper (explicit dtype coercion, None contract honored when the kernel symbol is absent, deliberate "demean_map" marker errors -> numpy fallback, others re-raised). demean_by_groups dispatches rust-then-numpy; warning contract, snap/LSMR stage, one-way delegation, n_effects unchanged. - diff_diff/_backend.py: demean_map imported INDEPENDENTLY so a stale or mixed-version extension missing only this symbol degrades to the numpy engine without disabling older Rust accelerations. Equivalence (tests/test_rust_backend.py::TestDemeanMapKernel, 16 cases): iteration-count EQUALITY then assert_allclose atol=1e-12 across weighted/unweighted x 2-way/3-way x balanced/unbalanced/contiguous(>100 iters), zero-weight inertness, NaN non-convergence parity, k=1/64, non-convergence flags, forced-fallback + stale-symbol fallback, warning parity, estimator-level ATT/SE + FE-spanned snap-decision parity. Outputs bit-identical on fixtures (asserted at allclose per policy). Sweeps: 960 passed (rust) + 875 passed (DIFF_DIFF_BACKEND=python). Measured on frozen code vs a fresh same-session numpy run (CVs <= 4.1%): firm_churn 48.98s -> 25.84s (1.90x; CV-adjusted lower bound 1.78x, submit gate >= 1.5x PASS), tail_stress 2.79x, geo 2.02x (now ahead of the pyfixest yardstick), scanner 1.59x, survey 1.44x, county 1.36x, guard_small unregressed. Identity vs committed after-baselines 1e-13-1e-16 incl. tail_stress. Memory trade-off disclosed: firm_churn peak RSS 13.4 -> 21.0 GB (kernel input copy + result; variable-chunking is the noted follow-up). Solver-bound county-class finding promoted to an Actionable TODO row. REGISTRY/doc-deps/CHANGELOG updated. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 13 + TODO.md | 1 + .../baselines/fe_absorption_rust.json | 151 +++ .../speed_review/gen_findings_tables.py | 18 +- diff_diff/_backend.py | 12 + diff_diff/utils.py | 165 ++- docs/doc-deps.yaml | 2 +- docs/methodology/REGISTRY.md | 12 + docs/performance-plan.md | 36 +- rust/src/demean.rs | 333 ++++++ rust/src/lib.rs | 9 +- tests/test_rust_backend.py | 998 +++++++++++++----- 12 files changed, 1415 insertions(+), 335 deletions(-) create mode 100644 benchmarks/speed_review/baselines/fe_absorption_rust.json create mode 100644 rust/src/demean.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c3e05b0e..0751f4aa1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **Rust `demean_map` kernel for FE absorption** (optional backend acceleration). When the Rust + extension is available, the method-of-alternating-projections sweeps in + `demean_by_groups`/`within_transform` run in a compiled kernel, rayon-parallel across the + demeaned variables, with the GIL released during compute. The kernel mirrors the canonical + numpy engine exactly (same sweep order, row-order bincount accumulation, zero-total-weight + inert guard, and `max|x - x_old| < tol` stopping rule with NaN-poisoning semantics); numpy + remains the reference implementation per the python-canonical policy, with equivalence tests + asserting iteration-count equality plus `assert_allclose` at atol=1e-12. No behavior change + when the extension is absent or `DIFF_DIFF_BACKEND=python`; kernel-side validation errors fall + back to the numpy engine. Affects every MAP consumer (TWFE, SunAbraham, Bacon, Wooldridge, + DiD/MPD `absorb=`, and the survey replicate refits). + ### Changed - **FE-absorption demeaning rewritten: factorize-once + `np.bincount` method of alternating projections** (`demean_by_groups` / `within_transform`). Each absorbed dimension is factorized diff --git a/TODO.md b/TODO.md index fb2d4130e..973f2cb92 100644 --- a/TODO.md +++ b/TODO.md @@ -47,6 +47,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| SunAbraham-family fits are now SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel. The parked "QR+SVD redundancy is minimal vs the SVD" note predates these measurements - the pivoted QR alone is now ~20% of the whole fit. Candidates: reuse the rank-detection QR factors for the solve, or a Cholesky fast path for well-conditioned tall-skinny designs. Measured attribution in docs/performance-plan.md (FE-absorption baseline). | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Heavy | Medium | | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium | | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | diff --git a/benchmarks/speed_review/baselines/fe_absorption_rust.json b/benchmarks/speed_review/baselines/fe_absorption_rust.json new file mode 100644 index 000000000..58e4518e4 --- /dev/null +++ b/benchmarks/speed_review/baselines/fe_absorption_rust.json @@ -0,0 +1,151 @@ +{ + "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": "rust", + "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.2803785413990308, + "se": 0.03124486461021617, + "fit_median_s": 1.7567, + "fit_min_s": 1.7285, + "fit_max_s": 1.8617, + "fit_cv": 0.0302, + "noisy": false, + "n_timed_fits": 9, + "datagen_s": 0.005, + "peak_rss_mb": 4237.1, + "backend_resolved": "rust", + "warnings": [] + }, + { + "scenario": "firm_churn", + "n_obs": 2400000, + "checksum": 240997.06223418942, + "att": 0.31443271844971477, + "se": 0.005837281701253889, + "fit_median_s": 25.8387, + "fit_min_s": 24.6121, + "fit_max_s": 26.7088, + "fit_cv": 0.041, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.069, + "peak_rss_mb": 20963.6, + "backend_resolved": "rust", + "warnings": [] + }, + { + "scenario": "scanner_twfe", + "n_obs": 3255000, + "checksum": 76832.22571739106, + "att": 0.25116413149501093, + "se": 0.0024149387421606118, + "fit_median_s": 0.6239, + "fit_min_s": 0.6215, + "fit_max_s": 0.6252, + "fit_cv": 0.003, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.076, + "peak_rss_mb": 1477.0, + "backend_resolved": "rust", + "warnings": [] + }, + { + "scenario": "geo_experiment", + "n_obs": 5000000, + "checksum": 443696.3834000407, + "att": 0.25194372599916653, + "se": 0.0017886314723251327, + "fit_median_s": 0.4839, + "fit_min_s": 0.483, + "fit_max_s": 0.4845, + "fit_cv": 0.0016, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.062, + "peak_rss_mb": 1744.8, + "backend_resolved": "rust", + "warnings": [ + "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", + "Regressor(s) ['treated', 'post'] are collinear with the absorbed fixed effects (absorb=['store', 'week']): the" + ] + }, + { + "scenario": "survey_absorb", + "n_obs": 500000, + "checksum": -176482.6014001448, + "att": 0.24554568144130481, + "se": 0.007662698712656388, + "fit_median_s": 2.8354, + "fit_min_s": 2.8354, + "fit_max_s": 2.8357, + "fit_cv": 0.0001, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.106, + "peak_rss_mb": 3587.3, + "backend_resolved": "rust", + "warnings": [ + "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", + "Regressor(s) ['treated', 'post'] are collinear with the absorbed fixed effects (absorb=['state', 'month']): th" + ] + }, + { + "scenario": "tail_stress", + "n_obs": 5000000, + "checksum": 376703.7248141762, + "att": 0.2502082560083505, + "se": 0.004392830594147583, + "fit_median_s": 11.365, + "fit_min_s": 11.3007, + "fit_max_s": 11.4856, + "fit_cv": 0.0082, + "noisy": false, + "n_timed_fits": 3, + "datagen_s": 0.066, + "peak_rss_mb": 1841.8, + "backend_resolved": "rust", + "warnings": [ + "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", + "Regressor(s) ['treated', 'post'] are collinear with the absorbed fixed effects (absorb=['store', 'week']): the" + ] + }, + { + "scenario": "guard_small", + "n_obs": 20000, + "checksum": -571.6190664591859, + "att": 0.2515087011451809, + "se": 0.028726073333820955, + "fit_median_s": 0.0034, + "fit_min_s": 0.0032, + "fit_max_s": 0.0036, + "fit_cv": 0.0346, + "noisy": false, + "n_timed_fits": 21, + "datagen_s": 0.001, + "peak_rss_mb": 148.0, + "backend_resolved": "rust", + "warnings": [] + } + ] +} \ No newline at end of file diff --git a/benchmarks/speed_review/gen_findings_tables.py b/benchmarks/speed_review/gen_findings_tables.py index 2643b2fa6..da0554368 100644 --- a/benchmarks/speed_review/gen_findings_tables.py +++ b/benchmarks/speed_review/gen_findings_tables.py @@ -244,12 +244,13 @@ def cell(rec): qual_txt = f", {', '.join(quals)}" if quals else "" return f"{rec['fit_median_s']:.3f} (cv {rec['fit_cv']:.1%}{qual_txt})" + rust = load_suite("fe_absorption_rust") rows = [ - "| Scenario | n rows | Before (s) | After (s) | Speedup | pyfixest (s) |", - "|---|---:|---:|---:|---:|---:|", + "| Scenario | n rows | Before (s) | After (s) | Speedup | Rust (s) | Rust speedup | pyfixest (s) |", + "|---|---:|---:|---:|---:|---:|---:|---:|", ] for scen, display in FE_SCENARIO_DISPLAY.items(): - b, a, y = before.get(scen), after.get(scen), yard.get(scen) + b, a, y, r = before.get(scen), after.get(scen), yard.get(scen), rust.get(scen) n_obs = (b or a or y or {}).get("n_obs") n_col = f"{n_obs:,}" if n_obs else "-" speedup = ( @@ -257,10 +258,19 @@ def cell(rec): if (b and a and a.get("fit_median_s")) else "-" ) - rows.append(f"| {display} | {n_col} | {cell(b)} | {cell(a)} | {speedup} | {cell(y)} |") + rust_speedup = ( + f"{a['fit_median_s'] / r['fit_median_s']:.1f}x" + if (a and r and r.get("fit_median_s")) + else "-" + ) + rows.append( + f"| {display} | {n_col} | {cell(b)} | {cell(a)} | {speedup} | " + f"{cell(r)} | {rust_speedup} | {cell(y)} |" + ) rows.append("") rows.append( "*noisy* = CV above the 10% unusable threshold from the noise protocol. " + "*Rust speedup* is vs the After column (numpy engine, same code). " "*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 " diff --git a/diff_diff/_backend.py b/diff_diff/_backend.py index f763f87c2..2c15de5f2 100644 --- a/diff_diff/_backend.py +++ b/diff_diff/_backend.py @@ -65,6 +65,14 @@ _rust_sc_weight_fw_weighted_with_convergence = None _rust_backend_info = None +# FE-absorption MAP demeaning kernel: imported independently so a stale or +# mixed-version extension missing only this newer symbol degrades to the +# numpy demeaning engine WITHOUT disabling the older Rust accelerations. +try: + from diff_diff._rust_backend import demean_map as _rust_demean_map +except ImportError: + _rust_demean_map = None + # Determine final backend based on environment variable and availability if _backend_env == "python": # Force pure Python mode - disable Rust even if available @@ -73,6 +81,8 @@ _rust_project_simplex = None _rust_solve_ols = None _rust_compute_robust_vcov = None + # FE-absorption MAP demeaning kernel + _rust_demean_map = None # TROP estimator acceleration (local method) _rust_unit_distance_matrix = None _rust_loocv_grid_search = None @@ -124,6 +134,8 @@ def rust_backend_info(): "_rust_project_simplex", "_rust_solve_ols", "_rust_compute_robust_vcov", + # FE-absorption MAP demeaning kernel + "_rust_demean_map", # TROP estimator acceleration (local method) "_rust_unit_distance_matrix", "_rust_loocv_grid_search", diff --git a/diff_diff/utils.py b/diff_diff/utils.py index e4a44c201..1e13a7bfb 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -13,6 +13,7 @@ # Import Rust backend if available (from _backend to avoid circular imports) from diff_diff._backend import ( HAS_RUST_BACKEND, + _rust_demean_map, _rust_project_simplex, _rust_sdid_unit_weights, _rust_compute_time_weights, @@ -2641,6 +2642,111 @@ def demean_by_group( return data, n_effects +def _demean_map_numpy( + x_cols: List[np.ndarray], + codes_list: List[np.ndarray], + n_groups_list: List[int], + weights: Optional[np.ndarray], + tol: float, + max_iter: int, +) -> Tuple[List[np.ndarray], List[int]]: + """Canonical numpy MAP engine over pre-factorized group codes. + + The reference implementation for the demeaning contract (python-canonical + policy): the Rust kernel mirrors this exactly (sweep order, row-order + bincount accumulation, ``max|x - x_old| < tol`` convergence per column) + and equivalence tests compare against THIS function directly. + + Returns (demeaned columns, iterations per column; -1 = not converged). + """ + demeaned: List[np.ndarray] = [] + iters: List[int] = [] + if weights is not None: + w = np.asarray(weights, dtype=np.float64) + # Cache per-group weight sums once (invariant across variables/iterations). + w_sums = [ + np.bincount(codes, weights=w, minlength=n_g) + for codes, n_g in zip(codes_list, n_groups_list) + ] + for x0 in x_cols: + x = np.asarray(x0, dtype=np.float64).copy() + it_out = -1 + for _iter in range(max_iter): + x_old = x.copy() + for codes, n_g, w_sum in zip(codes_list, n_groups_list, w_sums): + wx_sum = np.bincount(codes, weights=w * x, minlength=n_g) + # Guard zero-total-weight groups (survey subpopulation / + # zero-weight domain padding): leave such rows unchanged + # (mean 0) so they remain inert in the downstream WLS + # instead of poisoning the design with NaN/Inf. + means = np.divide(wx_sum, w_sum, out=np.zeros_like(wx_sum), where=w_sum > 0) + x = x - means[codes] + if np.max(np.abs(x - x_old)) < tol: + it_out = _iter + 1 + break + demeaned.append(x) + iters.append(it_out) + else: + counts = [ + np.bincount(codes, minlength=n_g).astype(np.float64) + for codes, n_g in zip(codes_list, n_groups_list) + ] + for x0 in x_cols: + x = np.asarray(x0, dtype=np.float64).copy() + it_out = -1 + for _iter in range(max_iter): + x_old = x.copy() + for codes, n_g, cnt in zip(codes_list, n_groups_list, counts): + means = np.bincount(codes, weights=x, minlength=n_g) / cnt + x = x - means[codes] + if np.max(np.abs(x - x_old)) < tol: + it_out = _iter + 1 + break + demeaned.append(x) + iters.append(it_out) + return demeaned, iters + + +def _demean_map_rust( + x_cols: List[np.ndarray], + codes_list: List[np.ndarray], + n_groups_list: List[int], + weights: Optional[np.ndarray], + tol: float, + max_iter: int, +) -> Optional[Tuple[List[np.ndarray], List[int]]]: + """Marshal to the Rust ``demean_map`` kernel. + + Returns None to signal "use the canonical numpy engine" (kernel absent, + degenerate shapes, or a deliberate kernel-side validation error). Dtypes + are coerced explicitly BEFORE the call - never rely on exception handling + for dtype mismatches. + """ + if _rust_demean_map is None: + return None + if not x_cols or x_cols[0].shape[0] == 0: + return None + x_mat = np.ascontiguousarray(np.column_stack(x_cols), dtype=np.float64) + codes_mat = np.ascontiguousarray(np.column_stack(codes_list), dtype=np.int64) + w = None if weights is None else np.ascontiguousarray(weights, dtype=np.float64) + try: + out, iters = _rust_demean_map( # type: ignore[misc] + x_mat, + codes_mat, + [int(g) for g in n_groups_list], + w, + float(tol), + int(max_iter), + ) + except ValueError as e: + if "demean_map" in str(e): + # deliberate kernel-side validation marker -> numpy fallback + return None + raise + # F-order result: per-column views are contiguous, no extra copy + return [out[:, j] for j in range(out.shape[1])], [int(i) for i in iters] + + def demean_by_groups( data: pd.DataFrame, variables: List[str], @@ -2774,55 +2880,18 @@ def demean_by_groups( codes_list.append(codes.astype(np.intp, copy=False)) n_groups_list.append(len(uniques)) n_effects = sum(n_g - 1 for n_g in n_groups_list) - demeaned_values: List[np.ndarray] = [] - non_converged_vars: List[str] = [] - - if weights is not None: - w = np.asarray(weights, dtype=np.float64) - # Cache per-group weight sums once (invariant across variables/iterations). - w_sums = [ - np.bincount(codes, weights=w, minlength=n_g) - for codes, n_g in zip(codes_list, n_groups_list) - ] - for var in variables: - x = data[var].values.astype(np.float64) - converged = False - for _iter in range(max_iter): - x_old = x.copy() - for codes, n_g, w_sum in zip(codes_list, n_groups_list, w_sums): - wx_sum = np.bincount(codes, weights=w * x, minlength=n_g) - # Guard zero-total-weight groups (survey subpopulation / - # zero-weight domain padding): leave such rows unchanged - # (mean 0) so they remain inert in the downstream WLS - # instead of poisoning the design with NaN/Inf. - means = np.divide(wx_sum, w_sum, out=np.zeros_like(wx_sum), where=w_sum > 0) - x = x - means[codes] - if np.max(np.abs(x - x_old)) < tol: - converged = True - break - if not converged: - non_converged_vars.append(var) - demeaned_values.append(x) - else: - counts = [ - np.bincount(codes, minlength=n_g).astype(np.float64) - for codes, n_g in zip(codes_list, n_groups_list) - ] - for var in variables: - x = data[var].values.astype(np.float64) - converged = False - for _iter in range(max_iter): - x_old = x.copy() - for codes, n_g, cnt in zip(codes_list, n_groups_list, counts): - means = np.bincount(codes, weights=x, minlength=n_g) / cnt - x = x - means[codes] - if np.max(np.abs(x - x_old)) < tol: - converged = True - break - if not converged: - non_converged_vars.append(var) - demeaned_values.append(x) + x_cols = [data[var].values for var in variables] + result = None + if HAS_RUST_BACKEND and _rust_demean_map is not None: + # Rust kernel: identical sweep order, accumulation order, and + # convergence criterion (rayon-parallel across columns). None means + # "use the canonical numpy engine". + result = _demean_map_rust(x_cols, codes_list, n_groups_list, weights, tol, max_iter) + if result is None: + result = _demean_map_numpy(x_cols, codes_list, n_groups_list, weights, tol, max_iter) + demeaned_values, _iters = result + non_converged_vars = [v for v, it in zip(variables, _iters) if it < 0] if non_converged_vars: warn_if_not_converged( diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index a773db435..176e74310 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -904,7 +904,7 @@ sources: - path: docs/methodology/REGISTRY.md section: "Absorbed Fixed Effects with Survey Weights" type: methodology - note: "demean_by_groups / within_transform: MAP convergence contract (max_iter/tol), bincount accumulation numerics, NaN-group-key guard, snap_absorbed_regressors FE-spanned snap + exact LSMR span confirmation" + note: "demean_by_groups / within_transform: MAP convergence contract (max_iter/tol), bincount accumulation numerics, NaN-group-key guard, snap_absorbed_regressors FE-spanned snap + exact LSMR span confirmation, rust demean_map dispatch" - path: docs/practitioner_getting_started.rst section: "Step 4: Check Whether the Result Is Trustworthy" type: user_guide diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 38e0c3558..33ab7fd32 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -4339,6 +4339,18 @@ unequal selection probabilities). (drift compounds across MAP iterations), not bit-for-bit; estimator estimates are validated unchanged at the FE-absorption benchmark identity gate (`benchmarks/speed_review/bench_fe_absorption.py --check-estimates`). +- **Note:** When the optional Rust backend is available, the MAP sweeps run in + the compiled `demean_map` kernel (rayon-parallel across the demeaned + variables). The kernel mirrors the canonical numpy engine + (`diff_diff.utils._demean_map_numpy`) exactly: same per-variable independent + convergence loops, same dimension sweep order, the same row-order + scatter-add accumulation as `np.bincount`, division by the per-group sums, + the same zero-total-weight inert-row guard, and the same + `max|x - x_old| < tol` stopping rule with NaN-poisoning semantics. Per the + python-canonical policy, numpy is the reference implementation; equivalence + tests assert iteration-count equality plus `assert_allclose` at atol=1e-12 + (never a bit-identity claim). `DIFF_DIFF_BACKEND=python` disables the + kernel; any kernel-side validation error falls back to the numpy engine. - **Edge case:** NaN in an absorbed group column raises a `ValueError` naming the column. `pd.factorize` codes NaN keys as -1, which would otherwise silently index the last group's mean; the prior pandas behavior was itself silently bad diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 864c30ddc..ec546a7b1 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -66,20 +66,34 @@ in the existing `rust/` backend is the candidate next step, to be scoped only if the remaining gap matters in practice (pyfixest itself validates that architecture). +**Landed (v3.6.x): Rust `demean_map` kernel** (optional backend; numpy +canonical). Rayon-parallel MAP sweeps across the demeaned variables with the +GIL released; equivalence locked by iteration-count equality + +assert_allclose(atol=1e-12) against `_demean_map_numpy`. Measured on frozen +code against a fresh same-session numpy-backend run (medians, CVs <= 4.1%): +firm_churn 48.98s -> 25.84s (1.90x; CV-adjusted lower bound 1.78x - the +submit gate), tail_stress 31.74s -> 11.37s (2.79x), geo_experiment 0.976s -> +0.484s (2.02x, now ahead of the pyfixest yardstick's 0.623s), scanner 1.59x, +survey 1.44x, county 1.36x, guard_small unregressed. Identity vs the +committed after-baselines: 1e-13-1e-16 on every scenario incl. tail_stress. +Known trade-off: the kernel holds one owned input copy plus the result +(firm_churn peak RSS 13.4 GB -> 21.0 GB); chunking variables through the +kernel is the noted follow-up if that bites a real workload. + ### 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%) | 2.388 (cv 2.6%) | 1.6x | 0.190 (cv 4.0%, proxy) | -| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | 49.502 (cv 1.9%) | 1.9x | 1.912 (cv 20.3%, noisy, proxy) | -| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | 0.983 (cv 0.1%) | 1.6x | 0.644 (cv 12.3%, noisy) | -| 10. Geo experiment 5M orders (DiD absorb) | 5,000,000 | 2.630 (cv 0.6%) | 0.973 (cv 0.4%) | 2.7x | 0.623 (cv 7.8%) | -| 11. Survey BRR replicates (DiD absorb) | 500,000 | 7.385 (cv 8.6%) | 4.077 (cv 0.1%) | 1.8x | - | -| 12. Correlated-FE stress (DiD absorb) | 5,000,000 | 26.271 (cv 0.4%) | 31.765 (cv 0.6%) | 0.8x | 3.075 (cv 1.5%) | -| 13. Small-panel guard (TWFE) | 20,000 | 0.005 (cv 4.1%) | 0.004 (cv 4.4%) | 1.3x | 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`). +| Scenario | n rows | Before (s) | After (s) | Speedup | Rust (s) | Rust speedup | pyfixest (s) | +|---|---:|---:|---:|---:|---:|---:|---:| +| 7. County policy event study (SunAbraham) | 177,289 | 3.865 (cv 2.5%) | 2.388 (cv 2.6%) | 1.6x | 1.757 (cv 3.0%) | 1.4x | 0.190 (cv 4.0%, proxy) | +| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | 49.502 (cv 1.9%) | 1.9x | 25.839 (cv 4.1%) | 1.9x | 1.912 (cv 20.3%, noisy, proxy) | +| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | 0.983 (cv 0.1%) | 1.6x | 0.624 (cv 0.3%) | 1.6x | 0.644 (cv 12.3%, noisy) | +| 10. Geo experiment 5M orders (DiD absorb) | 5,000,000 | 2.630 (cv 0.6%) | 0.973 (cv 0.4%) | 2.7x | 0.484 (cv 0.2%) | 2.0x | 0.623 (cv 7.8%) | +| 11. Survey BRR replicates (DiD absorb) | 500,000 | 7.385 (cv 8.6%) | 4.077 (cv 0.1%) | 1.8x | 2.835 (cv 0.0%) | 1.4x | - | +| 12. Correlated-FE stress (DiD absorb) | 5,000,000 | 26.271 (cv 0.4%) | 31.765 (cv 0.6%) | 0.8x | 11.365 (cv 0.8%) | 2.8x | 3.075 (cv 1.5%) | +| 13. Small-panel guard (TWFE) | 20,000 | 0.005 (cv 4.1%) | 0.004 (cv 4.4%) | 1.3x | 0.003 (cv 3.5%) | 1.0x | 0.007 (cv 2.2%) | + +*noisy* = CV above the 10% unusable threshold from the noise protocol. *Rust speedup* is vs the After column (numpy engine, same code). *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 diff --git a/rust/src/demean.rs b/rust/src/demean.rs new file mode 100644 index 000000000..034773b5d --- /dev/null +++ b/rust/src/demean.rs @@ -0,0 +1,333 @@ +//! Fixed-effects absorption: method-of-alternating-projections demeaning. +//! +//! Exact mirror of the canonical numpy engine `_demean_map_numpy` in +//! `diff_diff/utils.py` (python-canonical policy): same per-column +//! independent convergence loops, same dimension sweep order, the same +//! row-order scatter-add accumulation as `np.bincount`, division by the +//! per-group sums (never multiplication by a reciprocal), and the same +//! `max|x - x_old| < tol` stopping rule with NaN-poisoning semantics (a NaN +//! delta compares false, so a NaN-carrying column never converges - matching +//! numpy's `np.max(...) < tol`). Parallelism is rayon across columns, which +//! preserves the per-variable contract exactly. + +use ndarray::{Array1, Array2, ShapeBuilder}; +use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use rayon::prelude::*; + +/// Demean one column in place by alternating projections. +/// +/// Returns the 1-based iteration count at convergence, or -1 if the column +/// did not converge within `max_iter` (including the NaN case). +#[allow(clippy::too_many_arguments)] +fn demean_column( + x: &mut [f64], + prev: &mut [f64], + sums: &mut [Vec], + codes: &[Vec], + denoms: &[Vec], + weights: Option<&[f64]>, + tol: f64, + max_iter: usize, +) -> i64 { + let n = x.len(); + for iter in 0..max_iter { + prev.copy_from_slice(x); + for (d, codes_d) in codes.iter().enumerate() { + let sums_d = &mut sums[d]; + sums_d.iter_mut().for_each(|s| *s = 0.0); + // Row-order scatter-add: the same accumulation order as + // np.bincount(codes, weights=...). + match weights { + Some(w) => { + for i in 0..n { + sums_d[codes_d[i]] += w[i] * x[i]; + } + } + None => { + for i in 0..n { + sums_d[codes_d[i]] += x[i]; + } + } + } + // means = sums / denom; zero-total-weight groups stay at mean 0 + // (rows inert) - the np.divide(where=w_sum > 0) guard semantics. + let denom_d = &denoms[d]; + for g in 0..sums_d.len() { + if denom_d[g] > 0.0 { + sums_d[g] /= denom_d[g]; + } else { + sums_d[g] = 0.0; + } + } + for i in 0..n { + x[i] -= sums_d[codes_d[i]]; + } + } + // all(|delta| < tol) is exactly numpy's max(|delta|) < tol, + // INCLUDING NaN poisoning: a NaN delta compares false. + let converged = x.iter().zip(prev.iter()).all(|(a, b)| (a - b).abs() < tol); + if converged { + return (iter + 1) as i64; + } + } + -1 +} + +/// Method-of-alternating-projections demeaning over pre-factorized codes. +/// +/// # Arguments +/// * `x` - (n, k) float64 matrix; columns are the variables to demean +/// * `codes` - (n, d) int64 factorized group codes, all in [0, n_groups[j]) +/// * `n_groups` - group count per absorbed dimension +/// * `weights` - optional observation weights (weighted group means) +/// * `tol` - convergence tolerance on max|x - x_old| per full sweep cycle +/// * `max_iter` - iteration cap per column +/// +/// # Returns +/// Tuple of (demeaned (n, k) float64 array in Fortran order so per-column +/// views are contiguous, iterations-per-column int64 array; -1 marks a +/// column that did not converge). +#[pyfunction] +#[pyo3(signature = (x, codes, n_groups, weights=None, tol=1e-10, max_iter=10_000))] +#[allow(clippy::type_complexity)] +pub fn demean_map<'py>( + py: Python<'py>, + x: PyReadonlyArray2<'py, f64>, + codes: PyReadonlyArray2<'py, i64>, + n_groups: Vec, + weights: Option>, + tol: f64, + max_iter: usize, +) -> PyResult<(Bound<'py, PyArray2>, Bound<'py, PyArray1>)> { + let x_arr = x.as_array(); + let codes_arr = codes.as_array(); + let n = x_arr.nrows(); + let k = x_arr.ncols(); + let d = codes_arr.ncols(); + + if codes_arr.nrows() != n { + return Err(PyValueError::new_err( + "demean_map: x and codes must have the same number of rows", + )); + } + if n_groups.len() != d { + return Err(PyValueError::new_err( + "demean_map: n_groups length must match the number of code columns", + )); + } + if n == 0 || k == 0 || d == 0 { + return Err(PyValueError::new_err( + "demean_map: empty input (no rows, columns, or dimensions)", + )); + } + + // Owned copies FIRST: PyReadonlyArray borrows cannot cross py.detach + // (compiler-enforced by the Ungil bound). Codes are validated while + // copying - an out-of-range code would otherwise be a silent + // wrong-answer scatter. + let mut codes_owned: Vec> = Vec::with_capacity(d); + for j in 0..d { + let n_g = n_groups[j]; + let mut col = Vec::with_capacity(n); + for i in 0..n { + let c = codes_arr[(i, j)]; + if c < 0 || (c as usize) >= n_g { + return Err(PyValueError::new_err(format!( + "demean_map: code {} out of range [0, {}) in dimension {}", + c, n_g, j + ))); + } + col.push(c as usize); + } + codes_owned.push(col); + } + + let w_owned: Option> = match &weights { + Some(w) => { + let wv = w.as_array(); + if wv.len() != n { + return Err(PyValueError::new_err( + "demean_map: weights length must match the number of rows", + )); + } + Some(wv.to_vec()) + } + None => None, + }; + + // Per-dimension denominators, computed once (counts, or weight sums with + // the zero-total-weight guard applied at division time). + let denoms: Vec> = codes_owned + .iter() + .zip(n_groups.iter()) + .map(|(codes_d, &n_g)| { + let mut den = vec![0.0f64; n_g]; + match &w_owned { + Some(w) => { + for i in 0..n { + den[codes_d[i]] += w[i]; + } + } + None => { + for i in 0..n { + den[codes_d[i]] += 1.0; + } + } + } + den + }) + .collect(); + + // One owned working copy of the data, column-major so each rayon task + // owns a contiguous column (this buffer becomes the result - no second + // full-matrix transient). + let mut cols: Vec> = (0..k) + .map(|j| (0..n).map(|i| x_arr[(i, j)]).collect()) + .collect(); + + let w_slice = w_owned.as_deref(); + let codes_ref = &codes_owned; + let denoms_ref = &denoms; + let n_groups_ref = &n_groups; + + // Release the GIL for the compute: rayon across columns, each with its + // own convergence loop and scratch buffers. + let iters: Vec = py.detach(|| { + cols.par_iter_mut() + .with_min_len(1) + .map(|col| { + let mut prev = vec![0.0f64; n]; + let mut sums: Vec> = + n_groups_ref.iter().map(|&n_g| vec![0.0f64; n_g]).collect(); + demean_column( + col, &mut prev, &mut sums, codes_ref, denoms_ref, w_slice, tol, max_iter, + ) + }) + .collect() + }); + + // Flatten column-major into one buffer (progressively dropping source + // columns), then expose as a Fortran-order (n, k) array: per-column + // Python views are contiguous and no transpose copy is needed. + let mut flat: Vec = Vec::with_capacity(n * k); + for col in cols { + flat.extend_from_slice(&col); + } + let out = Array2::from_shape_vec((n, k).f(), flat) + .map_err(|e| PyValueError::new_err(format!("demean_map: shape error: {}", e)))?; + let iters_arr = Array1::from(iters); + Ok((out.into_pyarray(py), iters_arr.into_pyarray(py))) +} + +#[cfg(test)] +mod tests { + use super::demean_column; + + fn run( + x: &mut Vec, + codes: Vec>, + n_groups: Vec, + weights: Option>, + tol: f64, + max_iter: usize, + ) -> i64 { + let n = x.len(); + let denoms: Vec> = codes + .iter() + .zip(n_groups.iter()) + .map(|(codes_d, &n_g)| { + let mut den = vec![0.0f64; n_g]; + for i in 0..n { + den[codes_d[i]] += weights.as_ref().map_or(1.0, |w| w[i]); + } + den + }) + .collect(); + let mut prev = vec![0.0f64; n]; + let mut sums: Vec> = n_groups.iter().map(|&g| vec![0.0f64; g]).collect(); + demean_column( + x, + &mut prev, + &mut sums, + &codes, + &denoms, + weights.as_deref(), + tol, + max_iter, + ) + } + + #[test] + fn one_way_demean_exact() { + // groups {0: [1, 3], 1: [5, 7]} -> means 2 and 6 + let mut x = vec![1.0, 3.0, 5.0, 7.0]; + let it = run(&mut x, vec![vec![0, 0, 1, 1]], vec![2], None, 1e-12, 100); + assert!(it >= 1); + assert_eq!(x, vec![-1.0, 1.0, -1.0, 1.0]); + } + + #[test] + fn weighted_zero_weight_group_inert() { + // group 1 has zero total weight -> its rows keep their values + let mut x = vec![1.0, 3.0, 5.0, 7.0]; + let it = run( + &mut x, + vec![vec![0, 0, 1, 1]], + vec![2], + Some(vec![1.0, 1.0, 0.0, 0.0]), + 1e-12, + 100, + ); + assert!(it >= 1); + assert_eq!(x, vec![-1.0, 1.0, 5.0, 7.0]); + } + + #[test] + fn non_convergence_reports_minus_one() { + // two-way unbalanced with a starved budget and an impossible tol + let mut x = vec![1.0, 2.0, 3.0, 4.0, 5.0]; + let it = run( + &mut x, + vec![vec![0, 0, 1, 1, 2], vec![0, 1, 0, 1, 1]], + vec![3, 2], + None, + 0.0, // impossible: |delta| < 0 is never true + 1, + ); + assert_eq!(it, -1); + } + + #[test] + fn nan_in_variable_never_converges() { + let mut x = vec![1.0, f64::NAN, 3.0, 4.0]; + let it = run(&mut x, vec![vec![0, 0, 1, 1]], vec![2], None, 1e-8, 50); + assert_eq!(it, -1); // NaN delta compares false -> poisons convergence + } + + #[test] + fn two_way_matches_naive_reference() { + // small unbalanced two-way panel; reference = many-iteration run + let x0 = vec![2.0, -1.0, 0.5, 3.0, -2.0, 1.0, 4.0]; + let codes = vec![ + vec![0usize, 0, 1, 1, 2, 2, 2], + vec![0usize, 1, 0, 2, 1, 2, 0], + ]; + let mut a = x0.clone(); + let it_a = run(&mut a, codes.clone(), vec![3, 3], None, 1e-13, 10_000); + assert!(it_a >= 1); + // group means ~0 in every dimension + for (codes_d, &n_g) in codes.iter().zip([3usize, 3].iter()) { + let mut sums = vec![0.0f64; n_g]; + let mut cnt = vec![0.0f64; n_g]; + for i in 0..a.len() { + sums[codes_d[i]] += a[i]; + cnt[codes_d[i]] += 1.0; + } + for g in 0..n_g { + assert!((sums[g] / cnt[g]).abs() < 1e-10); + } + } + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index cc8fd13af..07e87b5d5 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -15,6 +15,7 @@ use pyo3::prelude::*; use std::collections::HashMap; mod bootstrap; +mod demean; mod linalg; mod trop; mod weights; @@ -28,6 +29,9 @@ fn _rust_backend(m: &Bound<'_, PyModule>) -> PyResult<()> { m )?)?; + // FE-absorption MAP demeaning kernel + m.add_function(wrap_pyfunction!(demean::demean_map, m)?)?; + // Simplex projection m.add_function(wrap_pyfunction!(weights::project_simplex, m)?)?; @@ -38,7 +42,10 @@ fn _rust_backend(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(weights::sc_weight_fw, m)?)?; m.add_function(wrap_pyfunction!(weights::sc_weight_fw_with_convergence, m)?)?; m.add_function(wrap_pyfunction!(weights::sc_weight_fw_weighted, m)?)?; - m.add_function(wrap_pyfunction!(weights::sc_weight_fw_weighted_with_convergence, m)?)?; + m.add_function(wrap_pyfunction!( + weights::sc_weight_fw_weighted_with_convergence, + m + )?)?; // Linear algebra operations m.add_function(wrap_pyfunction!(linalg::solve_ols, m)?)?; diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 1aa64b208..72511c70c 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -10,6 +10,7 @@ """ import numpy as np +import pandas as pd import pytest from diff_diff import HAS_RUST_BACKEND @@ -128,8 +129,18 @@ def test_bootstrap_weights_bit_identity_snapshot(self): ), "mammen": np.array( [ - [1.618033988749895, -0.6180339887498949, 1.618033988749895, -0.6180339887498949], - [-0.6180339887498949, -0.6180339887498949, 1.618033988749895, 1.618033988749895], + [ + 1.618033988749895, + -0.6180339887498949, + 1.618033988749895, + -0.6180339887498949, + ], + [ + -0.6180339887498949, + -0.6180339887498949, + 1.618033988749895, + 1.618033988749895, + ], ] ), "webb": np.array( @@ -346,8 +357,10 @@ def test_near_singular_matrix_lu_fallback(self): # Verify residuals are correct given coefficients expected_residuals = y - X @ coeffs np.testing.assert_array_almost_equal( - residuals, expected_residuals, decimal=8, - err_msg="Residuals should match y - X @ coeffs" + residuals, + expected_residuals, + decimal=8, + err_msg="Residuals should match y - X @ coeffs", ) def test_high_condition_number_matrix(self): @@ -431,8 +444,10 @@ def test_rank_deficient_matrix_produces_valid_coefficients(self): # Residuals should be correct given coefficients expected_residuals = y - X @ coeffs np.testing.assert_array_almost_equal( - residuals, expected_residuals, decimal=8, - err_msg="Residuals should match y - X @ coeffs" + residuals, + expected_residuals, + decimal=8, + err_msg="Residuals should match y - X @ coeffs", ) def test_multiperiod_did_like_design_matrix(self): @@ -491,9 +506,9 @@ def test_multiperiod_did_like_design_matrix(self): assert np.all(np.abs(coeffs) < 1e6), f"Coefficients are unreasonably large: {coeffs}" # Treatment effect (last coefficient) should be close to true effect - assert abs(coeffs[-1] - true_effect) < 2.0, ( - f"Treatment effect {coeffs[-1]} is too far from true effect {true_effect}" - ) + assert ( + abs(coeffs[-1] - true_effect) < 2.0 + ), f"Treatment effect {coeffs[-1]} is too far from true effect {true_effect}" @pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") @@ -518,12 +533,10 @@ def test_solve_ols_coefficients_match(self): numpy_coeffs, numpy_resid, numpy_vcov = numpy_fn(X, y, cluster_ids=None) np.testing.assert_array_almost_equal( - rust_coeffs, numpy_coeffs, decimal=8, - err_msg="OLS coefficients should match" + rust_coeffs, numpy_coeffs, decimal=8, err_msg="OLS coefficients should match" ) np.testing.assert_array_almost_equal( - rust_resid, numpy_resid, decimal=8, - err_msg="OLS residuals should match" + rust_resid, numpy_resid, decimal=8, err_msg="OLS residuals should match" ) def test_solve_ols_with_clusters_match(self): @@ -542,13 +555,11 @@ def test_solve_ols_with_clusters_match(self): numpy_coeffs, _, numpy_vcov = numpy_fn(X, y, cluster_ids=cluster_ids) np.testing.assert_array_almost_equal( - rust_coeffs, numpy_coeffs, decimal=8, - err_msg="Clustered OLS coefficients should match" + rust_coeffs, numpy_coeffs, decimal=8, err_msg="Clustered OLS coefficients should match" ) # VCoV may differ slightly due to implementation details np.testing.assert_array_almost_equal( - rust_vcov, numpy_vcov, decimal=5, - err_msg="Clustered OLS VCoV should match" + rust_vcov, numpy_vcov, decimal=5, err_msg="Clustered OLS VCoV should match" ) def test_rank_deficient_ols_residuals_match(self): @@ -597,8 +608,10 @@ def test_rank_deficient_ols_residuals_match(self): # Residuals should be very close (this is the key equivalence check) # Both approaches should produce the same fitted values and residuals np.testing.assert_array_almost_equal( - rust_resid, numpy_resid, decimal=5, - err_msg="Residuals should match despite different coefficient representations" + rust_resid, + numpy_resid, + decimal=5, + err_msg="Residuals should match despite different coefficient representations", ) def test_multiperiod_did_design_residuals_equivalence(self): @@ -648,25 +661,27 @@ def test_multiperiod_did_design_residuals_equivalence(self): # Rust should produce finite treatment effect rust_effect = rust_coeffs[-1] assert np.isfinite(rust_effect), "Rust treatment effect should be finite" - assert abs(rust_effect - true_effect) < 2.0, ( - f"Rust treatment effect {rust_effect} too far from true {true_effect}" - ) + assert ( + abs(rust_effect - true_effect) < 2.0 + ), f"Rust treatment effect {rust_effect} too far from true {true_effect}" # NumPy treatment effect should be close (may be finite or NaN depending on rank) numpy_effect = numpy_coeffs[-1] if np.isfinite(numpy_effect): - assert abs(numpy_effect - true_effect) < 2.0, ( - f"NumPy treatment effect {numpy_effect} too far from true {true_effect}" - ) + assert ( + abs(numpy_effect - true_effect) < 2.0 + ), f"NumPy treatment effect {numpy_effect} too far from true {true_effect}" # Effects should be close to each other - assert abs(rust_effect - numpy_effect) < 0.5, ( - f"Rust ({rust_effect}) and NumPy ({numpy_effect}) effects should match" - ) + assert ( + abs(rust_effect - numpy_effect) < 0.5 + ), f"Rust ({rust_effect}) and NumPy ({numpy_effect}) effects should match" # Residuals should be very close (key equivalence check) np.testing.assert_array_almost_equal( - rust_resid, numpy_resid, decimal=5, - err_msg="Residuals should match for MultiPeriodDiD-like design" + rust_resid, + numpy_resid, + decimal=5, + err_msg="Residuals should match for MultiPeriodDiD-like design", ) # ========================================================================= @@ -687,8 +702,7 @@ def test_robust_vcov_hc1_match(self): numpy_vcov = numpy_fn(X, residuals, None) np.testing.assert_array_almost_equal( - rust_vcov, numpy_vcov, decimal=8, - err_msg="HC1 robust VCoV should match" + rust_vcov, numpy_vcov, decimal=8, err_msg="HC1 robust VCoV should match" ) def test_robust_vcov_clustered_match(self): @@ -707,8 +721,7 @@ def test_robust_vcov_clustered_match(self): numpy_vcov = numpy_fn(X, residuals, cluster_ids) np.testing.assert_array_almost_equal( - rust_vcov, numpy_vcov, decimal=6, - err_msg="Cluster-robust VCoV should match" + rust_vcov, numpy_vcov, decimal=6, err_msg="Cluster-robust VCoV should match" ) # ========================================================================= @@ -744,10 +757,10 @@ def test_bootstrap_weights_mammen_properties(self): mean = weights.mean() assert abs(mean) < 0.02, f"Mammen mean should be ~0, got {mean}" - second_moment = (weights ** 2).mean() + second_moment = (weights**2).mean() assert abs(second_moment - 1.0) < 0.02, f"Mammen E[w^2] should be ~1, got {second_moment}" - third_moment = (weights ** 3).mean() + third_moment = (weights**3).mean() assert abs(third_moment - 1.0) < 0.1, f"Mammen E[w^3] should be ~1, got {third_moment}" def test_bootstrap_weights_webb_properties(self): @@ -793,8 +806,10 @@ def test_simplex_projection_match(self): numpy_proj = numpy_fn(v) np.testing.assert_array_almost_equal( - rust_proj, numpy_proj, decimal=10, - err_msg=f"Simplex projection mismatch for input {v}" + rust_proj, + numpy_proj, + decimal=10, + err_msg=f"Simplex projection mismatch for input {v}", ) def test_nan_vcov_fallback_to_python(self): @@ -834,17 +849,16 @@ def test_nan_vcov_fallback_to_python(self): # Check if fallback warning was emitted fallback_warning_emitted = any( - "Re-running with Python backend" in str(warning.message) - for warning in w + "Re-running with Python backend" in str(warning.message) for warning in w ) # Key invariants that must hold regardless of which backend is used: # 1. Coefficients must be finite (either via Rust SVD or Python R-style) finite_coeffs = coeffs[np.isfinite(coeffs)] - assert len(finite_coeffs) >= 3, \ - "At least 3 coefficients should be finite (identifiable)" - assert np.all(np.abs(finite_coeffs) < 1e10), \ - f"Finite coefficients should be reasonable, got {finite_coeffs}" + assert len(finite_coeffs) >= 3, "At least 3 coefficients should be finite (identifiable)" + assert np.all( + np.abs(finite_coeffs) < 1e10 + ), f"Finite coefficients should be reasonable, got {finite_coeffs}" # 2. If vcov has any finite values, they should correspond to finite coefficients if vcov is not None: @@ -853,8 +867,9 @@ def test_nan_vcov_fallback_to_python(self): if finite_coef_mask[i]: # This coefficient's variance should be finite var_i = vcov[i, i] - assert np.isfinite(var_i) or np.isnan(var_i), \ - f"Variance for finite coef {i} should be finite or NaN (dropped)" + assert np.isfinite(var_i) or np.isnan( + var_i + ), f"Variance for finite coef {i} should be finite or NaN (dropped)" # 3. Residuals must always be finite assert np.all(np.isfinite(residuals)), "Residuals should be finite" @@ -865,20 +880,22 @@ def test_nan_vcov_fallback_to_python(self): nan_vcov_diag_indices = set(np.where(np.isnan(np.diag(vcov)))[0]) # NaN in vcov diagonal should correspond exactly to NaN coefficients - assert nan_vcov_diag_indices == nan_coef_indices, \ - f"NaN vcov diagonal {nan_vcov_diag_indices} should match " \ + assert nan_vcov_diag_indices == nan_coef_indices, ( + f"NaN vcov diagonal {nan_vcov_diag_indices} should match " f"NaN coefficients {nan_coef_indices}" + ) # 5. If fallback warning was emitted, R-style handling MUST have occurred # This verifies that the fallback actually applies R-style NaN handling # (not minimum-norm solution which would have all finite coefficients) if fallback_warning_emitted: - assert np.any(np.isnan(coeffs)), \ - "Fallback warning emitted but no NaN coefficients - " \ - "R-style handling was not applied" - assert vcov is not None and np.any(np.isnan(vcov)), \ - "Fallback warning emitted but vcov has no NaN - " \ + assert np.any(np.isnan(coeffs)), ( + "Fallback warning emitted but no NaN coefficients - " "R-style handling was not applied" + ) + assert vcov is not None and np.any(np.isnan(vcov)), ( + "Fallback warning emitted but vcov has no NaN - " "R-style handling was not applied" + ) @pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") @@ -941,8 +958,7 @@ def test_unit_distance_matrix_matches_numpy(self): numpy_dist = trop._compute_all_unit_distances(Y, D, n_units, n_periods) np.testing.assert_array_almost_equal( - rust_dist, numpy_dist, decimal=10, - err_msg="Distance matrices should match" + rust_dist, numpy_dist, decimal=10, err_msg="Distance matrices should match" ) def test_unit_distance_excludes_treated(self): @@ -985,9 +1001,15 @@ def test_loocv_grid_search_returns_valid_params(self): lambda_nn = np.array([0.0, 0.1], dtype=np.float64) best_lt, best_lu, best_ln, score, n_valid, n_attempted, first_failed = loocv_grid_search( - Y, D, control_mask, time_dist, - lambda_time, lambda_unit, lambda_nn, - 100, 1e-6, + Y, + D, + control_mask, + time_dist, + lambda_time, + lambda_unit, + lambda_nn, + 100, + 1e-6, ) # Check returned parameters are from the grid @@ -1033,9 +1055,18 @@ def test_bootstrap_variance_shape(self): n_control=5, n_treated=1, n_bootstrap=n_bootstrap, seed=42 ) estimates, se = bootstrap_trop_variance( - Y, D, control_mask, time_dist, - 1.0, 1.0, 0.1, # lambda values - n_bootstrap, 100, 1e-6, ctrl_idx, trt_idx, + Y, + D, + control_mask, + time_dist, + 1.0, + 1.0, + 0.1, # lambda values + n_bootstrap, + 100, + 1e-6, + ctrl_idx, + trt_idx, ) # Should return array of bootstrap estimates and SE @@ -1067,12 +1098,32 @@ def test_bootstrap_reproducibility(self): n_control=5, n_treated=1, n_bootstrap=20, seed=42 ) est1, se1 = bootstrap_trop_variance( - Y, D, control_mask, time_dist, - 1.0, 1.0, 0.1, 20, 100, 1e-6, ctrl_idx_a, trt_idx_a, + Y, + D, + control_mask, + time_dist, + 1.0, + 1.0, + 0.1, + 20, + 100, + 1e-6, + ctrl_idx_a, + trt_idx_a, ) est2, se2 = bootstrap_trop_variance( - Y, D, control_mask, time_dist, - 1.0, 1.0, 0.1, 20, 100, 1e-6, ctrl_idx_b, trt_idx_b, + Y, + D, + control_mask, + time_dist, + 1.0, + 1.0, + 0.1, + 20, + 100, + 1e-6, + ctrl_idx_b, + trt_idx_b, ) np.testing.assert_array_almost_equal(est1, est2) @@ -1098,8 +1149,18 @@ def test_bootstrap_rejects_negative_index(self): ctrl_idx[2, 3] = -1 # negative with pytest.raises(ValueError, match="control_indices.*out-of-range"): bootstrap_trop_variance( - Y, D, control_mask, time_dist, - 1.0, 1.0, 0.1, 5, 100, 1e-6, ctrl_idx, trt_idx, + Y, + D, + control_mask, + time_dist, + 1.0, + 1.0, + 0.1, + 5, + 100, + 1e-6, + ctrl_idx, + trt_idx, ) def test_bootstrap_rejects_out_of_range_index(self): @@ -1122,8 +1183,18 @@ def test_bootstrap_rejects_out_of_range_index(self): trt_idx[1, 0] = 99 # >> n_treated=1 with pytest.raises(ValueError, match="treated_indices.*out-of-range"): bootstrap_trop_variance( - Y, D, control_mask, time_dist, - 1.0, 1.0, 0.1, 5, 100, 1e-6, ctrl_idx, trt_idx, + Y, + D, + control_mask, + time_dist, + 1.0, + 1.0, + 0.1, + 5, + 100, + 1e-6, + ctrl_idx, + trt_idx, ) @@ -1152,8 +1223,7 @@ def test_distance_matrix_matches_numpy(self): numpy_dist = trop._compute_all_unit_distances(Y, D, n_units, n_periods) np.testing.assert_array_almost_equal( - rust_dist, numpy_dist, decimal=10, - err_msg="Distance matrices should match exactly" + rust_dist, numpy_dist, decimal=10, err_msg="Distance matrices should match exactly" ) def test_trop_produces_valid_results(self): @@ -1172,13 +1242,14 @@ def test_trop_produces_valid_results(self): for i in range(n_units): for t in range(n_periods): is_treated = (i == 0) and (t >= 6) - y = 1.0 + 0.5 * i + 0.3 * t + (true_effect if is_treated else 0) + np.random.randn() * 0.5 - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': 1 if is_treated else 0 - }) + y = ( + 1.0 + + 0.5 * i + + 0.3 * t + + (true_effect if is_treated else 0) + + np.random.randn() * 0.5 + ) + data.append({"unit": i, "time": t, "outcome": y, "treated": 1 if is_treated else 0}) df = pd.DataFrame(data) @@ -1188,9 +1259,9 @@ def test_trop_produces_valid_results(self): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=20, - seed=42 + seed=42, ) - results = trop.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop.fit(df, "outcome", "treated", "unit", "time") # Check results are valid assert np.isfinite(results.att), "ATT should be finite" @@ -1204,8 +1275,9 @@ def test_trop_produces_valid_results(self): # - LOOCV-selected tuning parameters may not be optimal for small samples # This is a validity test, not a precision test - we're checking the # estimation produces sensible results, not exact recovery. - assert abs(results.att - true_effect) < 2.0, \ - f"ATT {results.att:.2f} should be close to true effect {true_effect}" + assert ( + abs(results.att - true_effect) < 2.0 + ), f"ATT {results.att:.2f} should be close to true effect {true_effect}" # Tuning parameters should be from the grid assert results.lambda_time in [0.0, 1.0] @@ -1238,9 +1310,14 @@ def test_loocv_grid_search_global_returns_valid_result(self): lambda_nn_grid = np.array([0.0, 0.1]) result = loocv_grid_search_global( - Y, D, control_mask, - lambda_time_grid, lambda_unit_grid, lambda_nn_grid, - 100, 1e-6, + Y, + D, + control_mask, + lambda_time_grid, + lambda_unit_grid, + lambda_nn_grid, + 100, + 1e-6, ) best_lt, best_lu, best_ln, best_score, n_valid, n_attempted, _ = result @@ -1275,14 +1352,24 @@ def test_loocv_grid_search_global_reproducible(self): lambda_nn_grid = np.array([0.0, 0.1]) result1 = loocv_grid_search_global( - Y, D, control_mask, - lambda_time_grid, lambda_unit_grid, lambda_nn_grid, - 50, 1e-6, + Y, + D, + control_mask, + lambda_time_grid, + lambda_unit_grid, + lambda_nn_grid, + 50, + 1e-6, ) result2 = loocv_grid_search_global( - Y, D, control_mask, - lambda_time_grid, lambda_unit_grid, lambda_nn_grid, - 50, 1e-6, + Y, + D, + control_mask, + lambda_time_grid, + lambda_unit_grid, + lambda_nn_grid, + 50, + 1e-6, ) # Without subsampling, results should be deterministic @@ -1314,7 +1401,16 @@ def test_bootstrap_trop_variance_global_shape(self): n_control=n_units - n_treated, n_treated=n_treated, n_bootstrap=50, seed=42 ) estimates, se = bootstrap_trop_variance_global( - Y, D, 0.5, 0.5, 0.1, 50, 50, 1e-6, ctrl_idx, trt_idx, + Y, + D, + 0.5, + 0.5, + 0.1, + 50, + 50, + 1e-6, + ctrl_idx, + trt_idx, ) assert isinstance(estimates, np.ndarray) @@ -1342,10 +1438,28 @@ def test_bootstrap_trop_variance_global_reproducible(self): n_control=n_units - n_treated, n_treated=n_treated, n_bootstrap=50, seed=42 ) est1, se1 = bootstrap_trop_variance_global( - Y, D, 0.5, 0.5, 0.1, 50, 50, 1e-6, ctrl_a, trt_a, + Y, + D, + 0.5, + 0.5, + 0.1, + 50, + 50, + 1e-6, + ctrl_a, + trt_a, ) est2, se2 = bootstrap_trop_variance_global( - Y, D, 0.5, 0.5, 0.1, 50, 50, 1e-6, ctrl_b, trt_b, + Y, + D, + 0.5, + 0.5, + 0.1, + 50, + 50, + 1e-6, + ctrl_b, + trt_b, ) np.testing.assert_array_almost_equal(est1, est2) @@ -1368,7 +1482,16 @@ def test_bootstrap_global_rejects_negative_index(self): ctrl_idx[3, 2] = -5 # negative with pytest.raises(ValueError, match="control_indices.*out-of-range"): bootstrap_trop_variance_global( - Y, D, 0.5, 0.5, 0.1, 10, 50, 1e-6, ctrl_idx, trt_idx, + Y, + D, + 0.5, + 0.5, + 0.1, + 10, + 50, + 1e-6, + ctrl_idx, + trt_idx, ) def test_bootstrap_global_rejects_out_of_range_index(self): @@ -1388,7 +1511,16 @@ def test_bootstrap_global_rejects_out_of_range_index(self): trt_idx[5, 1] = 99 # >> n_treated=4 with pytest.raises(ValueError, match="treated_indices.*out-of-range"): bootstrap_trop_variance_global( - Y, D, 0.5, 0.5, 0.1, 10, 50, 1e-6, ctrl_idx, trt_idx, + Y, + D, + 0.5, + 0.5, + 0.1, + 10, + 50, + 1e-6, + ctrl_idx, + trt_idx, ) @@ -1417,12 +1549,14 @@ def test_trop_global_produces_valid_results(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_effect - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1432,9 +1566,9 @@ def test_trop_global_produces_valid_results(self): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=30, - seed=42 + seed=42, ) - results = trop.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop.fit(df, "outcome", "treated", "unit", "time") # Check results are valid assert np.isfinite(results.att), "ATT should be finite" @@ -1469,12 +1603,14 @@ def test_trop_global_and_local_agree_in_direction(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_effect - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1485,9 +1621,9 @@ def test_trop_global_and_local_agree_in_direction(self): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=20, - seed=42 + seed=42, ) - results_global = trop_global.fit(df, 'outcome', 'treated', 'unit', 'time') + results_global = trop_global.fit(df, "outcome", "treated", "unit", "time") # Fit with local method trop_local = TROP( @@ -1496,9 +1632,9 @@ def test_trop_global_and_local_agree_in_direction(self): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=20, - seed=42 + seed=42, ) - results_local = trop_local.fit(df, 'outcome', 'treated', 'unit', 'time') + results_local = trop_local.fit(df, "outcome", "treated", "unit", "time") # Both should have same sign (both positive for true_effect=2.0) assert np.sign(results_global.att) == np.sign(results_local.att) @@ -1523,12 +1659,14 @@ def test_trop_global_handles_nan_outcomes(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_effect - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1536,10 +1674,10 @@ def test_trop_global_handles_nan_outcomes(self): # Set 5% of control pre-treatment observations to NaN nan_indices = [] for idx, row in df.iterrows(): - if row['treated'] == 0 and row['time'] < (n_periods - n_post): + if row["treated"] == 0 and row["time"] < (n_periods - n_post): if np.random.rand() < 0.05: nan_indices.append(idx) - df.loc[nan_indices, 'outcome'] = np.nan + df.loc[nan_indices, "outcome"] = np.nan n_nan = len(nan_indices) assert n_nan > 0, "Should have introduced some NaN values" @@ -1550,9 +1688,9 @@ def test_trop_global_handles_nan_outcomes(self): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=20, - seed=42 + seed=42, ) - results = trop.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop.fit(df, "outcome", "treated", "unit", "time") # Results should be finite (NaN observations are excluded) assert np.isfinite(results.att), f"ATT {results.att} should be finite with NaN data" @@ -1591,24 +1729,30 @@ def test_trop_global_no_valid_pre_unit_gets_zero_weight(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_effect - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) # Set ALL pre-period outcomes to NaN for one control unit (unit n_treated) # This unit has no valid pre-period data and should get zero weight control_unit_with_no_pre = n_treated # First control unit - pre_mask = (df['unit'] == control_unit_with_no_pre) & (df['time'] < (n_periods - n_post)) - df.loc[pre_mask, 'outcome'] = np.nan + pre_mask = (df["unit"] == control_unit_with_no_pre) & (df["time"] < (n_periods - n_post)) + df.loc[pre_mask, "outcome"] = np.nan # Verify we set NaN correctly - unit_pre_data = df[(df['unit'] == control_unit_with_no_pre) & (df['time'] < (n_periods - n_post))] - assert unit_pre_data['outcome'].isna().all(), "Control unit should have all NaN in pre-period" + unit_pre_data = df[ + (df["unit"] == control_unit_with_no_pre) & (df["time"] < (n_periods - n_post)) + ] + assert ( + unit_pre_data["outcome"].isna().all() + ), "Control unit should have all NaN in pre-period" # Fit with global method - should handle gracefully trop = TROP( @@ -1617,9 +1761,9 @@ def test_trop_global_no_valid_pre_unit_gets_zero_weight(self): lambda_unit_grid=[0.5, 1.0], lambda_nn_grid=[0.0], n_bootstrap=20, - seed=42 + seed=42, ) - results = trop.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop.fit(df, "outcome", "treated", "unit", "time") # Results should be finite - the unit with no valid pre-period data # should get zero weight and not break estimation @@ -1628,8 +1772,9 @@ def test_trop_global_no_valid_pre_unit_gets_zero_weight(self): # ATT should be in reasonable range of true effect # The no-valid-pre unit getting zero weight shouldn't corrupt the estimate - assert abs(results.att - true_effect) < 1.5, \ - f"ATT {results.att:.2f} should be close to true effect {true_effect}" + assert ( + abs(results.att - true_effect) < 1.5 + ), f"ATT {results.att:.2f} should be close to true effect {true_effect}" def test_trop_global_nan_exclusion_rust_python_parity(self): """Test Rust and Python backends produce matching results with NaN data. @@ -1660,23 +1805,25 @@ def test_trop_global_nan_exclusion_rust_python_parity(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_effect - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) # Introduce scattered NaN values (5% of control pre-period observations) np.random.seed(123) # Different seed for NaN placement for idx, row in df.iterrows(): - if row['treated'] == 0 and row['time'] < (n_periods - n_post): + if row["treated"] == 0 and row["time"] < (n_periods - n_post): if np.random.rand() < 0.05: - df.loc[idx, 'outcome'] = np.nan + df.loc[idx, "outcome"] = np.nan - n_nan = df['outcome'].isna().sum() + n_nan = df["outcome"].isna().sum() assert n_nan > 0, "Should have some NaN values" # Common TROP parameters @@ -1686,25 +1833,28 @@ def test_trop_global_nan_exclusion_rust_python_parity(self): lambda_unit_grid=[0.5, 1.0], lambda_nn_grid=[0.0], n_bootstrap=20, - seed=42 + seed=42, ) # Run with Rust backend (current default when available) trop_rust = TROP(**trop_params) - results_rust = trop_rust.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") # Run with Python-only backend using mock.patch to avoid module reload issues # (Module reload breaks isinstance() checks in other tests due to class identity) from unittest.mock import patch import sys - trop_global_module = sys.modules['diff_diff.trop_global'] - with patch.object(trop_global_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_global_module, '_rust_loocv_grid_search_global', None), \ - patch.object(trop_global_module, '_rust_bootstrap_trop_variance_global', None): + trop_global_module = sys.modules["diff_diff.trop_global"] + + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + ): trop_python = TROP(**trop_params) - results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_python = trop_python.fit(df.copy(), "outcome", "treated", "unit", "time") # Both should produce finite results assert np.isfinite(results_rust.att), f"Rust ATT {results_rust.att} should be finite" @@ -1713,9 +1863,10 @@ def test_trop_global_nan_exclusion_rust_python_parity(self): # ATT estimates should be close (within reasonable tolerance) # Allow some difference due to LOOCV randomness and numerical differences att_diff = abs(results_rust.att - results_python.att) - assert att_diff < 0.5, \ - f"Rust ATT ({results_rust.att:.3f}) and Python ATT ({results_python.att:.3f}) " \ + assert att_diff < 0.5, ( + f"Rust ATT ({results_rust.att:.3f}) and Python ATT ({results_python.att:.3f}) " f"differ by {att_diff:.3f}, should be < 0.5" + ) # Both should recover true effect direction assert results_rust.att > 0, f"Rust ATT {results_rust.att} should be positive" @@ -1749,12 +1900,14 @@ def test_trop_global_treated_pre_nan_rust_python_parity(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_effect - data.append({ - 'unit': i, - 'time': t, - 'outcome': y, - 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1762,11 +1915,11 @@ def test_trop_global_treated_pre_nan_rust_python_parity(self): # This makes average_treated[3] = NaN target_period = 3 treated_units = list(range(n_treated)) - mask = df['unit'].isin(treated_units) & (df['time'] == target_period) - df.loc[mask, 'outcome'] = np.nan + mask = df["unit"].isin(treated_units) & (df["time"] == target_period) + df.loc[mask, "outcome"] = np.nan # Verify we set NaN correctly - n_nan = df.loc[mask, 'outcome'].isna().sum() + n_nan = df.loc[mask, "outcome"].isna().sum() assert n_nan == n_treated, f"Should have {n_treated} NaN, got {n_nan}" # Common TROP parameters @@ -1776,25 +1929,28 @@ def test_trop_global_treated_pre_nan_rust_python_parity(self): lambda_unit_grid=[1.0], lambda_nn_grid=[0.0], n_bootstrap=20, - seed=42 + seed=42, ) # Run with Rust backend (current default when available) trop_rust = TROP(**trop_params) - results_rust = trop_rust.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") # Run with Python-only backend using mock.patch to avoid module reload issues # (Module reload breaks isinstance() checks in other tests due to class identity) from unittest.mock import patch import sys - trop_global_module = sys.modules['diff_diff.trop_global'] - with patch.object(trop_global_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_global_module, '_rust_loocv_grid_search_global', None), \ - patch.object(trop_global_module, '_rust_bootstrap_trop_variance_global', None): + trop_global_module = sys.modules["diff_diff.trop_global"] + + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + ): trop_python = TROP(**trop_params) - results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_python = trop_python.fit(df.copy(), "outcome", "treated", "unit", "time") # Both should produce finite results assert np.isfinite(results_rust.att), f"Rust ATT {results_rust.att} should be finite" @@ -1802,9 +1958,10 @@ def test_trop_global_treated_pre_nan_rust_python_parity(self): # ATT estimates should be close (within reasonable tolerance) att_diff = abs(results_rust.att - results_python.att) - assert att_diff < 0.5, \ - f"Rust ATT ({results_rust.att:.3f}) and Python ATT ({results_python.att:.3f}) " \ + assert att_diff < 0.5, ( + f"Rust ATT ({results_rust.att:.3f}) and Python ATT ({results_python.att:.3f}) " f"differ by {att_diff:.3f}, should be < 0.5" + ) def test_trop_global_solver_parity_no_lowrank(self): """Test Rust/Python solver parity for no-lowrank path (lambda_nn >= 1e10). @@ -1831,10 +1988,14 @@ def test_trop_global_solver_parity_no_lowrank(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += 2.0 - data.append({ - 'unit': i, 'time': t, - 'outcome': y, 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) # Fixed lambda with lambda_nn=inf (no low-rank) @@ -1849,32 +2010,37 @@ def test_trop_global_solver_parity_no_lowrank(self): # Rust backend trop_rust = TROP(**trop_params) - results_rust = trop_rust.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") # Python-only backend - trop_global_module = sys.modules['diff_diff.trop_global'] - with patch.object(trop_global_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_global_module, '_rust_loocv_grid_search_global', None), \ - patch.object(trop_global_module, '_rust_bootstrap_trop_variance_global', None): + trop_global_module = sys.modules["diff_diff.trop_global"] + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + ): trop_python = TROP(**trop_params) - results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_python = trop_python.fit(df.copy(), "outcome", "treated", "unit", "time") # ATT should match closely - assert abs(results_rust.att - results_python.att) < 1e-6, \ - f"No-lowrank ATT mismatch: Rust={results_rust.att:.8f}, Python={results_python.att:.8f}" + assert ( + abs(results_rust.att - results_python.att) < 1e-6 + ), f"No-lowrank ATT mismatch: Rust={results_rust.att:.8f}, Python={results_python.att:.8f}" # Unit and time effects should match for key in results_rust.unit_effects: r_val = results_rust.unit_effects[key] p_val = results_python.unit_effects[key] - assert abs(r_val - p_val) < 1e-6, \ - f"Unit effect mismatch for {key}: Rust={r_val:.8f}, Python={p_val:.8f}" + assert ( + abs(r_val - p_val) < 1e-6 + ), f"Unit effect mismatch for {key}: Rust={r_val:.8f}, Python={p_val:.8f}" for key in results_rust.time_effects: r_val = results_rust.time_effects[key] p_val = results_python.time_effects[key] - assert abs(r_val - p_val) < 1e-6, \ - f"Time effect mismatch for {key}: Rust={r_val:.8f}, Python={p_val:.8f}" + assert ( + abs(r_val - p_val) < 1e-6 + ), f"Time effect mismatch for {key}: Rust={r_val:.8f}, Python={p_val:.8f}" def test_trop_global_solver_parity_with_lowrank(self): """Test Rust/Python solver parity for with-lowrank path (finite lambda_nn). @@ -1902,10 +2068,14 @@ def test_trop_global_solver_parity_with_lowrank(self): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += 2.0 - data.append({ - 'unit': i, 'time': t, - 'outcome': y, 'treated': treatment_indicator, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) # Fixed lambda with finite lambda_nn (low-rank enabled) @@ -1920,26 +2090,30 @@ def test_trop_global_solver_parity_with_lowrank(self): # Rust backend trop_rust = TROP(**trop_params) - results_rust = trop_rust.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") # Python-only backend - trop_global_module = sys.modules['diff_diff.trop_global'] - with patch.object(trop_global_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_global_module, '_rust_loocv_grid_search_global', None), \ - patch.object(trop_global_module, '_rust_bootstrap_trop_variance_global', None): + trop_global_module = sys.modules["diff_diff.trop_global"] + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + ): trop_python = TROP(**trop_params) - results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time') + results_python = trop_python.fit(df.copy(), "outcome", "treated", "unit", "time") # ATT should match closely - assert abs(results_rust.att - results_python.att) < 1e-6, \ - f"With-lowrank ATT mismatch: Rust={results_rust.att:.8f}, Python={results_python.att:.8f}" + assert ( + abs(results_rust.att - results_python.att) < 1e-6 + ), f"With-lowrank ATT mismatch: Rust={results_rust.att:.8f}, Python={results_python.att:.8f}" # Unit and time effects should match for key in results_rust.unit_effects: r_val = results_rust.unit_effects[key] p_val = results_python.unit_effects[key] - assert abs(r_val - p_val) < 1e-6, \ - f"Unit effect mismatch for {key}: Rust={r_val:.8f}, Python={p_val:.8f}" + assert ( + abs(r_val - p_val) < 1e-6 + ), f"Unit effect mismatch for {key}: Rust={r_val:.8f}, Python={p_val:.8f}" @pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") @@ -1955,8 +2129,9 @@ def test_noise_level_matches_numpy(self): Y_pre = np.random.randn(10, 5) rust_nl = rust_fn(Y_pre) numpy_nl = numpy_fn(Y_pre) - assert abs(rust_nl - numpy_nl) < 1e-10, \ - f"Noise levels differ: rust={rust_nl}, numpy={numpy_nl}" + assert ( + abs(rust_nl - numpy_nl) < 1e-10 + ), f"Noise levels differ: rust={rust_nl}, numpy={numpy_nl}" def test_noise_level_single_period(self): """Test noise level returns 0 for single pre-period.""" @@ -1986,8 +2161,7 @@ def test_sc_weight_fw_matches_numpy(self): rust_w = rust_fn(Y, 0.5, True, None, 1e-3, 1000) numpy_w = numpy_fn(Y, 0.5, True, None, 1e-3, 1000) np.testing.assert_array_almost_equal( - rust_w, numpy_w, decimal=6, - err_msg="Frank-Wolfe weights should match" + rust_w, numpy_w, decimal=6, err_msg="Frank-Wolfe weights should match" ) def test_sc_weight_fw_with_init_weights(self): @@ -2001,8 +2175,7 @@ def test_sc_weight_fw_with_init_weights(self): rust_w = rust_fn(Y, 0.2, True, init_w, 1e-3, 500) numpy_w = numpy_fn(Y, 0.2, True, init_w, 1e-3, 500) np.testing.assert_array_almost_equal( - rust_w, numpy_w, decimal=6, - err_msg="Frank-Wolfe with init weights should match" + rust_w, numpy_w, decimal=6, err_msg="Frank-Wolfe with init weights should match" ) def test_time_weights_on_simplex(self): @@ -2031,21 +2204,17 @@ def test_time_weights_match_numpy(self): max_iter = 1000 # Rust implementation (2-pass with sparsification) - rust_w = rust_fn(Y_pre, Y_post, 0.01, True, min_decrease, - max_iter_pre, max_iter) + rust_w = rust_fn(Y_pre, Y_post, 0.01, True, min_decrease, max_iter_pre, max_iter) # Python implementation (manual 2-pass matching Rust) post_means = np.mean(Y_post, axis=0) Y_time = np.column_stack([Y_pre.T, post_means]) - lam = _sc_weight_fw_numpy(Y_time, 0.01, True, None, - min_decrease, max_iter_pre) + lam = _sc_weight_fw_numpy(Y_time, 0.01, True, None, min_decrease, max_iter_pre) lam = _sparsify(lam) - numpy_w = _sc_weight_fw_numpy(Y_time, 0.01, True, lam, - min_decrease, max_iter) + numpy_w = _sc_weight_fw_numpy(Y_time, 0.01, True, lam, min_decrease, max_iter) np.testing.assert_array_almost_equal( - rust_w, numpy_w, decimal=6, - err_msg="Time weights should match" + rust_w, numpy_w, decimal=6, err_msg="Time weights should match" ) def test_time_weights_single_preperiod(self): @@ -2089,8 +2258,7 @@ def test_unit_weights_match_numpy(self): numpy_w = _sc_weight_fw_numpy(Y_unit, 0.5, True, omega, 1e-3, 1000) np.testing.assert_array_almost_equal( - rust_w, numpy_w, decimal=6, - err_msg="Unit weights should match" + rust_w, numpy_w, decimal=6, err_msg="Unit weights should match" ) def test_unit_weights_single_control(self): @@ -2123,8 +2291,7 @@ def test_fw_gram_vs_standard_equivalence(self): # Weights must match to high precision np.testing.assert_array_almost_equal( - rust_w, numpy_w, decimal=6, - err_msg="Gram path weights should match Python" + rust_w, numpy_w, decimal=6, err_msg="Gram path weights should match Python" ) assert abs(rust_w.sum() - 1.0) < 1e-6 assert np.all(rust_w >= -1e-6) @@ -2147,8 +2314,7 @@ def test_fw_standard_path_equivalence(self): numpy_w = numpy_fn(Y, 0.5, True, None, 1e-5, 10000) np.testing.assert_array_almost_equal( - rust_w, numpy_w, decimal=6, - err_msg="Standard path weights should match Python" + rust_w, numpy_w, decimal=6, err_msg="Standard path weights should match Python" ) assert abs(rust_w.sum() - 1.0) < 1e-6 assert np.all(rust_w >= -1e-6) @@ -2169,8 +2335,10 @@ def test_sdid_intercept_false_rust_vs_python(self): rust_w_gram = rust_fn(Y_gram, 0.2, False, None, 1e-5, 10000) numpy_w_gram = numpy_fn(Y_gram, 0.2, False, None, 1e-5, 10000) np.testing.assert_array_almost_equal( - rust_w_gram, numpy_w_gram, decimal=6, - err_msg="Gram path intercept=false weights should match Python" + rust_w_gram, + numpy_w_gram, + decimal=6, + err_msg="Gram path intercept=false weights should match Python", ) # Standard path: T0 >= N @@ -2178,8 +2346,10 @@ def test_sdid_intercept_false_rust_vs_python(self): rust_w_std = rust_fn(Y_std, 0.2, False, None, 1e-5, 10000) numpy_w_std = numpy_fn(Y_std, 0.2, False, None, 1e-5, 10000) np.testing.assert_array_almost_equal( - rust_w_std, numpy_w_std, decimal=6, - err_msg="Standard path intercept=false weights should match Python" + rust_w_std, + numpy_w_std, + decimal=6, + err_msg="Standard path intercept=false weights should match Python", ) def test_full_sdid_rust_vs_python(self): @@ -2200,28 +2370,31 @@ def test_full_sdid_rust_vs_python(self): y = 1.0 + 0.5 * i + 0.3 * t + np.random.randn() * 0.3 if is_treated and post: y += true_effect - data.append({ - 'unit': i, 'time': t, 'outcome': y, - 'treated': 1 if is_treated else 0, - 'post': 1 if post else 0, - }) + data.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": 1 if is_treated else 0, + "post": 1 if post else 0, + } + ) df = pd.DataFrame(data) post_periods = list(range(n_pre, n_pre + n_post)) # Run with Rust backend sdid_rust = SyntheticDiD(variance_method="placebo", seed=42) - results_rust = sdid_rust.fit(df, 'outcome', 'treated', 'unit', 'time', post_periods) + results_rust = sdid_rust.fit(df, "outcome", "treated", "unit", "time", post_periods) # Run with Python backend - utils_mod = sys.modules['diff_diff.utils'] - with patch.object(utils_mod, 'HAS_RUST_BACKEND', False): + utils_mod = sys.modules["diff_diff.utils"] + with patch.object(utils_mod, "HAS_RUST_BACKEND", False): sdid_py = SyntheticDiD(variance_method="placebo", seed=42) - results_py = sdid_py.fit(df.copy(), 'outcome', 'treated', 'unit', 'time', post_periods) + results_py = sdid_py.fit(df.copy(), "outcome", "treated", "unit", "time", post_periods) # ATT should be very close np.testing.assert_almost_equal( - results_rust.att, results_py.att, decimal=4, - err_msg="Rust and Python ATT should match" + results_rust.att, results_py.att, decimal=4, err_msg="Rust and Python ATT should match" ) @@ -2266,8 +2439,10 @@ def _run_both_backends(self, X, y): ) linalg_module = sys.modules["diff_diff.linalg"] - with patch.object(linalg_module, "HAS_RUST_BACKEND", False), \ - patch.object(linalg_module, "_rust_solve_ols", None): + with ( + patch.object(linalg_module, "HAS_RUST_BACKEND", False), + patch.object(linalg_module, "_rust_solve_ols", None), + ): coef_py, resid_py, _ = solve_ols( X, y, skip_rank_check=True, vcov_type="hc1", return_vcov=False ) @@ -2278,11 +2453,13 @@ def test_mixed_scale_full_rank(self): should truncate the same singular values.""" rng = np.random.default_rng(11) n, k = 80, 3 - X = np.column_stack([ - rng.normal(0, 1, n), # unit scale - rng.normal(0, 1e6, n), # 1e6 scale - rng.normal(0, 1, n), # unit scale - ]) + X = np.column_stack( + [ + rng.normal(0, 1, n), # unit scale + rng.normal(0, 1e6, n), # 1e6 scale + rng.normal(0, 1, n), # unit scale + ] + ) y = 1.0 + 0.5 * X[:, 0] + 1e-7 * X[:, 1] + 0.3 * X[:, 2] + rng.normal(0, 0.1, n) coef_rust, coef_py = self._run_both_backends(X, y) @@ -2290,7 +2467,10 @@ def test_mixed_scale_full_rank(self): fitted_rust = X @ coef_rust fitted_py = X @ coef_py np.testing.assert_allclose( - fitted_rust, fitted_py, rtol=1e-6, atol=1e-8, + fitted_rust, + fitted_py, + rtol=1e-6, + atol=1e-8, err_msg="Rust vs Python fitted-value divergence on mixed-scale X", ) @@ -2310,7 +2490,10 @@ def test_near_singular_full_rank(self): fitted_rust = X @ coef_rust fitted_py = X @ coef_py np.testing.assert_allclose( - fitted_rust, fitted_py, rtol=1e-6, atol=1e-8, + fitted_rust, + fitted_py, + rtol=1e-6, + atol=1e-8, err_msg="Rust vs Python fitted-value divergence on near-singular X", ) @@ -2331,7 +2514,10 @@ def test_rank_deficient_collinear(self): fitted_py = X @ coef_py # Fitted values are the backend-invariant object under rank deficiency. np.testing.assert_allclose( - fitted_rust, fitted_py, rtol=1e-6, atol=1e-8, + fitted_rust, + fitted_py, + rtol=1e-6, + atol=1e-8, err_msg="Rust vs Python fitted-value divergence on rank-deficient X", ) @@ -2361,6 +2547,7 @@ def _make_correlated_panel(n_units=6, n_periods=5, n_treated=2): """Panel with two control units nearly parallel to each other, making the pre-period Y matrix near rank-deficient.""" import pandas as pd + rng = np.random.default_rng(13) data = [] shared_path = rng.normal(0, 1, n_periods) @@ -2414,16 +2601,20 @@ def test_grid_search_rank_deficient_Y(self): res_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") trop_global_module = sys.modules["diff_diff.trop_global"] - with patch.object(trop_global_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), \ - patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None): + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + ): trop_py = TROP(**trop_params) res_py = trop_py.fit(df.copy(), "outcome", "treated", "unit", "time") # Primary assertion: the ATT point estimate at the chosen λ matches. # This catches both (a) same λ chosen and (b) tied λ producing same fit. np.testing.assert_allclose( - res_rust.att, res_py.att, atol=1e-6, + res_rust.att, + res_py.att, + atol=1e-6, err_msg="Grid-search ATT divergence on rank-deficient Y: " f"Rust={res_rust.att:.8f}, Python={res_py.att:.8f}", ) @@ -2459,14 +2650,19 @@ def test_bootstrap_seed_reproducibility(self, seed): res_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") trop_global_module = sys.modules["diff_diff.trop_global"] - with patch.object(trop_global_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), \ - patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None): + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_loocv_grid_search_global", None), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + ): trop_py = TROP(**trop_params) res_py = trop_py.fit(df.copy(), "outcome", "treated", "unit", "time") np.testing.assert_allclose( - res_rust.se, res_py.se, atol=1e-14, rtol=1e-14, + res_rust.se, + res_py.se, + atol=1e-14, + rtol=1e-14, err_msg=f"Bootstrap SE divergence under seed={seed}: " f"Rust={res_rust.se:.16f}, Python={res_py.se:.16f}", ) @@ -2522,13 +2718,18 @@ def test_bootstrap_seed_reproducibility_local(self, seed, lambda_nn): res_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") trop_local_module = sys.modules["diff_diff.trop_local"] - with patch.object(trop_local_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_local_module, "_rust_bootstrap_trop_variance", None): + with ( + patch.object(trop_local_module, "HAS_RUST_BACKEND", False), + patch.object(trop_local_module, "_rust_bootstrap_trop_variance", None), + ): trop_py = TROP(**trop_params) res_py = trop_py.fit(df.copy(), "outcome", "treated", "unit", "time") np.testing.assert_allclose( - res_rust.se, res_py.se, atol=1e-5, rtol=1e-5, + res_rust.se, + res_py.se, + atol=1e-5, + rtol=1e-5, err_msg=f"Local-method bootstrap SE divergence under " f"seed={seed}, lambda_nn={lambda_nn}: " f"Rust={res_rust.se:.16f}, Python={res_py.se:.16f}", @@ -2587,15 +2788,20 @@ def test_local_method_main_fit_parity(self, lambda_nn, tol): trop_module = sys.modules["diff_diff.trop"] trop_local_module = sys.modules["diff_diff.trop_local"] - with patch.object(trop_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_module, "_rust_loocv_grid_search", None), \ - patch.object(trop_local_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_local_module, "_rust_bootstrap_trop_variance", None): + with ( + patch.object(trop_module, "HAS_RUST_BACKEND", False), + patch.object(trop_module, "_rust_loocv_grid_search", None), + patch.object(trop_local_module, "HAS_RUST_BACKEND", False), + patch.object(trop_local_module, "_rust_bootstrap_trop_variance", None), + ): trop_py = TROP(**trop_params) res_py = trop_py.fit(df.copy(), "outcome", "treated", "unit", "time") np.testing.assert_allclose( - res_rust.att, res_py.att, atol=tol, rtol=tol, + res_rust.att, + res_py.att, + atol=tol, + rtol=tol, err_msg=f"Local-method ATT divergence at lambda_nn={lambda_nn}: " f"Rust={res_rust.att:.16f}, Python={res_py.att:.16f}", ) @@ -2653,15 +2859,20 @@ def test_local_method_same_cohort_donor_parity(self, lambda_nn, tol): trop_module = sys.modules["diff_diff.trop"] trop_local_module = sys.modules["diff_diff.trop_local"] - with patch.object(trop_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_module, "_rust_loocv_grid_search", None), \ - patch.object(trop_local_module, "HAS_RUST_BACKEND", False), \ - patch.object(trop_local_module, "_rust_bootstrap_trop_variance", None): + with ( + patch.object(trop_module, "HAS_RUST_BACKEND", False), + patch.object(trop_module, "_rust_loocv_grid_search", None), + patch.object(trop_local_module, "HAS_RUST_BACKEND", False), + patch.object(trop_local_module, "_rust_bootstrap_trop_variance", None), + ): trop_py = TROP(**trop_params) res_py = trop_py.fit(df.copy(), "outcome", "treated", "unit", "time") np.testing.assert_allclose( - res_rust.att, res_py.att, atol=tol, rtol=tol, + res_rust.att, + res_py.att, + atol=tol, + rtol=tol, err_msg=f"Same-cohort donor ATT divergence at lambda_nn={lambda_nn}: " f"Rust={res_rust.att:.16f}, Python={res_py.att:.16f}", ) @@ -2699,3 +2910,250 @@ def test_linalg_works_without_rust(self): assert coeffs.shape == (k,) assert residuals.shape == (n,) assert vcov.shape == (k, k) + + +from diff_diff._backend import _rust_demean_map as _demean_map_symbol + + +@pytest.mark.skipif( + not HAS_RUST_BACKEND or _demean_map_symbol is None, + reason="Rust backend or demean_map kernel not available", +) +class TestDemeanMapKernel: + """Rust demean_map vs the canonical numpy engine (_demean_map_numpy). + + Contract: identical sweep order, row-order scatter-add accumulation, and + max|x - x_old| < tol convergence per column (incl. NaN poisoning). The + assertion order is iteration-count EQUALITY first (deterministic under + the pinned op-order contract), then allclose on outputs. + """ + + @staticmethod + def _fixture(kind, seed=0, k=3): + rng = np.random.default_rng(seed) + if kind == "balanced": + n_units, n_periods = 40, 8 + unit = np.repeat(np.arange(n_units), n_periods) + period = np.tile(np.arange(n_periods), n_units) + elif kind == "unbalanced": + n_units, n_periods = 60, 12 + unit = np.repeat(np.arange(n_units), n_periods) + period = np.tile(np.arange(n_periods), n_units) + keep = rng.random(unit.size) > 0.35 + unit, period = unit[keep], period[keep] + elif kind == "contiguous": # slow-convergence regime (>100 iterations) + n_units, n_periods, span = 120, 40, 6 + unit = np.repeat(np.arange(n_units), n_periods) + period = np.tile(np.arange(n_periods), n_units) + entry = rng.integers(0, n_periods - span, n_units) + keep = (period >= entry[unit]) & (period < entry[unit] + span) + unit, period = unit[keep], period[keep] + else: + raise ValueError(kind) + n = unit.size + x_cols = [rng.normal(size=n) for _ in range(k)] + codes_list = [ + pd.factorize(unit, sort=False)[0].astype(np.intp), + pd.factorize(period, sort=False)[0].astype(np.intp), + ] + n_groups = [len(np.unique(unit)), len(np.unique(period))] + w = rng.uniform(0.5, 2.0, n) + return x_cols, codes_list, n_groups, w + + @staticmethod + def _run_both(x_cols, codes_list, n_groups, weights, tol=1e-10, max_iter=10_000): + from diff_diff.utils import _demean_map_numpy, _demean_map_rust + + rust = _demean_map_rust(x_cols, codes_list, n_groups, weights, tol, max_iter) + assert rust is not None, "rust kernel unexpectedly fell back" + numpy_res = _demean_map_numpy(x_cols, codes_list, n_groups, weights, tol, max_iter) + return rust, numpy_res + + @pytest.mark.parametrize("kind", ["balanced", "unbalanced", "contiguous"]) + @pytest.mark.parametrize("weighted", [False, True]) + def test_equivalence_two_way(self, kind, weighted): + x_cols, codes_list, n_groups, w = self._fixture(kind) + (r_cols, r_iters), (p_cols, p_iters) = self._run_both( + x_cols, codes_list, n_groups, w if weighted else None + ) + assert r_iters == p_iters # deterministic under the pinned op order + if kind == "contiguous": + assert all(it > 100 or it < 0 for it in p_iters) or max(p_iters) > 100 + for rc, pc in zip(r_cols, p_cols): + np.testing.assert_allclose(rc, pc, rtol=0, atol=1e-12) + + def test_equivalence_three_way(self): + rng = np.random.default_rng(3) + x_cols, codes_list, n_groups, w = self._fixture("unbalanced", seed=3) + n = x_cols[0].size + firm = rng.integers(0, 7, n) + codes_list = codes_list + [pd.factorize(firm, sort=False)[0].astype(np.intp)] + n_groups = n_groups + [len(np.unique(firm))] + (r_cols, r_iters), (p_cols, p_iters) = self._run_both(x_cols, codes_list, n_groups, w) + assert r_iters == p_iters + for rc, pc in zip(r_cols, p_cols): + np.testing.assert_allclose(rc, pc, rtol=0, atol=1e-12) + + def test_zero_total_weight_group_rows_inert_parity(self): + x_cols, codes_list, n_groups, w = self._fixture("unbalanced", seed=4) + w = w.copy() + zero_rows = codes_list[0] == 0 + w[zero_rows] = 0.0 + (r_cols, r_iters), (p_cols, p_iters) = self._run_both(x_cols, codes_list, n_groups, w) + assert r_iters == p_iters + for rc, pc in zip(r_cols, p_cols): + np.testing.assert_allclose(rc, pc, rtol=0, atol=1e-12) + assert np.isfinite(rc).all() + + def test_nan_in_variable_never_converges_both(self): + x_cols, codes_list, n_groups, _ = self._fixture("unbalanced", seed=5, k=1) + x_cols[0][3] = np.nan + (_, r_iters), (_, p_iters) = self._run_both( + x_cols, codes_list, n_groups, None, tol=1e-8, max_iter=25 + ) + assert r_iters == [-1] + assert p_iters == [-1] + + @pytest.mark.parametrize("k", [1, 64]) + def test_column_counts(self, k): + x_cols, codes_list, n_groups, _ = self._fixture("unbalanced", seed=6, k=k) + (r_cols, r_iters), (p_cols, p_iters) = self._run_both(x_cols, codes_list, n_groups, None) + assert len(r_cols) == k and r_iters == p_iters + for rc, pc in zip(r_cols, p_cols): + np.testing.assert_allclose(rc, pc, rtol=0, atol=1e-12) + + def test_nonconvergence_flag_parity_at_max_iter_1(self): + x_cols, codes_list, n_groups, _ = self._fixture("unbalanced", seed=7) + (_, r_iters), (_, p_iters) = self._run_both( + x_cols, codes_list, n_groups, None, tol=1e-15, max_iter=1 + ) + assert r_iters == p_iters == [-1] * len(x_cols) + + def test_forced_fallback_runs_numpy_engine(self, monkeypatch): + """Wrapper returning None must route demean_by_groups to the numpy + engine and still produce a correct result.""" + import diff_diff.utils as utils_mod + from diff_diff.utils import demean_by_groups + + rng = np.random.default_rng(8) + df = pd.DataFrame( + { + "unit": np.repeat(np.arange(20), 5), + "period": np.tile(np.arange(5), 20), + "y": rng.normal(size=100), + } + ) + calls = {"numpy": 0} + orig_numpy = utils_mod._demean_map_numpy + + def counting_numpy(*a, **kw): + calls["numpy"] += 1 + return orig_numpy(*a, **kw) + + monkeypatch.setattr(utils_mod, "_demean_map_rust", lambda *a, **kw: None) + monkeypatch.setattr(utils_mod, "_demean_map_numpy", counting_numpy) + out, _ = demean_by_groups(df, ["y"], ["unit", "period"], suffix="_dm") + assert calls["numpy"] == 1 + assert np.abs(out["y_dm"].values.mean()) < 1e-12 + + def test_nonconvergence_warning_parity_under_rust(self): + """Same 'did not converge' warning, same variable names, via the + rust dispatch path.""" + from diff_diff.utils import demean_by_groups + + rng = np.random.default_rng(9) + df = pd.DataFrame( + { + "unit": np.repeat(np.arange(30), 6), + "period": np.tile(np.arange(6), 30), + } + ) + df = df[rng.random(len(df)) > 0.3].reset_index(drop=True) + df["y"] = rng.normal(size=len(df)) + with pytest.warns(UserWarning, match=r"\['y'\].*did not converge"): + demean_by_groups(df, ["y"], ["unit", "period"], suffix="_dm", max_iter=1, tol=1e-15) + + def test_estimator_level_att_parity(self, monkeypatch): + """SunAbraham + DiD(absorb=) ATT/SE identical across engines, + including the FE-spanned-covariate snap decision.""" + import warnings as _w + + import diff_diff.utils as utils_mod + from diff_diff import DifferenceInDifferences, SunAbraham + + rng = np.random.default_rng(10) + n_units, n_periods = 90, 10 + unit = np.repeat(np.arange(n_units), n_periods) + time_ = np.tile(np.arange(n_periods), n_units) + keep = rng.random(unit.size) > 0.25 + unit, time_ = unit[keep], time_[keep] + first = np.where(np.arange(n_units) % 3 == 0, 0, 5)[unit] + treated = (unit < n_units // 2).astype(int) + post = (time_ >= n_periods // 2).astype(int) + y = 0.3 * treated * post + rng.normal(size=unit.size) + df = pd.DataFrame( + { + "y": y, + "unit": unit, + "time": time_, + "first_treat": first, + "treated": treated, + "post": post, + } + ) + # FE-spanned covariate: snap decisions must match across engines + a = rng.normal(size=n_units) + b = rng.normal(size=n_periods) + df["xspan"] = a[unit] + b[time_] + + def fits(): + with _w.catch_warnings(): + _w.simplefilter("ignore") + sa = SunAbraham().fit( + df, outcome="y", unit="unit", time="time", first_treat="first_treat" + ) + did = DifferenceInDifferences().fit( + df, + outcome="y", + treatment="treated", + time="post", + absorb=["unit", "time"], + covariates=["xspan"], + ) + return sa, did + + sa_r, did_r = fits() + monkeypatch.setattr(utils_mod, "_rust_demean_map", None) + sa_p, did_p = fits() + assert sa_r.att == pytest.approx(sa_p.att, abs=1e-10) + assert sa_r.se == pytest.approx(sa_p.se, rel=1e-8) + assert did_r.att == pytest.approx(did_p.att, abs=1e-10) + assert did_r.se == pytest.approx(did_p.se, rel=1e-8) + # snap decision parity: spanned covariate NaN under BOTH engines + assert np.isnan(did_r.coefficients["xspan"]) + assert np.isnan(did_p.coefficients["xspan"]) + + def test_stale_symbol_none_falls_back_to_numpy(self, monkeypatch): + """A mixed-version extension missing demean_map (symbol None) must + route the PUBLIC entry point to the numpy engine, not raise.""" + import diff_diff.utils as utils_mod + from diff_diff.utils import demean_by_groups + + rng = np.random.default_rng(11) + df = pd.DataFrame( + { + "unit": np.repeat(np.arange(15), 4), + "period": np.tile(np.arange(4), 15), + "y": rng.normal(size=60), + } + ) + monkeypatch.setattr(utils_mod, "_rust_demean_map", None) + out, _ = demean_by_groups(df, ["y"], ["unit", "period"], suffix="_dm") + assert np.abs(out["y_dm"].values.mean()) < 1e-12 + # the wrapper itself honors its documented None contract too + assert ( + utils_mod._demean_map_rust( + [df["y"].values], [np.zeros(60, dtype=np.intp)], [1], None, 1e-10, 10 + ) + is None + )