Skip to content

Commit d17dc87

Browse files
igerberclaude
andcommitted
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

File tree

0 commit comments

Comments
 (0)