Add unified LOOCV for TROP joint method with Rust acceleration

[igerber]

Summary

• Implement proper leave-one-out cross-validation (LOOCV) for the TROP joint method, matching the twostep method's approach per the paper's Equation 5
• Add Rust backend acceleration for parallel LOOCV grid search (loocv_grid_search_joint())
• Add Rust acceleration for parallel bootstrap variance estimation (bootstrap_trop_variance_joint())
• Add Python fallback implementation (_loocv_score_joint()) when Rust unavailable
• Update _fit_joint() and _bootstrap_variance_joint() to use Rust when available

Methodology references (required if estimator / math changes)

• Method name(s): TROP (Triply Robust Panel) - Joint Optimization Method
• Paper / source link(s): Athey, S., Imbens, G. W., Qu, Z., & Viviano, D. (2025). Triply Robust Panel Estimators. Working Paper. https://arxiv.org/abs/2508.21536
• Any intentional deviations from the source (and why): None. The joint method now uses the same LOOCV framework (Equation 5) as the twostep method, computing pseudo-treatment effects at held-out control observations.

Validation

• Tests added/updated:
  • tests/test_trop.py: Added test_joint_loocv_selects_from_grid, test_joint_loocv_score_internal
  • tests/test_rust_backend.py: Added TestTROPJointRustBackend class (4 tests) and TestTROPJointRustVsNumpy class (2 tests)
• All 12 joint method tests pass in Python mode
• Rust code compiles cleanly (cargo check passes)

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Methodology

• P1 | Impact: Unbalanced panels with missing outcomes can yield NaN weights or biased estimates because Python propagates NaNs into weights/OLS while Rust imputes non-finite Y as 0.0; this violates the registry’s unbalanced-panel expectation and creates backend divergence. | Fix: mask missing outcomes consistently (e.g., np.nanmean for treated averages, set delta to 0 for missing Y, and exclude non-finite cells in WLS); mirror the same logic in Rust (skip non-finite rather than zero-impute). | Ref: diff_diff/trop.py:L860, diff_diff/trop.py:L1022, diff_diff/trop.py:L1099, rust/src/trop.rs:L1047, rust/src/trop.rs:L1136, rust/src/trop.rs:L1260, docs/methodology/REGISTRY.md:L572
• P1 | Impact: Python vs Rust low-rank updates scale the nuclear-norm threshold differently (lambda_nn vs eta*lambda_nn), so joint estimates/LOOCV can diverge by backend. | Fix: align threshold scaling across backends and document the scaling choice; add a small parity test. | Ref: diff_diff/trop.py:L1099, rust/src/trop.rs:L1260
• P1 | Impact: Python bootstrap uses treated_periods from the original fit even when bootstrap samples shift the first-treated period (staggered/resampling), moving the weight center and SE; Rust recomputes from D, so backends differ. | Fix: recompute treated_periods from each bootstrap sample’s D (or move computation inside _fit_joint_with_fixed_lambda) to match Rust. | Ref: diff_diff/trop.py:L1479, diff_diff/trop.py:L1619, rust/src/trop.rs:L1545

Documentation/Tests

• P2 | Impact: New joint-method tests cover balanced panels only and don’t exercise missing outcomes or backend parity, so the above regressions can slip. | Fix: add unbalanced/NaN tests for method='joint' and a small Rust-vs-Python parity check (ATT/LOOCV/SE). | Ref: tests/test_trop.py:L2660, tests/test_rust_backend.py:L1138

Code Quality

• P3 | Impact: Rust LOOCV docstring claims a two-stage search but the implementation is full grid search, which can mislead maintainers. | Fix: update the comment or implement the two-stage approach. | Ref: rust/src/trop.rs:L1429

Performance

• No issues found.

Maintainability

• No issues found.

Tech Debt

• No issues found.

Security

• No issues found.

Overall Assessment: ⚠️ Needs changes

Executive Summary

• Affected method: TROP joint optimization (method='joint') plus Rust LOOCV/bootstrap acceleration.
• Missing-outcome handling is inconsistent across Python/Rust and violates the registry’s unbalanced-panel expectation.
• Low-rank nuclear-norm threshold scaling differs across backends, leading to divergent estimates/LOOCV.
• Python bootstrap uses fixed treated-period length, diverging from Rust for staggered/adoption resamples.
• Tests don’t cover unbalanced panels or backend parity for the joint method.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 7e41d768577b090e041ee1ef7a8b104ecedb8db1

---

Methodology

