Add TROP estimator #72
Merged
+3,370
−2
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Implements the TROP estimator from Athey, Imbens, Qu & Viviano (2025) following the exact paper methodology: - Nuclear norm regularized factor model (interactive fixed effects) - Exponential distance-based unit weights: ω_j = exp(-λ_unit × d(j,i)) - Exponential time decay weights: θ_s = exp(-λ_time × |s-t|) - LOOCV tuning parameter selection over (λ_time, λ_unit, λ_nn) grid - Bootstrap and jackknife variance estimation New files: - diff_diff/trop.py: TROP class and TROPResults dataclass - tests/test_trop.py: 23 tests covering all functionality - docs/tutorials/10_trop.ipynb: Tutorial notebook Updated files: - diff_diff/__init__.py: Add TROP exports, bump version to 2.1.0 - pyproject.toml: Bump version to 2.1.0 - README.md: Add TROP documentation and API reference - CLAUDE.md: Add module documentation - CHANGELOG.md: Add v2.1.0 release notes Reference: https://arxiv.org/abs/2508.21536 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Key changes to match the reference implementation: 1. Unit distance computation: Now computes RMSE from the average of treated units over pre-treatment periods, matching the paper's Equation 3. 2. Time distance computation: Centers weights around the treatment period midpoint (T - T_treat/2) rather than individual observations. 3. Global weight matrix: Computes weights once as outer product of unit and time weights, matching reference implementation. 4. Numerical stability improvements: - Truncated SVD reconstruction using only non-zero singular values - Input/output sanitization for NaN/Inf values - Suppress expected numerical warnings during ill-conditioned matmul - Proper handling of edge cases in soft-thresholding All 23 tests pass with no warnings. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Include the Athey et al. (2025) TROP paper and a detailed analysis document comparing the reference implementation to the paper. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
igerber
pushed a commit
that referenced
this pull request
Jan 18, 2026
Review identifies critical methodology deviations from the paper: - Unit distance computed from treated average instead of pairwise - Time distance from treatment center instead of specific period - Global weights instead of observation-specific weights per (i,t) - Algorithm structure differs from paper's Algorithm 2 Recommends changes before merging to preserve triple robustness property.
This commit addresses the methodology review and aligns the TROP implementation
with the paper's specification (Athey, Imbens, Qu & Viviano 2025):
**Algorithm Changes:**
- Restructured fit() to follow Algorithm 2: for each treated (i,t), compute
observation-specific weights, fit model, compute τ̂_{it}, then average
- Changed unit distance from "average of treated" to pairwise RMSE between
each control unit j and specific treated unit i (Equation 3)
- Changed time distance from |s - (T - T_treat/2)| to simple |t - s| where
t is the specific treatment period (Equation 3)
- Added _compute_observation_weights() for per-(i,t) weight matrices
- Updated _loocv_score_obs_specific() to use observation-specific weights
- Updated bootstrap and jackknife variance methods accordingly
**New Methodology Tests:**
- test_limiting_case_uniform_weights: λ_unit = λ_time = λ_nn = 0 gives TWFE-like
- test_unit_weights_reduce_bias: unit weighting helps with heterogeneous controls
- test_time_weights_reduce_bias: time weighting helps with trending data
- test_factor_model_reduces_bias: nuclear norm helps with factor structure
- test_paper_dgp_recovery: validates against paper's Table 2 simulation DGP
**Performance Note:**
The per-observation model fitting is more computationally intensive but
provides the triple robustness property described in Theorem 5.1.
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Create docs/api/trop.rst with full API documentation - Add TROP and TROPResults to docs/api/index.rst - Include algorithm description, tuning parameters, and usage examples - Reference paper (Athey, Imbens, Qu & Viviano 2025) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Add 'treat' column (unit-level ever-treated indicator) to generate_factor_dgp() - Update SDID calls to use 'treat' instead of 'treated' - TROP continues to use 'treated' (observation-level indicator) - Add clarifying comments explaining the difference The tutorial now correctly demonstrates TROP vs SDID comparison, showing TROP's advantage under factor confounding (bias 0.08 vs 1.31). Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Addition of TROP estimator based on https://arxiv.org/abs/2508.21536