From 8067f140e058d71f6a5a89d24e43b61b20e6db1d Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 24 Apr 2026 08:38:41 -0400 Subject: [PATCH 1/6] Fix TROP local-method backend parity: drop Rust weight normalization + Python cache-fallthrough MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the local-method half of silent-failures audit finding #23 (RNG half closed in PR #354; grid-search half in PR #348). Two methodology fixes, both isolated to the local-method path — global is unaffected. 1. Rust weight-matrix normalization removed ------------------------------------------ `rust/src/trop.rs::compute_weight_matrix` no longer divides `time_weights` and `unit_weights` by their respective sums before the outer product. The paper's Equation 2/3 (Athey, Imbens, Qu, Viviano 2025) and REGISTRY.md Requirements checklist (`[x] Unit weights: exp(-λ_unit × distance) (unnormalized, matching Eq. 2)`) both specify raw-exponential weights; Python's `_compute_observation_weights` was already REGISTRY-compliant. Rust's normalization inflated the effective nuclear-norm penalty relative to the data-fit term, changing the regularization trade-off. User-visible effect: Rust local-method ATT values may shift for fits with `lambda_nn < infinity`. For `lambda_nn = infinity` (factor model disabled) outputs are unchanged — uniform weight scaling leaves the minimum-norm WLS argmin invariant. Rust LOOCV-selected lambdas may also shift on that boundary; both backends now converge on the same selection. Affects both local-method Rust call sites (LOOCV at trop.rs:459, bootstrap at trop.rs:1096). 2. Python `_compute_observation_weights` cache-fallthrough removed --------------------------------------------------------------- Removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`). Test changes ------------ - `tests/test_rust_backend.py::TestTROPRustEdgeCaseParity:: test_bootstrap_seed_reproducibility_local`: flipped from `xfail(strict=True)` to passing `assert_allclose` at `atol=1e-5` across seeds `[0, 42, 12345]`. Residual ~1e-7 gap is Rust `estimate_model` vs numpy `lstsq` roundoff that accumulates differently across per-replicate bootstrap fits; follow-up TODO row tracks unifying Rust to the `solve_wls_svd` path (same SVD helper the global-method uses since PR #348) for sub-1e-14 parity. - New `test_local_method_main_fit_parity`: parametrized over `(lambda_nn=inf, atol=1e-14)` and `(lambda_nn=0.1, atol=1e-10)`; asserts `atol=1e-14` bit-identity for the main-fit ATT at `lambda_nn=inf` (the regression guard for the normalization fix) and `atol=1e-10` for the finite-`lambda_nn` FISTA path. Verification ------------ Targeted regression sweep — all green: - 9 `TestTROPRustEdgeCaseParity` tests (grid-search + global bootstrap × 3 seeds + local bootstrap × 3 seeds + local main-fit × 2 regimes) - Full `test_rust_backend.py` suite: 92 passed - Full `test_trop.py` suite under Rust backend: 120 passed Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 2 + TODO.md | 4 +- diff_diff/trop_local.py | 44 +++--------------- rust/src/trop.rs | 21 +++------ tests/test_rust_backend.py | 92 +++++++++++++++++++++++++++++--------- 5 files changed, 85 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e40e919..15f3f876 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - SyntheticDiD bootstrap SE formula applies the `sqrt((r-1)/r)` correction matching R's synthdid and the placebo SE formula. - SyntheticDiD bootstrap now retries degenerate resamples (all-control or all-treated, or non-finite `τ_b`) until exactly `n_bootstrap` valid replicates are accumulated, matching R's `synthdid::bootstrap_sample` and Arkhangelsky et al. (2021) Algorithm 2. Previously the Python path counted attempts (with degenerate draws silently dropped), producing fewer valid replicates than requested. A bounded-attempt guard (`20 × n_bootstrap`) prevents pathological-input hangs. - **TROP global bootstrap SE backend parity under fixed seed** — Rust and Python backends now produce bit-identical bootstrap SE under the same `seed`. Previously Rust's `bootstrap_trop_variance_global` seeded `rand_xoshiro::Xoshiro256PlusPlus` per replicate while Python's fallback consumed `numpy.random.default_rng` (PCG64), producing ~28% SE divergence on tiny panels under `seed=42`. Fixed by extracting a shared `stratified_bootstrap_indices` helper in `diff_diff/bootstrap_utils.py` that pre-generates per-replicate stratified sample indices via numpy on the Python side; both backends consume the same integer arrays through the PyO3 surface. Sampling law (stratified: controls then treated, with replacement) is unchanged. Closes silent-failures audit finding #23 (bootstrap half; grid-search half closed in PR #348). Local-method TROP also adopts the Python-canonical index contract for the RNG layer, but separately-discovered backend divergences (Rust normalizes weight-matrix outer product, Python `_compute_observation_weights` reads stale `_precomputed` cache) prevent local-method bit-identity SE; tracked as a follow-up in `TODO.md`. +- **TROP local-method Rust weight-matrix no longer normalized** — `rust/src/trop.rs::compute_weight_matrix` no longer divides time-weights or unit-weights by their respective sums before the outer product. The paper's Equation 2/3 (Athey, Imbens, Qu, Viviano 2025) and REGISTRY.md Requirements checklist (line 2037: `[x] Unit weights: exp(-λ_unit × distance) (unnormalized, matching Eq. 2)`) both specify raw-exponential weights; Python's `_compute_observation_weights` was already REGISTRY-compliant. **User-visible effect**: Rust local-method ATT values may shift for any fit with `lambda_nn < infinity` — normalizing the weight-matrix inflated the effective nuclear-norm penalty relative to the data-fit term, changing the regularization trade-off. For `lambda_nn = infinity` (factor model disabled) outputs are unchanged because uniform weight scaling leaves the minimum-norm WLS argmin invariant. Rust LOOCV-selected lambdas may also shift on this boundary; both backends now converge on the same REGISTRY-compliant selection. +- **TROP local-method Python `_compute_observation_weights` no longer reads stale `_precomputed` cache** — removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`). Together with the normalization fix above, TROP local-method backends are now bit-identical for main-fit ATT (`atol=rtol=1e-14` regression guard in `test_local_method_main_fit_parity`); bootstrap SE parity is asserted at `atol=1e-5` to accommodate ~1e-7 roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq` (sub-1e-14 bootstrap parity tracked as a follow-up in `TODO.md` under "unify Rust local-method solver path"). Closes silent-failures audit finding #23 (local-method half; the RNG half closed in PR #354 and the grid-search half in PR #348). ### Changed - **`did_had_pretest_workflow(aggregate="event_study")` verdict no longer emits the "paper step 2 deferred to Phase 3 follow-up" caveat** — the joint pre-trends Stute test closes that gap. The two-period `aggregate="overall"` path retains the existing caveat since the joint variant does not apply to single-pre-period panels. Downstream code that greps verdict strings for the Phase 3 caveat will see it suppressed on the event-study path. diff --git a/TODO.md b/TODO.md index 0f15a8c6..4a5c86c7 100644 --- a/TODO.md +++ b/TODO.md @@ -83,8 +83,8 @@ Deferred items from PR reviews that were not addressed before merge. | Weighted CR2 Bell-McCaffrey cluster-robust (`vcov_type="hc2_bm"` + `cluster_ids` + `weights`) currently raises `NotImplementedError`. Weighted hat matrix and residual rebalancing need threading per clubSandwich WLS handling. | `linalg.py::_compute_cr2_bm` | Phase 1a | Medium | | Regenerate `benchmarks/data/clubsandwich_cr2_golden.json` from R (`Rscript benchmarks/R/generate_clubsandwich_golden.R`). Current JSON has `source: python_self_reference` as a stability anchor until an authoritative R run. | `benchmarks/R/generate_clubsandwich_golden.R` | Phase 1a | Medium | | `honest_did.py:1907` `np.linalg.solve(A_sys, b_sys) / except LinAlgError: continue` is a silent basis-rejection in the vertex-enumeration loop that is algorithmically intentional (try the next basis). Consider surfacing a count of rejected bases as a diagnostic when ARP enumeration exhausts, so users see when the vertex search was heavily constrained. Not a silent failure in the sense of the Phase 2 audit (the algorithm is supposed to skip), but the diagnostic would help debug borderline cases. | `honest_did.py` | #334 | Low | -| TROP local-method Rust vs Python bootstrap SE divergence beyond RNG: with stratified indices now Python-canonical (closes RNG axis-H), local-method SE still diverges by 10-20% on tiny panels. Two downstream causes: (a) Rust `compute_weight_matrix` (`rust/src/trop.rs:573-606`) normalizes time_weights and unit_weights to sum to 1 before the outer product; Python `_compute_observation_weights` (`diff_diff/trop_local.py:489`) does not normalize. (b) Python `_compute_observation_weights` reads `self._precomputed["Y"]`, `["D"]`, `["time_dist_matrix"]` (lines 423-458) — original-panel cache — rather than the bootstrap-sample data passed via the function arguments, so unit-distance computation in the Python fallback uses stale data. `@pytest.mark.xfail(strict=True)` in `tests/test_rust_backend.py::TestTROPRustEdgeCaseParity::test_bootstrap_seed_reproducibility_local` baselines the gap across `seeds=[0, 42, 12345]`. Affects local-method TROP backend parity broadly, not just bootstrap. | `trop_local.py`, `rust/src/trop.rs` | follow-up | Medium | -| Rust multiplier-bootstrap weight RNG (`generate_bootstrap_weights_batch` in `rust/src/bootstrap.rs:9-10, 57-75`) uses `Xoshiro256PlusPlus::seed_from_u64(seed + i)` per row for Rademacher/Mammen/Webb generation. If any Python caller (SDID / efficient-DiD multiplier bootstrap) has a numpy-canonical equivalent, the two backends likely diverge under the same seed. Audit Python callers (`diff_diff/sdid.py`, `diff_diff/efficient_did_bootstrap.py`, `diff_diff/bootstrap_utils.py::generate_bootstrap_weights_batch_numpy`) for parity-test gaps. Same fix shape as TROP RNG parity (PR #NNN): pre-generate weights in Python via numpy and pass them to Rust through PyO3. | `rust/src/bootstrap.rs`, `diff_diff/bootstrap_utils.py` | follow-up | Medium | +| Unify Rust local-method `estimate_model` solver path to `solve_wls_svd` (the same SVD helper used by the global-method since PR #348) for sub-1e-14 bootstrap SE parity. Current local-method bootstrap parity test (`tests/test_rust_backend.py::TestTROPRustEdgeCaseParity::test_bootstrap_seed_reproducibility_local`) passes at `atol=1e-5` — the residual ~1e-7 gap is roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq`, which accumulates differently across per-replicate bootstrap fits. Main-fit ATT is bit-identical at `atol=1e-14` (see `test_local_method_main_fit_parity`), so this is a bootstrap-aggregation-only concern; not a user-visible correctness bug. | `rust/src/trop.rs::estimate_model`, `rust/src/linalg.rs::solve_wls_svd` | follow-up | Low | +| Rust multiplier-bootstrap weight RNG (`generate_bootstrap_weights_batch` in `rust/src/bootstrap.rs:9-10, 57-75`) uses `Xoshiro256PlusPlus::seed_from_u64(seed + i)` per row for Rademacher/Mammen/Webb generation. If any Python caller (SDID / efficient-DiD multiplier bootstrap) has a numpy-canonical equivalent, the two backends likely diverge under the same seed. Audit Python callers (`diff_diff/sdid.py`, `diff_diff/efficient_did_bootstrap.py`, `diff_diff/bootstrap_utils.py::generate_bootstrap_weights_batch_numpy`) for parity-test gaps. Same fix shape as TROP RNG parity (PR #354): pre-generate weights in Python via numpy and pass them to Rust through PyO3. | `rust/src/bootstrap.rs`, `diff_diff/bootstrap_utils.py` | follow-up | Medium | | `bias_corrected_local_linear`: extend golden parity to `kernel="triangular"` and `kernel="uniform"` (currently epa-only; all three kernels share `kernel_W` and the `lprobust` math, so parity is expected but not separately asserted). | `benchmarks/R/generate_nprobust_lprobust_golden.R`, `tests/test_bias_corrected_lprobust.py` | Phase 1c | Low | | `bias_corrected_local_linear`: expose `vce in {"hc0", "hc1", "hc2", "hc3"}` on the public wrapper once R parity goldens exist (currently raises `NotImplementedError`). The port-level `lprobust` and `lprobust_res` already support all four; expanding the public surface requires a golden generator for each hc mode and a decision on hc2/hc3 q-fit leverage (R reuses p-fit `hii` for q-fit residuals; whether to match that or stage-match deserves a derivation before the wrapper advertises CCT-2014 conformance). | `diff_diff/local_linear.py::bias_corrected_local_linear`, `benchmarks/R/generate_nprobust_lprobust_golden.R`, `tests/test_bias_corrected_lprobust.py` | Phase 1c | Medium | | `bias_corrected_local_linear`: support `weights=` once survey-design adaptation lands. nprobust's `lprobust` has no weight argument so there is no parity anchor; derivation needed. | `diff_diff/local_linear.py`, `diff_diff/_nprobust_port.py::lprobust` | Phase 1c | Medium | diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index 39cce9fc..555f8650 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -391,7 +391,11 @@ def _compute_observation_weights( weights for ALL units where D[t, j] = 0 at the target period, not just never-treated units. - Uses pre-computed structures when available for efficiency. + Always computes from the function-argument ``Y, D``; does not read + ``self._precomputed``. Under bootstrap the caller passes resampled + ``Y, D``, and a prior version of this method silently fell through to + the original-panel cache via a ``_precomputed`` branch, producing + stale unit distances. Parameters ---------- @@ -420,44 +424,6 @@ def _compute_observation_weights( np.ndarray Weight matrix (n_periods x n_units) for observation (i, t). """ - # Use pre-computed structures when available - if self._precomputed is not None: - # Time weights from pre-computed time distance matrix - # time_dist_matrix[t, s] = |t - s| - time_weights = np.exp(-lambda_time * self._precomputed["time_dist_matrix"][t, :]) - - # Unit weights - computed for ALL units where D[t, j] = 0 - # (Issue A fix: includes pre-treatment obs of eventually-treated units) - unit_weights = np.zeros(n_units) - D_stored = self._precomputed["D"] - Y_stored = self._precomputed["Y"] - - # Valid control units at time t: D[t, j] == 0 - valid_control_at_t = D_stored[t, :] == 0 - - if lambda_unit == 0: - # Uniform weights when lambda_unit = 0 - # All units not treated at time t get weight 1 - unit_weights[valid_control_at_t] = 1.0 - else: - # Use observation-specific distances with target period excluded - # (Issue B fix: compute exact per-observation distance) - for j in range(n_units): - if valid_control_at_t[j] and j != i: - # Compute distance excluding target period t - dist = self._compute_unit_distance_for_obs(Y_stored, D_stored, j, i, t) - if np.isinf(dist): - unit_weights[j] = 0.0 - else: - unit_weights[j] = np.exp(-lambda_unit * dist) - - # Treated unit i gets weight 1 - unit_weights[i] = 1.0 - - # Weight matrix: outer product (n_periods x n_units) - return np.outer(time_weights, unit_weights) - - # Fallback: compute from scratch (used in bootstrap) # Time distance: |t - s| following paper's Equation 3 (page 7) dist_time = np.abs(np.arange(n_periods) - t) time_weights = np.exp(-lambda_time * dist_time) diff --git a/rust/src/trop.rs b/rust/src/trop.rs index ac98b44f..706f3958 100644 --- a/rust/src/trop.rs +++ b/rust/src/trop.rs @@ -565,19 +565,15 @@ fn compute_weight_matrix( time_dist: &ArrayView2, ) -> Array2 { // Time weights for this target period: θ_s = exp(-λ_time × |t - s|) - let mut time_weights: Array1 = Array1::from_shape_fn(n_periods, |s| { + // Unnormalized per REGISTRY Eq. 2/3. + let time_weights: Array1 = Array1::from_shape_fn(n_periods, |s| { let dist = time_dist[[target_period, s]] as f64; (-lambda_time * dist).exp() }); - // Normalize time weights to sum to 1 - let time_sum: f64 = time_weights.sum(); - if time_sum > 0.0 { - time_weights /= time_sum; - } - // Unit weights: ω_j = exp(-λ_unit × dist(j, i)) - // Paper alignment: compute for ALL units, let control masking handle exclusion + // Paper alignment: compute for ALL units, let control masking handle exclusion. + // Unnormalized per REGISTRY Eq. 2/3. let mut unit_weights = Array1::::zeros(n_units); if lambda_unit == 0.0 { @@ -600,14 +596,7 @@ fn compute_weight_matrix( // Target unit gets weight 1 (will be masked out in estimation anyway) unit_weights[target_unit] = 1.0; - // Normalize unit weights to sum to 1 - let unit_sum: f64 = unit_weights.sum(); - if unit_sum > 0.0 { - unit_weights /= unit_sum; - } - - // Outer product: W[t, i] = time_weights[t] * unit_weights[i] - // Result is normalized since both components sum to 1 + // Outer product: W[t, i] = θ_s × ω_j (raw exponentials, unnormalized) let mut weight_matrix = Array2::::zeros((n_periods, n_units)); for t in 0..n_periods { for i in 0..n_units { diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 1ecc8af8..09ff503d 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -2420,28 +2420,25 @@ def test_bootstrap_seed_reproducibility(self, seed): f"Rust={res_rust.se:.16f}, Python={res_py.se:.16f}", ) - @pytest.mark.xfail( - strict=True, - reason="Local-method TROP has additional Rust-vs-Python backend " - "divergences beyond RNG (weight-matrix normalization in " - "`compute_weight_matrix` is Rust-only; Python " - "`_compute_observation_weights` reads `self._precomputed['Y']` " - "instead of the bootstrap-sample `Y`). RNG parity infrastructure " - "is in place (both backends now consume the same Python-canonical " - "stratified indices), but post-resampling computations diverge so " - "SE is not yet bit-identical. Tracked as a separate finding in " - "TODO.md. This xfail baselines the gap so we notice when both " - "downstream divergences are closed.", - ) @pytest.mark.parametrize("seed", [0, 42, 12345]) def test_bootstrap_seed_reproducibility_local(self, seed): - """Bootstrap SE parity under a fixed seed (local method). - - Currently xfail(strict=True) — see decorator. The local-method - path has additional backend divergences beyond RNG that this PR - does not address. Companion to the global-method parity test - above; once the downstream divergences close, this test will - XPASS and should be promoted to a regression guard. + """Backend-invariant bootstrap SE parity for the local method. + + Post-methodology-alignment regression guard. With the Rust weight- + matrix normalization dropped and the Python ``_precomputed`` + cache-fallthrough removed, local-method Rust and Python bootstraps + consume bit-identical stratified indices AND bit-identical raw- + exponential weights. Main-fit ATT is bit-identical across backends + (see ``test_local_method_main_fit_parity``), but per-replicate + bootstrap fits route through Rust's ``estimate_model`` vs numpy's + ``lstsq``, which use different matrix factorization paths and + accumulate different BLAS roundoff. Empirically the residual gap + is ~1e-7 relative; asserted at ``atol=1e-5`` which is ~100x the + observed gap and comfortable across CI runner variance. + + Follow-up to tighten to ``atol=1e-14``: unify Rust ``estimate_model`` + to use ``solve_wls_svd`` (the same SVD path used by global-method + since PR #348). Tracked in ``TODO.md``. """ import sys from unittest.mock import patch @@ -2468,11 +2465,64 @@ def test_bootstrap_seed_reproducibility_local(self, seed): 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-5, rtol=1e-5, err_msg=f"Local-method bootstrap SE divergence under seed={seed}: " f"Rust={res_rust.se:.16f}, Python={res_py.se:.16f}", ) + @pytest.mark.parametrize( + "lambda_nn,tol", + [ + (np.inf, 1e-14), + (0.1, 1e-10), + ], + ) + def test_local_method_main_fit_parity(self, lambda_nn, tol): + """Backend-invariant ATT parity for the local-method main fit. + + Companion to the bootstrap seed-parity test above. Exercises both + regimes: ``lambda_nn=inf`` (no-lowrank, bit-identical minimum-norm + WLS argmin under aligned raw-exponential weights) and a finite + ``lambda_nn`` (with-lowrank, FISTA inner loop; tolerance relaxed + to ``1e-10`` because FISTA iteration ordering and BLAS reduction + ordering introduce sub-1e-10 noise across Rust faer and numpy BLAS + paths). + + Regression guard for the normalization fix and cache-fallthrough + fix landed in this PR. Before the fix, Rust ATT diverged from + Python ATT by O(10%) at finite ``lambda_nn`` and O(0) at + ``lambda_nn=inf``; after the fix both regimes match to tolerance. + """ + import sys + from unittest.mock import patch + + from diff_diff import TROP + + df = self._make_correlated_panel(n_units=6, n_periods=6, n_treated=2) + trop_params = dict( + method="local", + lambda_time_grid=[1.0], + lambda_unit_grid=[1.0], + lambda_nn_grid=[lambda_nn], + n_bootstrap=2, # minimum allowed; we assert ATT, not SE + seed=42, + ) + + trop_rust = TROP(**trop_params) + 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): + 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, + err_msg=f"Local-method ATT divergence at lambda_nn={lambda_nn}: " + f"Rust={res_rust.att:.16f}, Python={res_py.att:.16f}", + ) + class TestFallbackWhenNoRust: """Test that pure Python fallback works when Rust is unavailable.""" From f3000331525ac26c554884c738b1172866b13666 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 24 Apr 2026 09:16:16 -0400 Subject: [PATCH 2/6] Address PR #358 AI review: remove Python same-period donor gate + regression MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the second P1 from the review. Python `_compute_observation_weights` had an extra `valid_control_at_t = D[t, :] == 0` gate that zeroed ω_j for units treated at the target period (other than the target unit itself). Rust's `compute_weight_matrix` has no such gate — per the paper's Eq. 2/3 and `docs/methodology/REGISTRY.md` TROP section, `ω_j = exp(-λ_unit × dist(j, i))` is distance-based for all `j ≠ i` and the treated-cell exclusion is the `(1 - W_{js})` factor applied inside `_estimate_model` via the control mask, not an extra target-period unit-weight gate. The empirical impact of removing the gate is zero on the ATT point estimate: same-cohort donors' pre-treatment rows are exactly absorbed by their own unit fixed effect `alpha_j` without propagating into `mu`, `beta`, or other units' parameters — adding them to the fit changes which rows are scored but not the solution the fit converges to. Verified: the flipped bootstrap-seed parity test, the main-fit parity test at `lambda_nn=inf` (`atol=1e-14`) and at `lambda_nn=0.1` (`atol=1e-10`), and the new same-cohort regression test (below) all pass before and after the gate removal. The change is structural alignment with the paper and Rust, not a numerical behavior shift. Test addition ------------- `TestTROPRustEdgeCaseParity::test_local_method_same_cohort_donor_parity` isolates the scenario the gate used to handle differently from Rust: a fixture with three treated units sharing one cohort (all treated at `t=5`) and three controls. Before the gate was removed, Python's and Rust's same-target-period donors diverged in which rows contributed to the fit; the tests prove the ATT point estimate was never affected (pre-treatment rows absorbed by `alpha_j`), and now both backends also agree structurally. Parametrized over the same regime split as the main-fit parity test (`lambda_nn=inf` → `atol=1e-14`, `lambda_nn=0.1` → `atol=1e-10`). Note on the other P1 in the review (HAD rollback claim): that finding was a phantom caused by a stale branch base — PR #353 (HAD joint Stute pretest) landed on `origin/main` between this branch's cut and the review run, so the PR diff against current `origin/main` appeared to "delete" the PR #353 additions. Resolved by rebasing onto the updated `origin/main` before this push. Co-Authored-By: Claude Opus 4.7 (1M context) --- diff_diff/trop_local.py | 21 +++++++------ tests/test_rust_backend.py | 63 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 10 deletions(-) diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index 555f8650..d1fdfd36 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -428,19 +428,20 @@ def _compute_observation_weights( dist_time = np.abs(np.arange(n_periods) - t) time_weights = np.exp(-lambda_time * dist_time) - # Unit weights - computed for ALL units where D[t, j] = 0 - # (Issue A fix: includes pre-treatment obs of eventually-treated units) + # Unit weights ω_j = exp(-λ_unit × dist(j, i)) for all j ≠ i per Eq. 2/3. + # No target-period gate: same-cohort donors enter with distance-based + # weight (their pre-treatment rows contribute via theta_s * omega_j; + # their post-treatment cells are zeroed by the control mask (1-D_{js}) + # applied inside ``_estimate_model``). Matches Rust's compute_weight_matrix. unit_weights = np.zeros(n_units) - # Valid control units at time t: D[t, j] == 0 - valid_control_at_t = D[t, :] == 0 - if lambda_unit == 0: - # Uniform weights when lambda_unit = 0 - unit_weights[valid_control_at_t] = 1.0 + # Uniform weights when lambda_unit = 0 — all units get 1. + # Control masking in _estimate_model handles treated-cell exclusion. + unit_weights[:] = 1.0 else: for j in range(n_units): - if valid_control_at_t[j] and j != i: + if j != i: # Compute distance excluding target period t (Issue B fix) dist = self._compute_unit_distance_for_obs(Y, D, j, i, t) if np.isinf(dist): @@ -448,8 +449,8 @@ def _compute_observation_weights( else: unit_weights[j] = np.exp(-lambda_unit * dist) - # Treated unit i gets weight 1 (or could be omitted since we fit on controls) - # We include treated unit's own observation for model fitting + # Target unit gets weight 1 (will be masked out in estimation via + # the control mask, matching the paper's (1-W_{js}) factor). unit_weights[i] = 1.0 # Weight matrix: outer product (n_periods x n_units) diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 09ff503d..59be799d 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -2523,6 +2523,69 @@ def test_local_method_main_fit_parity(self, lambda_nn, tol): f"Rust={res_rust.att:.16f}, Python={res_py.att:.16f}", ) + @pytest.mark.parametrize("lambda_nn,tol", [(np.inf, 1e-14), (0.1, 1e-10)]) + def test_local_method_same_cohort_donor_parity(self, lambda_nn, tol): + """Backend-invariant ATT when multiple units share the same treatment cohort. + + Isolates the ``D[t, j] == 1`` target-period case the prior ``_compute_ + observation_weights`` gate silently dropped: three treated units all + starting at ``t=3``. Under the paper's Eq. 2/3, ``ω_j`` is + distance-based for all ``j ≠ i`` (same-cohort donors included); their + pre-treatment rows contribute via ``θ_s · ω_j`` and post-treatment + cells are zeroed by the control mask ``(1 - W_{js})``. Python now + matches this convention (gate removed). This regression asserts that + the main-fit ATT is bit-identical across backends on a fixture where + the gate would previously have excluded donors' pre-treatment + information. + """ + import sys + from unittest.mock import patch + + import pandas as pd + + from diff_diff import TROP + + # 3 treated units + 3 controls, all treated units share cohort at t=3. + # Pre-treatment trajectories are distinct so same-cohort donors carry + # non-trivial information; without the fix Python and Rust would have + # different effective donor pools at each treated-observation target. + rng = np.random.default_rng(7) + rows = [] + for i in range(6): + is_treated = i < 3 + base = rng.normal(0, 1, 8) + for t in range(8): + y = 3.0 + i * 0.2 + 0.5 * t + base[t] + treated = 1 if (is_treated and t >= 5) else 0 + if treated: + y += 1.5 + rows.append({"unit": i, "time": t, "outcome": y, "treated": treated}) + df = pd.DataFrame(rows) + + trop_params = dict( + method="local", + lambda_time_grid=[1.0], + lambda_unit_grid=[1.0], + lambda_nn_grid=[lambda_nn], + n_bootstrap=2, + seed=42, + ) + + trop_rust = TROP(**trop_params) + 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): + 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, + err_msg=f"Same-cohort donor ATT divergence at lambda_nn={lambda_nn}: " + f"Rust={res_rust.att:.16f}, Python={res_py.att:.16f}", + ) + class TestFallbackWhenNoRust: """Test that pure Python fallback works when Rust is unavailable.""" From 962d5d58b4f740f6a7bf764cc696a76fb607394b Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 24 Apr 2026 09:26:04 -0400 Subject: [PATCH 3/6] Address PR #358 AI review round 2: doc consistency + finite-lambda_nn bootstrap case MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P3 — update stale comments/docstrings/changelog that still described the old pre-fix behavior: - `diff_diff/trop_local.py` `_compute_observation_weights` docstring (L388-L392) rewritten: weights are assigned for every `j != i` with distance-based ω_j; treated-cell exclusion happens via the `(1-W_{js})` factor inside `_estimate_model` via the control mask, not via a target-period donor gate. - `rust/src/trop.rs` `compute_weight_matrix` paper-alignment comments (L550-L554) rewritten: weights are unnormalized raw exponentials per REGISTRY Eq. 2/3; ALL donor units with `j != target_unit` get distance-based weights; treated-cell exclusion happens via the control mask inside `estimate_model`. - `CHANGELOG.md` local-method parity statement rewritten: main-fit ATT parity is `atol=rtol=1e-14` for `lambda_nn=inf` and `atol=1e-10` for finite `lambda_nn`, not uniformly "bit-identical". P2 — extend `test_bootstrap_seed_reproducibility_local` to also cover finite `lambda_nn=0.1`. Previously parametrized only on seed with `lambda_nn=inf`; the finite regime is where the Rust weight- normalization change has its biggest numerical effect, so adding one `(seed=42, lambda_nn=0.1)` case closes the coverage gap the reviewer flagged. Assertion tolerance unchanged (`atol=1e-5`). All 4 parametrized cases pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 2 +- diff_diff/trop_local.py | 11 ++++---- rust/src/trop.rs | 6 +++-- tests/test_rust_backend.py | 51 ++++++++++++++++++++++++-------------- 4 files changed, 43 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 15f3f876..fd308a9b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,7 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - SyntheticDiD bootstrap now retries degenerate resamples (all-control or all-treated, or non-finite `τ_b`) until exactly `n_bootstrap` valid replicates are accumulated, matching R's `synthdid::bootstrap_sample` and Arkhangelsky et al. (2021) Algorithm 2. Previously the Python path counted attempts (with degenerate draws silently dropped), producing fewer valid replicates than requested. A bounded-attempt guard (`20 × n_bootstrap`) prevents pathological-input hangs. - **TROP global bootstrap SE backend parity under fixed seed** — Rust and Python backends now produce bit-identical bootstrap SE under the same `seed`. Previously Rust's `bootstrap_trop_variance_global` seeded `rand_xoshiro::Xoshiro256PlusPlus` per replicate while Python's fallback consumed `numpy.random.default_rng` (PCG64), producing ~28% SE divergence on tiny panels under `seed=42`. Fixed by extracting a shared `stratified_bootstrap_indices` helper in `diff_diff/bootstrap_utils.py` that pre-generates per-replicate stratified sample indices via numpy on the Python side; both backends consume the same integer arrays through the PyO3 surface. Sampling law (stratified: controls then treated, with replacement) is unchanged. Closes silent-failures audit finding #23 (bootstrap half; grid-search half closed in PR #348). Local-method TROP also adopts the Python-canonical index contract for the RNG layer, but separately-discovered backend divergences (Rust normalizes weight-matrix outer product, Python `_compute_observation_weights` reads stale `_precomputed` cache) prevent local-method bit-identity SE; tracked as a follow-up in `TODO.md`. - **TROP local-method Rust weight-matrix no longer normalized** — `rust/src/trop.rs::compute_weight_matrix` no longer divides time-weights or unit-weights by their respective sums before the outer product. The paper's Equation 2/3 (Athey, Imbens, Qu, Viviano 2025) and REGISTRY.md Requirements checklist (line 2037: `[x] Unit weights: exp(-λ_unit × distance) (unnormalized, matching Eq. 2)`) both specify raw-exponential weights; Python's `_compute_observation_weights` was already REGISTRY-compliant. **User-visible effect**: Rust local-method ATT values may shift for any fit with `lambda_nn < infinity` — normalizing the weight-matrix inflated the effective nuclear-norm penalty relative to the data-fit term, changing the regularization trade-off. For `lambda_nn = infinity` (factor model disabled) outputs are unchanged because uniform weight scaling leaves the minimum-norm WLS argmin invariant. Rust LOOCV-selected lambdas may also shift on this boundary; both backends now converge on the same REGISTRY-compliant selection. -- **TROP local-method Python `_compute_observation_weights` no longer reads stale `_precomputed` cache** — removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`). Together with the normalization fix above, TROP local-method backends are now bit-identical for main-fit ATT (`atol=rtol=1e-14` regression guard in `test_local_method_main_fit_parity`); bootstrap SE parity is asserted at `atol=1e-5` to accommodate ~1e-7 roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq` (sub-1e-14 bootstrap parity tracked as a follow-up in `TODO.md` under "unify Rust local-method solver path"). Closes silent-failures audit finding #23 (local-method half; the RNG half closed in PR #354 and the grid-search half in PR #348). +- **TROP local-method Python `_compute_observation_weights` no longer reads stale `_precomputed` cache** — removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`). Together with the normalization fix above, TROP local-method backend parity on the main-fit ATT is regime-dependent: `atol=rtol=1e-14` for `lambda_nn=inf` (no nuclear-norm regularization, uniform weight scaling leaves the WLS argmin invariant) and `atol=1e-10` for finite `lambda_nn` (FISTA inner loop + BLAS reduction ordering introduce sub-1e-10 roundoff across Rust `faer` vs numpy paths). Bootstrap SE parity is asserted at `atol=1e-5` to accommodate ~1e-7 roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq` that accumulates across per-replicate fits; sub-1e-14 bootstrap parity is tracked as a follow-up in `TODO.md` under "unify Rust local-method solver path". Closes silent-failures audit finding #23 (local-method half; the RNG half closed in PR #354 and the grid-search half in PR #348). ### Changed - **`did_had_pretest_workflow(aggregate="event_study")` verdict no longer emits the "paper step 2 deferred to Phase 3 follow-up" caveat** — the joint pre-trends Stute test closes that gap. The two-period `aggregate="overall"` path retains the existing caveat since the joint variant does not apply to single-pre-period panels. Downstream code that greps verdict strings for the Phase 3 caveat will see it suppressed on the event-study path. diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index d1fdfd36..231e313e 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -385,11 +385,12 @@ def _compute_observation_weights( - Time weights theta_s^{i,t} = exp(-lambda_time * |t - s|) - Unit weights omega_j^{i,t} = exp(-lambda_unit * dist_unit_{-t}(j, i)) - IMPORTANT (Issue A fix): The paper's objective sums over ALL observations - where (1 - W_js) is non-zero, which includes pre-treatment observations of - eventually-treated units since W_js = 0 for those. This method computes - weights for ALL units where D[t, j] = 0 at the target period, not just - never-treated units. + Weights are assigned for every unit ``j != i`` (distance-based, per + Eq. 2/3). Treated-cell exclusion is handled by the `(1 - W_{js})` + factor applied inside ``_estimate_model`` via the control mask, not + by gating ``ω_j`` on ``D[t, j]``. Same-cohort donors therefore + contribute via their pre-treatment rows, and future-cohort donors + contribute via rows where both units are still untreated. Always computes from the function-argument ``Y, D``; does not read ``self._precomputed``. Under bootstrap the caller passes resampled diff --git a/rust/src/trop.rs b/rust/src/trop.rs index 706f3958..28cc2392 100644 --- a/rust/src/trop.rs +++ b/rust/src/trop.rs @@ -548,9 +548,11 @@ fn compute_unit_distance_for_obs( /// Unit weights: ω_j = exp(-λ_unit × dist(j, i)) /// /// Paper alignment notes: -/// - ALL units get weights (not just those untreated at target period) +/// - ALL units with `j != target_unit` get distance-based weights +/// (same-cohort donors contribute via their pre-treatment rows) /// - The (1 - D_js) masking in the loss naturally excludes treated cells -/// - Weights are normalized to sum to 1 (probability weights) +/// via the control mask applied inside `estimate_model` +/// - Weights are unnormalized raw exponentials per REGISTRY Eq. 2/3 /// - Distance excludes target period t per Equation 3 #[allow(clippy::too_many_arguments)] fn compute_weight_matrix( diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 59be799d..d54b5888 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -2420,25 +2420,37 @@ def test_bootstrap_seed_reproducibility(self, seed): f"Rust={res_rust.se:.16f}, Python={res_py.se:.16f}", ) - @pytest.mark.parametrize("seed", [0, 42, 12345]) - def test_bootstrap_seed_reproducibility_local(self, seed): + @pytest.mark.parametrize( + "seed,lambda_nn", + [ + (0, np.inf), + (42, np.inf), + (12345, np.inf), + (42, 0.1), + ], + ) + def test_bootstrap_seed_reproducibility_local(self, seed, lambda_nn): """Backend-invariant bootstrap SE parity for the local method. - Post-methodology-alignment regression guard. With the Rust weight- - matrix normalization dropped and the Python ``_precomputed`` - cache-fallthrough removed, local-method Rust and Python bootstraps - consume bit-identical stratified indices AND bit-identical raw- - exponential weights. Main-fit ATT is bit-identical across backends - (see ``test_local_method_main_fit_parity``), but per-replicate - bootstrap fits route through Rust's ``estimate_model`` vs numpy's - ``lstsq``, which use different matrix factorization paths and - accumulate different BLAS roundoff. Empirically the residual gap - is ~1e-7 relative; asserted at ``atol=1e-5`` which is ~100x the - observed gap and comfortable across CI runner variance. - - Follow-up to tighten to ``atol=1e-14``: unify Rust ``estimate_model`` - to use ``solve_wls_svd`` (the same SVD path used by global-method - since PR #348). Tracked in ``TODO.md``. + Post-methodology-alignment regression guard covering both the + ``lambda_nn=inf`` regime (no-lowrank path, closed by the Python + ``_precomputed`` cache-fallthrough removal) and the finite + ``lambda_nn`` regime (with-lowrank FISTA path, closed by the Rust + weight-matrix normalization removal). With the RNG fix from + PR #354 plus both methodology fixes landed here, local-method + Rust and Python bootstraps consume bit-identical stratified + indices AND bit-identical raw-exponential weights. Main-fit ATT + is bit-identical (see ``test_local_method_main_fit_parity``), + but per-replicate bootstrap fits route through Rust's + ``estimate_model`` vs numpy's ``lstsq``, which use different + matrix factorization paths and accumulate different BLAS + roundoff. Empirically the residual gap is ~1e-7 relative; + asserted at ``atol=1e-5`` which is ~100x the observed gap and + comfortable across CI runner variance. + + Follow-up to tighten to ``atol=1e-14``: unify Rust + ``estimate_model`` to use ``solve_wls_svd`` (the same SVD path + used by global-method since PR #348). Tracked in ``TODO.md``. """ import sys from unittest.mock import patch @@ -2450,7 +2462,7 @@ def test_bootstrap_seed_reproducibility_local(self, seed): method="local", lambda_time_grid=[1.0], lambda_unit_grid=[1.0], - lambda_nn_grid=[np.inf], + lambda_nn_grid=[lambda_nn], n_bootstrap=10, seed=seed, ) @@ -2466,7 +2478,8 @@ def test_bootstrap_seed_reproducibility_local(self, seed): np.testing.assert_allclose( res_rust.se, res_py.se, atol=1e-5, rtol=1e-5, - err_msg=f"Local-method bootstrap SE divergence under seed={seed}: " + 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}", ) From a5182c609d89ebcd660205cfaa91100ed6d52a06 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 24 Apr 2026 09:47:30 -0400 Subject: [PATCH 4/6] Address PR #358 AI review round 3: fix main-fit test patching + stale status text MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P2 — the new main-fit parity tests patched `diff_diff.trop_local. HAS_RUST_BACKEND` but the LOOCV dispatch lives in `diff_diff.trop` at line 572-583, so the "Python" branch was still running Rust LOOCV. Combined with singleton lambda grids (`[1.0]`, `[1.0]`, `[lambda_nn]`), the tests didn't actually exercise the Rust `compute_weight_matrix` surface the normalization fix changed. Fixed in both `test_local_method_main_fit_parity` and `test_local_method_same_cohort_donor_parity`: - Patch `diff_diff.trop.HAS_RUST_BACKEND` and `diff_diff.trop._rust_loocv_grid_search` in addition to the existing `diff_diff.trop_local` patches, so the Python branch forces Python LOOCV + Python per-observation weights end-to-end. - Expand grids to `[0.1, 1.0, 10.0]` (3x3) so LOOCV selection is non-trivial and actually scores candidates via the backend under test. P3 — stale status text updated: - `diff_diff/trop_local.py` bootstrap-index-generation comment (L940-L947) no longer says local SE parity is a queued follow-up; now correctly describes the full RNG + normalization + cache-fallthrough alignment and the bootstrap-SE tolerance regime. - `TODO.md` solver-path follow-up row no longer says main-fit ATT is uniformly `1e-14` bit-identical; rewritten to reflect the regime- dependent tolerance (`1e-14` for `lambda_nn=inf`, `1e-10` for finite `lambda_nn`) that's asserted in the test suite. Verification: all 12 `TestTROPRustEdgeCaseParity` tests pass under the tightened patching + multi-candidate grids. Co-Authored-By: Claude Opus 4.7 (1M context) --- TODO.md | 2 +- diff_diff/trop_local.py | 12 +++++------- tests/test_rust_backend.py | 28 ++++++++++++++++++++++------ 3 files changed, 28 insertions(+), 14 deletions(-) diff --git a/TODO.md b/TODO.md index 4a5c86c7..3c4853d6 100644 --- a/TODO.md +++ b/TODO.md @@ -83,7 +83,7 @@ Deferred items from PR reviews that were not addressed before merge. | Weighted CR2 Bell-McCaffrey cluster-robust (`vcov_type="hc2_bm"` + `cluster_ids` + `weights`) currently raises `NotImplementedError`. Weighted hat matrix and residual rebalancing need threading per clubSandwich WLS handling. | `linalg.py::_compute_cr2_bm` | Phase 1a | Medium | | Regenerate `benchmarks/data/clubsandwich_cr2_golden.json` from R (`Rscript benchmarks/R/generate_clubsandwich_golden.R`). Current JSON has `source: python_self_reference` as a stability anchor until an authoritative R run. | `benchmarks/R/generate_clubsandwich_golden.R` | Phase 1a | Medium | | `honest_did.py:1907` `np.linalg.solve(A_sys, b_sys) / except LinAlgError: continue` is a silent basis-rejection in the vertex-enumeration loop that is algorithmically intentional (try the next basis). Consider surfacing a count of rejected bases as a diagnostic when ARP enumeration exhausts, so users see when the vertex search was heavily constrained. Not a silent failure in the sense of the Phase 2 audit (the algorithm is supposed to skip), but the diagnostic would help debug borderline cases. | `honest_did.py` | #334 | Low | -| Unify Rust local-method `estimate_model` solver path to `solve_wls_svd` (the same SVD helper used by the global-method since PR #348) for sub-1e-14 bootstrap SE parity. Current local-method bootstrap parity test (`tests/test_rust_backend.py::TestTROPRustEdgeCaseParity::test_bootstrap_seed_reproducibility_local`) passes at `atol=1e-5` — the residual ~1e-7 gap is roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq`, which accumulates differently across per-replicate bootstrap fits. Main-fit ATT is bit-identical at `atol=1e-14` (see `test_local_method_main_fit_parity`), so this is a bootstrap-aggregation-only concern; not a user-visible correctness bug. | `rust/src/trop.rs::estimate_model`, `rust/src/linalg.rs::solve_wls_svd` | follow-up | Low | +| Unify Rust local-method `estimate_model` solver path to `solve_wls_svd` (the same SVD helper used by the global-method since PR #348) for sub-1e-14 bootstrap SE parity. Current local-method bootstrap parity test (`tests/test_rust_backend.py::TestTROPRustEdgeCaseParity::test_bootstrap_seed_reproducibility_local`) passes at `atol=1e-5` — the residual ~1e-7 gap is roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq`, which accumulates differently across per-replicate bootstrap fits. Main-fit ATT parity is regime-dependent (`atol=1e-14` for `lambda_nn=inf`, `atol=1e-10` for finite `lambda_nn` — see `test_local_method_main_fit_parity`); the bootstrap gap is a same-solver-path roundoff concern and not a user-visible correctness bug. | `rust/src/trop.rs::estimate_model`, `rust/src/linalg.rs::solve_wls_svd` | follow-up | Low | | Rust multiplier-bootstrap weight RNG (`generate_bootstrap_weights_batch` in `rust/src/bootstrap.rs:9-10, 57-75`) uses `Xoshiro256PlusPlus::seed_from_u64(seed + i)` per row for Rademacher/Mammen/Webb generation. If any Python caller (SDID / efficient-DiD multiplier bootstrap) has a numpy-canonical equivalent, the two backends likely diverge under the same seed. Audit Python callers (`diff_diff/sdid.py`, `diff_diff/efficient_did_bootstrap.py`, `diff_diff/bootstrap_utils.py::generate_bootstrap_weights_batch_numpy`) for parity-test gaps. Same fix shape as TROP RNG parity (PR #354): pre-generate weights in Python via numpy and pass them to Rust through PyO3. | `rust/src/bootstrap.rs`, `diff_diff/bootstrap_utils.py` | follow-up | Medium | | `bias_corrected_local_linear`: extend golden parity to `kernel="triangular"` and `kernel="uniform"` (currently epa-only; all three kernels share `kernel_W` and the `lprobust` math, so parity is expected but not separately asserted). | `benchmarks/R/generate_nprobust_lprobust_golden.R`, `tests/test_bias_corrected_lprobust.py` | Phase 1c | Low | | `bias_corrected_local_linear`: expose `vce in {"hc0", "hc1", "hc2", "hc3"}` on the public wrapper once R parity goldens exist (currently raises `NotImplementedError`). The port-level `lprobust` and `lprobust_res` already support all four; expanding the public surface requires a golden generator for each hc mode and a decision on hc2/hc3 q-fit leverage (R reuses p-fit `hii` for q-fit residuals; whether to match that or stage-match deserves a derivation before the wrapper advertises CCT-2014 conformance). | `diff_diff/local_linear.py::bias_corrected_local_linear`, `benchmarks/R/generate_nprobust_lprobust_golden.R`, `tests/test_bias_corrected_lprobust.py` | Phase 1c | Medium | diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index 231e313e..6efa3d7a 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -938,13 +938,11 @@ def _bootstrap_variance( n_control_units = len(control_units) # Pre-generate stratified bootstrap indices via numpy (Python-canonical RNG). - # This aligns the RNG layer between backends. For the global method the RNG - # was the only divergence, so global SE is now bit-identical under the same - # seed. For the local method two downstream divergences remain (Rust weight- - # matrix normalization + Python `_compute_observation_weights` reading the - # stale `_precomputed` cache) — tracked in TODO.md; until those land, local - # bootstrap SE still differs across backends. Silent-failures finding #23 - # (bootstrap half) closed for global; local follow-up queued. + # Aligns the RNG layer between backends. Combined with the Rust weight- + # matrix de-normalization and the Python `_compute_observation_weights` + # cache-fallthrough removal (also shipped with this parity work), local- + # method Rust and Python produce matching bootstrap SE up to solver-path + # roundoff (~1e-7); asserted at atol=1e-5 in the parity regression guard. rng = np.random.default_rng(self.seed) control_idx, treated_idx = stratified_bootstrap_indices( rng, n_control_units, n_treated_units, self.n_bootstrap diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index d54b5888..48d1f0e5 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -2505,6 +2505,13 @@ def test_local_method_main_fit_parity(self, lambda_nn, tol): fix landed in this PR. Before the fix, Rust ATT diverged from Python ATT by O(10%) at finite ``lambda_nn`` and O(0) at ``lambda_nn=inf``; after the fix both regimes match to tolerance. + + Uses multi-candidate lambda grids so LOOCV selection exercises + Rust's `compute_weight_matrix` (the surface the normalization fix + changed). Patches both the LOOCV dispatch in ``diff_diff.trop`` + and the bootstrap dispatch in ``diff_diff.trop_local`` on the + Python side so the comparison is Rust-LOOCV-and-fit vs + Python-LOOCV-and-fit end-to-end. """ import sys from unittest.mock import patch @@ -2512,10 +2519,13 @@ def test_local_method_main_fit_parity(self, lambda_nn, tol): from diff_diff import TROP df = self._make_correlated_panel(n_units=6, n_periods=6, n_treated=2) + # Multi-candidate grids so LOOCV selection isn't trivial; the Rust + # weight-normalization fix changes per-lambda LOOCV scores and thus + # potentially the selected lambda. trop_params = dict( method="local", - lambda_time_grid=[1.0], - lambda_unit_grid=[1.0], + lambda_time_grid=[0.1, 1.0, 10.0], + lambda_unit_grid=[0.1, 1.0, 10.0], lambda_nn_grid=[lambda_nn], n_bootstrap=2, # minimum allowed; we assert ATT, not SE seed=42, @@ -2524,8 +2534,11 @@ def test_local_method_main_fit_parity(self, lambda_nn, tol): trop_rust = TROP(**trop_params) res_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") + trop_module = sys.modules["diff_diff.trop"] trop_local_module = sys.modules["diff_diff.trop_local"] - with patch.object(trop_local_module, "HAS_RUST_BACKEND", False), \ + 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") @@ -2577,8 +2590,8 @@ def test_local_method_same_cohort_donor_parity(self, lambda_nn, tol): trop_params = dict( method="local", - lambda_time_grid=[1.0], - lambda_unit_grid=[1.0], + lambda_time_grid=[0.1, 1.0, 10.0], + lambda_unit_grid=[0.1, 1.0, 10.0], lambda_nn_grid=[lambda_nn], n_bootstrap=2, seed=42, @@ -2587,8 +2600,11 @@ def test_local_method_same_cohort_donor_parity(self, lambda_nn, tol): trop_rust = TROP(**trop_params) res_rust = trop_rust.fit(df.copy(), "outcome", "treated", "unit", "time") + trop_module = sys.modules["diff_diff.trop"] trop_local_module = sys.modules["diff_diff.trop_local"] - with patch.object(trop_local_module, "HAS_RUST_BACKEND", False), \ + 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") From b6417dca22570f068a8c71b947691e05b05b84e1 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 24 Apr 2026 10:07:25 -0400 Subject: [PATCH 5/6] Address PR #358 AI review round 4: expand CHANGELOG for same-cohort donor gate removal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P3 — CHANGELOG entry previously described only the `_precomputed` cache-fallthrough fix; the diff also removes the target-period `D[t,j]==0` donor gate in `_compute_observation_weights`, which the new `test_local_method_same_cohort_donor_parity` regression explicitly exercises. Release-notes readers could miss that Python local main-fit behavior is now structurally aligned with Rust + paper on same-cohort panels. Rewrote the Python-side Fixed bullet into two numbered sub-items (cache-fallthrough + donor-gate removal) with a note that the main-fit ATT is empirically unchanged on tested fixtures because same-cohort pre-treatment rows are exactly absorbed by their own unit fixed effect `alpha_j` without propagating into shared parameters — so this is structural alignment, not a numerical shift. The same-cohort regression test name is referenced so readers can find the guard. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fd308a9b..d35ad6b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,7 +20,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - SyntheticDiD bootstrap now retries degenerate resamples (all-control or all-treated, or non-finite `τ_b`) until exactly `n_bootstrap` valid replicates are accumulated, matching R's `synthdid::bootstrap_sample` and Arkhangelsky et al. (2021) Algorithm 2. Previously the Python path counted attempts (with degenerate draws silently dropped), producing fewer valid replicates than requested. A bounded-attempt guard (`20 × n_bootstrap`) prevents pathological-input hangs. - **TROP global bootstrap SE backend parity under fixed seed** — Rust and Python backends now produce bit-identical bootstrap SE under the same `seed`. Previously Rust's `bootstrap_trop_variance_global` seeded `rand_xoshiro::Xoshiro256PlusPlus` per replicate while Python's fallback consumed `numpy.random.default_rng` (PCG64), producing ~28% SE divergence on tiny panels under `seed=42`. Fixed by extracting a shared `stratified_bootstrap_indices` helper in `diff_diff/bootstrap_utils.py` that pre-generates per-replicate stratified sample indices via numpy on the Python side; both backends consume the same integer arrays through the PyO3 surface. Sampling law (stratified: controls then treated, with replacement) is unchanged. Closes silent-failures audit finding #23 (bootstrap half; grid-search half closed in PR #348). Local-method TROP also adopts the Python-canonical index contract for the RNG layer, but separately-discovered backend divergences (Rust normalizes weight-matrix outer product, Python `_compute_observation_weights` reads stale `_precomputed` cache) prevent local-method bit-identity SE; tracked as a follow-up in `TODO.md`. - **TROP local-method Rust weight-matrix no longer normalized** — `rust/src/trop.rs::compute_weight_matrix` no longer divides time-weights or unit-weights by their respective sums before the outer product. The paper's Equation 2/3 (Athey, Imbens, Qu, Viviano 2025) and REGISTRY.md Requirements checklist (line 2037: `[x] Unit weights: exp(-λ_unit × distance) (unnormalized, matching Eq. 2)`) both specify raw-exponential weights; Python's `_compute_observation_weights` was already REGISTRY-compliant. **User-visible effect**: Rust local-method ATT values may shift for any fit with `lambda_nn < infinity` — normalizing the weight-matrix inflated the effective nuclear-norm penalty relative to the data-fit term, changing the regularization trade-off. For `lambda_nn = infinity` (factor model disabled) outputs are unchanged because uniform weight scaling leaves the minimum-norm WLS argmin invariant. Rust LOOCV-selected lambdas may also shift on this boundary; both backends now converge on the same REGISTRY-compliant selection. -- **TROP local-method Python `_compute_observation_weights` no longer reads stale `_precomputed` cache** — removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`). Together with the normalization fix above, TROP local-method backend parity on the main-fit ATT is regime-dependent: `atol=rtol=1e-14` for `lambda_nn=inf` (no nuclear-norm regularization, uniform weight scaling leaves the WLS argmin invariant) and `atol=1e-10` for finite `lambda_nn` (FISTA inner loop + BLAS reduction ordering introduce sub-1e-10 roundoff across Rust `faer` vs numpy paths). Bootstrap SE parity is asserted at `atol=1e-5` to accommodate ~1e-7 roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq` that accumulates across per-replicate fits; sub-1e-14 bootstrap parity is tracked as a follow-up in `TODO.md` under "unify Rust local-method solver path". Closes silent-failures audit finding #23 (local-method half; the RNG half closed in PR #354 and the grid-search half in PR #348). +- **TROP local-method Python `_compute_observation_weights` now uses the function-argument `Y, D` and treats all non-target units as donors** — two coupled changes that bring Python structurally in line with Rust and the paper's Eq. 2/3: + 1. Removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`). + 2. Removed the `valid_control_at_t = D[t, :] == 0` target-period donor gate that zeroed `ω_j` for any unit `j` treated at the target period (other than the target unit itself). Per REGISTRY Eq. 2/3 and Rust's `compute_weight_matrix`, `ω_j = exp(-λ_unit × dist(j, i))` for all `j ≠ i`; treated-cell exclusion happens via the `(1 − W_{js})` factor applied inside `_estimate_model`. Same-cohort donors now contribute via their pre-treatment rows. Empirically the main-fit ATT is unchanged on tested fixtures because same-cohort pre-treatment observations are exactly absorbed by their own unit fixed effect `alpha_j` without propagating into `mu`, `beta`, or other units' parameters — so this change is structural alignment rather than a numerical shift in output. Users on same-cohort panels with very few controls may still see tiny differences in edge cases; the new `test_local_method_same_cohort_donor_parity` regression guards the aligned behavior. + Together with the normalization fix above, TROP local-method backend parity on the main-fit ATT is regime-dependent: `atol=rtol=1e-14` for `lambda_nn=inf` (no nuclear-norm regularization, uniform weight scaling leaves the WLS argmin invariant) and `atol=1e-10` for finite `lambda_nn` (FISTA inner loop + BLAS reduction ordering introduce sub-1e-10 roundoff across Rust `faer` vs numpy paths). Bootstrap SE parity is asserted at `atol=1e-5` to accommodate ~1e-7 roundoff between Rust's `estimate_model` matrix factorization and numpy's `lstsq` that accumulates across per-replicate fits; sub-1e-14 bootstrap parity is tracked as a follow-up in `TODO.md` under "unify Rust local-method solver path". Closes silent-failures audit finding #23 (local-method half; the RNG half closed in PR #354 and the grid-search half in PR #348). ### Changed - **`did_had_pretest_workflow(aggregate="event_study")` verdict no longer emits the "paper step 2 deferred to Phase 3 follow-up" caveat** — the joint pre-trends Stute test closes that gap. The two-period `aggregate="overall"` path retains the existing caveat since the joint variant does not apply to single-pre-period panels. Downstream code that greps verdict strings for the Phase 3 caveat will see it suppressed on the event-study path. From af6ac60a8d8f8b688f6a60febc089a23937ea773 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 24 Apr 2026 10:45:26 -0400 Subject: [PATCH 6/6] Address PR #358 AI review round 5: align PR #354 CHANGELOG tail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P3 — the PR #354 [Unreleased] Fixed entry (line 21) said local-method bit-identity SE remained blocked by the Rust-normalization and Python cache-fallthrough divergences and was "tracked as a follow-up in TODO.md." With the two TROP-local Fixed entries that this PR adds (lines 22-27) closing exactly those divergences, the PR #354 tail sentence is now internally inconsistent with the surrounding entries. Rewritten to say the RNG half of finding #23 is closed here (bootstrap contract), grid-search half was closed in PR #348, and the local- method methodology half is closed by the two Fixed entries that follow in the same release. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d35ad6b4..d83b93fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - SyntheticDiD `variance_method="bootstrap"` now computes p-values from the analytical normal-theory formula using the bootstrap SE (matching R's `synthdid::vcov()` convention), rather than an empirical null-distribution formula that is not valid for bootstrap draws. `is_significant` and `significance_stars` are derived from `p_value` and will also change for bootstrap fits. Placebo and jackknife are unchanged. Point estimates are unaffected. - SyntheticDiD bootstrap SE formula applies the `sqrt((r-1)/r)` correction matching R's synthdid and the placebo SE formula. - SyntheticDiD bootstrap now retries degenerate resamples (all-control or all-treated, or non-finite `τ_b`) until exactly `n_bootstrap` valid replicates are accumulated, matching R's `synthdid::bootstrap_sample` and Arkhangelsky et al. (2021) Algorithm 2. Previously the Python path counted attempts (with degenerate draws silently dropped), producing fewer valid replicates than requested. A bounded-attempt guard (`20 × n_bootstrap`) prevents pathological-input hangs. -- **TROP global bootstrap SE backend parity under fixed seed** — Rust and Python backends now produce bit-identical bootstrap SE under the same `seed`. Previously Rust's `bootstrap_trop_variance_global` seeded `rand_xoshiro::Xoshiro256PlusPlus` per replicate while Python's fallback consumed `numpy.random.default_rng` (PCG64), producing ~28% SE divergence on tiny panels under `seed=42`. Fixed by extracting a shared `stratified_bootstrap_indices` helper in `diff_diff/bootstrap_utils.py` that pre-generates per-replicate stratified sample indices via numpy on the Python side; both backends consume the same integer arrays through the PyO3 surface. Sampling law (stratified: controls then treated, with replacement) is unchanged. Closes silent-failures audit finding #23 (bootstrap half; grid-search half closed in PR #348). Local-method TROP also adopts the Python-canonical index contract for the RNG layer, but separately-discovered backend divergences (Rust normalizes weight-matrix outer product, Python `_compute_observation_weights` reads stale `_precomputed` cache) prevent local-method bit-identity SE; tracked as a follow-up in `TODO.md`. +- **TROP global bootstrap SE backend parity under fixed seed** — Rust and Python backends now produce bit-identical bootstrap SE under the same `seed`. Previously Rust's `bootstrap_trop_variance_global` seeded `rand_xoshiro::Xoshiro256PlusPlus` per replicate while Python's fallback consumed `numpy.random.default_rng` (PCG64), producing ~28% SE divergence on tiny panels under `seed=42`. Fixed by extracting a shared `stratified_bootstrap_indices` helper in `diff_diff/bootstrap_utils.py` that pre-generates per-replicate stratified sample indices via numpy on the Python side; both backends consume the same integer arrays through the PyO3 surface. Sampling law (stratified: controls then treated, with replacement) is unchanged. Closes the bootstrap-RNG half of silent-failures audit finding #23 (grid-search half closed in PR #348; local-method methodology half closed by the two Fixed entries below). Local-method TROP also adopts the Python-canonical index contract for the RNG layer here. - **TROP local-method Rust weight-matrix no longer normalized** — `rust/src/trop.rs::compute_weight_matrix` no longer divides time-weights or unit-weights by their respective sums before the outer product. The paper's Equation 2/3 (Athey, Imbens, Qu, Viviano 2025) and REGISTRY.md Requirements checklist (line 2037: `[x] Unit weights: exp(-λ_unit × distance) (unnormalized, matching Eq. 2)`) both specify raw-exponential weights; Python's `_compute_observation_weights` was already REGISTRY-compliant. **User-visible effect**: Rust local-method ATT values may shift for any fit with `lambda_nn < infinity` — normalizing the weight-matrix inflated the effective nuclear-norm penalty relative to the data-fit term, changing the regularization trade-off. For `lambda_nn = infinity` (factor model disabled) outputs are unchanged because uniform weight scaling leaves the minimum-norm WLS argmin invariant. Rust LOOCV-selected lambdas may also shift on this boundary; both backends now converge on the same REGISTRY-compliant selection. - **TROP local-method Python `_compute_observation_weights` now uses the function-argument `Y, D` and treats all non-target units as donors** — two coupled changes that bring Python structurally in line with Rust and the paper's Eq. 2/3: 1. Removed the `if self._precomputed is not None:` branch that silently substituted `self._precomputed["Y"]` / `["D"]` / `["time_dist_matrix"]` (original-panel cache populated during main fit) for the function-argument `Y, D`. Under bootstrap, `_fit_with_fixed_lambda` computes fresh `Y, D` from the resampled `boot_data` and passes them in; the helper was discarding those and recomputing unit distances from the original panel, so Python's local bootstrap resampled units but reused stale unit-distance weights. Rust's bootstrap was already correct (always consumed `y_boot, d_boot`).