Commit d17dc87
Optimize TROP estimator performance with vectorized operations
Implement performance optimizations for the TROP estimator while preserving
exact methodology from Athey, Imbens, Qu & Viviano (2025). Target: 5-10x speedup.
Key optimizations:
- Vectorized data matrix construction using pandas pivot() instead of iterrows
- Pre-computed structures (_precompute_structures) for distance matrices,
time distances, and observation lists reused across LOOCV iterations
- Vectorized unit distance computation with symmetric matrix storage
- Optimized observation weights using pre-computed distance matrices
- Vectorized alternating minimization with matrix operations instead of
nested loops for alpha, beta, and L updates
- LOOCV optimization using shared pre-computed structures
Added TestOptimizationEquivalence test class with 5 tests verifying:
- Pre-computed structures consistency
- Vectorized alternating minimization convergence
- Weight computation correctness
- Pivot vs iterrows equivalence
- Reproducibility with seed
All 33 tests pass (28 original + 5 new equivalence tests).
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>1 parent f2058e9 commit d17dc87
2 files changed
Lines changed: 567 additions & 128 deletions
0 commit comments