• Severity: P1; Impact: TROP joint low-rank estimation (and LOOCV/variance paths) treats NaN outcomes as 0, so missing observations influence mu/alpha/beta/L/tau, contradicting the docstring claim that NaNs are zero-weighted and the registry’s unbalanced-panel support; Fix: mask delta with np.isfinite(Y) and carry that mask through the low-rank updates in both Python and Rust so missing observations are excluded. diff_diff/trop.py:1153 diff_diff/trop.py:1163 rust/src/trop.rs:1279
• Severity: P1; Impact: Joint jackknife variance imputes the left-out unit with column means, so it still contributes as a control unit and can understate SEs; Fix: implement a true leave-one-unit-out by zeroing weights (or dropping the unit column) without imputation, then reuse the masked missing-data handling. diff_diff/trop.py:1734 diff_diff/trop.py:1739
• Severity: P1; Impact: Joint bootstrap/jackknife reuse treated_periods from the original sample; with staggered adoption, resamples that drop early-treated units shift the treated block, so time-weight centers are inconsistent with the estimator definition; Fix: recompute treated_periods from each resampled D in _fit_joint_with_fixed_lambda and the Rust bootstrap loop (or document a fixed-timing assumption). diff_diff/trop.py:1621 diff_diff/trop.py:1679 rust/src/trop.rs:1577
• Severity: P2; Impact: Unit-weight RMSE gives distance 0 for units with no valid pre-period data, and Python/Rust diverge when treated averages are missing, yielding backend-dependent weights for unbalanced panels; Fix: compute RMSE only where both treated averages and unit outcomes are finite, and set weight 0 when no valid pre data, aligned across backends. diff_diff/trop.py:914 diff_diff/trop.py:927 rust/src/trop.rs:1077 rust/src/trop.rs:1108

Code Quality

• No issues found.

Performance

• No issues found.

Maintainability

• No issues found.

Tech Debt

• No issues found.

Security

• No issues found.

Documentation/Tests

• Severity: P3; Impact: Notebook summary states weights “sum to one,” which is not true for the joint weights (and is not guaranteed in twostep); Fix: revise the summary or normalize weights to match the statement. docs/tutorials/10_trop.ipynb:696

Open Questions

• Is the joint bootstrap/jackknife intended to keep treated_periods fixed across resamples (assume common adoption), or should it recompute from each resampled D matrix?
• Is imputing NaNs to 0 in the joint low-rank path intended? If so, should REGISTRY/docs explicitly note that method="joint" does not support unbalanced panels?

Overall Assessment

• ⚠️ Needs changes

Executive Summary

• Joint low-rank path includes missing outcomes as zeros, which can bias estimates, LOOCV, and SEs in unbalanced panels.
• Jackknife variance is not a true leave-one-unit-out because the dropped unit is imputed back in.
• Bootstrap/jackknife reuse original treated_periods, miscentering time weights under staggered adoption.
• Unit-weight RMSE for missing pre data yields backend-dependent weights.
• Notebook summary overstates weight normalization.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: f58d690ad7cc6141f41758eec4ce26be33e9558f

---

Methodology

• Severity: P1 – Impact: Rust LOOCV/bootstrapping for the joint method treats NaN outcomes as zero with non‑zero weights, which violates the “unbalanced panels supported” expectation and diverges from Python handling; this can bias tuning and SE when outcomes are missing – Fix: mask weights where Y is non‑finite in rust/src/trop.rs (solve_joint_no_lowrank, solve_joint_with_lowrank, and any shared weighting path) to mirror diff_diff/trop.py NaN masking in TROP._solve_joint_with_lowrank/_solve_joint_no_lowrank.
• Severity: P1 – Impact: Units with no valid pre‑period data receive full weight in Rust (n_valid==0 → dist=0 → weight=1), while Python sets zero weight; this can distort joint weights and LOOCV/SE in unbalanced panels – Fix: set delta_unit[i]=0.0 when n_valid==0 in rust/src/trop.rs compute_joint_weights to match diff_diff/trop.py TROP._compute_joint_weights.

Code Quality

• No findings.

Performance

• No findings.

Maintainability

• No findings.

Tech Debt

• No findings.

Security

• No findings.

Documentation/Tests

• Severity: P2 – Impact: Rust backend behavior for NaN handling and no‑valid‑pre units is untested, so regressions in LOOCV/bootstrapping may slip through – Fix: add Rust parity tests that assert NaNs are excluded and no‑pre units get zero weight when using loocv_grid_search_joint/bootstrap_trop_variance_joint (tests/test_rust_backend.py).

Open Questions

• Assumption: Rust joint backend should match Python’s missing‑data behavior for weighting; confirm this is intended before changing.

Overall Assessment
⚠️ Needs changes

Executive Summary

• Rust joint LOOCV/bootstrapping currently includes NaN outcomes with positive weight, which can bias tuning/SE on unbalanced panels.
• Rust joint weights give full weight to units with no valid pre‑period data, conflicting with Python behavior and intended edge‑case handling.
• Add Rust parity tests for NaN exclusion and no‑pre unit weighting to prevent regressions.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: cef9d31bf40e5cce2395a67e84acfdb03aa9c9d1

---

Overall assessment: ⚠️ Needs changes

Executive summary

• Joint TROP method is integrated across Python/Rust and documented, but two methodology‑critical gaps remain.
• Missing assumption check for simultaneous adoption in method="joint" risks invalid weights/SEs under staggered adoption.
• Rust joint estimator treats NaN outcomes as zeros (weights still applied), diverging from Python and “unbalanced panels supported.”
• Unit‑weight RMSE uses periods with undefined treated averages, skewing weights and creating Python/Rust inconsistency.

Methodology

• P1 | Impact: Joint method assumes simultaneous adoption per Methodology Registry, but no runtime check/warn; staggered adoption yields incorrect weights and invalid bootstrap/jackknife SEs. | Fix: compute per‑unit first treatment period; if multiple unique starts, raise or warn and require method="twostep" (or skip joint variance). | Ref: diff_diff/trop.py:L1268-L1318
• P1 | Impact: Rust joint estimation includes NaN outcomes with positive weights (imputed as 0), biasing τ/SE in unbalanced panels and breaking parity with Python. | Fix: mask out NaN outcomes by zeroing delta (or skip terms) in solve_joint_no_lowrank and solve_joint_with_lowrank; ensure LOOCV/bootstraps use the same mask. | Ref: rust/src/trop.rs:L1140-L1216, rust/src/trop.rs:L1264-L1315
• P2 | Impact: Unit‑weight RMSE includes pre‑periods where treated average is undefined, distorting distances and making Python/Rust inconsistent under missing treated data. | Fix: require both average_treated[t] and Y[t,i] finite for numerator/denominator; in Rust track treated counts per period and skip count==0 periods. | Ref: diff_diff/trop.py:L912-L941, rust/src/trop.rs:L1077-L1110

Code Quality

• No findings.

Performance

• No findings.

Maintainability

• No findings.

Tech Debt

• No findings.

Security

• No findings.

Documentation/Tests

• No findings.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: a7a115de65232078e2d0d418986f1e0d8038b440

---

Overall assessment: ⚠️ Needs changes

Executive Summary

• P1: Joint-method unit-weight RMSE in Python doesn’t exclude pre-periods where the treated average is undefined, inflating weights and diverging from registry/Rust behavior; this can bias LOOCV/ATT in unbalanced panels.
• P2: Tests cover NaNs in control observations only; missing treated-pre NaN coverage leaves the above weighting edge case untested.
• The rest of the joint LOOCV flow and simultaneous-adoption checks appear aligned with the Methodology Registry.

Methodology

• Severity: P1; Impact: The registry defines unit weights as RMSE over pre-periods against the treated average, but Python counts pre-periods even when the treated average is NaN, shrinking distances and overweighting units; this also creates Python/Rust divergence in unbalanced panels; Fix: mask on np.isfinite(average_treated) when computing diff_sq/valid_count, and treat “no valid pre-period pairs” as invalid for LOOCV (docs/methodology/REGISTRY.md:L595-L599, diff_diff/trop.py:L912-L941, rust/src/trop.rs:L1074-L1112).

Code Quality

• None.

Performance

• None.

Maintainability

• Severity: P3; Impact: Rust docstring for loocv_grid_search_joint claims a two-stage coordinate-descent search, but the implementation is full parallel grid search, which can confuse future maintenance and contradict registry text; Fix: update the docstring to match current behavior or implement the described two-stage algorithm (rust/src/trop.rs:L1420-L1493).

Tech Debt

• None.

Security

• None.

Documentation/Tests

• Severity: P2; Impact: New NaN tests only set NaNs in control observations, so treated-pre NaNs (which drive average_treated to NaN and trigger the weighting bug above) remain untested; Fix: add a test where at least one pre-period has all treated outcomes missing and assert consistent handling across Python/Rust or that LOOCV rejects the case (tests/test_trop.py:L2952-L2997, tests/test_rust_backend.py:L1367-L1404).

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 4a27ec58dbcc56e3c1659a41861abebdfdc6711e

---

Overall assessment: ⚠️ Needs changes
Executive summary:

• Affected method: TROP joint optimization (method="joint"); cross-checked against docs/methodology/REGISTRY.md.
• P1: Simultaneous-adoption check can false‑positive on unbalanced panels with missing treatment entries.
• P1: Python joint solver can proceed with all‑zero weights (no valid pre data), yielding misleading estimates and Rust/Python divergence.
• P2: Rust vs Python parity tests may not actually switch to the Python backend.

Methodology

• P1 – Simultaneous‑adoption guard ignores missing treatment entries: first_treat_by_unit uses D with NaNs filled as 0, so unbalanced panels can be falsely flagged as staggered adoption, conflicting with the registry’s “missing observations don’t trigger false violations.” Impact: valid data rejected or forced into method="twostep". Fix: compute first‑treatment periods using D_raw/missing_mask (skip missing) or apply observed_mask when determining per‑unit first treatment. (diff_diff/trop.py:L1270-L1335, docs/methodology/REGISTRY.md:L569-L572)
• P1 – Zero‑weight fits not treated as invalid in Python joint solver: if treated pre‑period data are all missing, _compute_joint_weights can yield all‑zero weights but _solve_joint_no_lowrank still returns coefficients, while Rust rejects via sum_w < 1e-10. Impact: misleading ATT/SE and backend divergence in missing‑pre edge cases. Fix: add a sum(weights) == 0 / np.any(delta) guard that raises ValueError (so LOOCV returns inf) and align Python with Rust’s failure behavior. (diff_diff/trop.py:L899-L1124, rust/src/trop.rs:L1182-L1184)

Code Quality

• No issues found.

Performance

• No issues found.

Maintainability

• No issues found.

Tech Debt

• No issues found.

Security

• No issues found.

Documentation/Tests

• P2 – Backend parity tests don’t actually switch to Python: after setting DIFF_DIFF_BACKEND=python, the tests reload diff_diff._backend but keep the already‑imported TROP, so HAS_RUST_BACKEND may remain true and the “python” run can still use Rust. Impact: parity tests can pass without exercising the Python fallback. Fix: reload diff_diff.trop (or diff_diff) after the env change and re‑bind TROP, or delay importing TROP until after the backend switch. (tests/test_rust_backend.py:L1505-L1576)

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 9eeb4363fb767d44f0c7681f03ff63ee65058891

---

Overall assessment: ⚠️ Needs changes

Executive Summary

• P1: Python-only LOOCV subsampling for the joint method likely errors when control_obs exceeds max_loocv_samples, breaking non‑Rust usage.
• Methodology alignment: Joint LOOCV equation, weight construction, and simultaneous‑adoption guard match docs/methodology/REGISTRY.md.
• Tests are strong for Rust parity, but no Python‑only LOOCV subsampling test to catch the above.

Methodology

• No issues found; implementation aligns with REGISTRY and in‑code docstrings for joint LOOCV and weights.

Code Quality

• P1: Impact: Python fallback LOOCV sampling can raise ValueError: a must be 1-dimensional when control_obs (list of tuples) is subsampled, preventing method="joint" from running without Rust on typical panels. Concrete fix: sample indices (as in twostep) or convert to a 1D object array before choice. Location: diff_diff/trop.py:_fit_joint (LOOCV fallback control_obs subsampling block).

Performance

• No issues found.

Maintainability

• No issues found.

Tech Debt

• No issues found.

Security

• No issues found.

Documentation/Tests

• P3: Impact: Missing regression test for Python‑only joint LOOCV subsampling allows the P1 to slip when Rust is unavailable. Concrete fix: add a test that sets DIFF_DIFF_BACKEND=python, creates >max_loocv_samples control observations, and asserts fit() completes. Location: tests/test_trop.py (new test).

[igerber]

/ai-review

[igerber]

/ai-review

[igerber]

/ai-review

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 4f28c0589a2ae291ffb2a1b2eca86c097be5b1c0

---

Overall Assessment: ✅ Looks good

Executive Summary

• Methodology for joint LOOCV/weights and simultaneous adoption aligns with docs/methodology/REGISTRY.md; no deviations found.
• P2 performance risk: Python fallback builds a dense WLS design matrix that can blow memory for large panels when Rust isn’t available.
• P3 maintainability: joint data prep duplicates twostep setup, increasing risk of drift.

Methodology

• No findings; LOOCV and weight computation match the registry and in-code docstrings.

Code Quality

• No findings.

Performance

• P2 Impact: Python fallback allocates a dense n_obs x (n_units+n_periods) matrix for WLS, which can be prohibitive on large panels when Rust is unavailable. Fix: implement a sparse/iterative solver in Python (mirror Rust coordinate descent) or gate to Rust for large n_units*n_periods. (diff_diff/trop.py:L1087-L1116)

Maintainability

• P3 Impact: Panel setup and validation logic is duplicated between _fit_joint and fit, raising drift risk as validation rules evolve. Fix: extract shared prep/validation into a helper used by both paths. (diff_diff/trop.py:L1258-L1323, diff_diff/trop.py:L1884-L1913)

Tech Debt

• No findings.

Security

• No findings.

Documentation/Tests

• No findings; docs and tests added for joint method and Rust parity.