From f5941233e580c7b452a7b6205770bb3fc4d99f6e Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 31 Jan 2026 17:58:23 -0500 Subject: [PATCH 1/3] Align TROP lambda conventions with paper (Athey et al. 2025) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix disabled-parameter semantics for λ_time and λ_unit to match the paper's convention: 0.0 means uniform weights (exp(-0×dist)=1), not ∞. Only λ_nn=∞ remains valid (disables factor model). Add ValueError validation rejecting inf in time/unit grids, remove dead ∞→0 conversions in Python and Rust, fix Stage 1 LOOCV fixed values, and update Registry equations to match Section 2.2 and Eq. 2-3 of the paper. Co-Authored-By: Claude Opus 4.5 --- diff_diff/trop.py | 119 ++++++++++++++++--------------- docs/methodology/REGISTRY.md | 53 ++++++++------ rust/src/trop.rs | 56 ++++++--------- tests/test_trop.py | 131 ++++++++++++++++++++--------------- 4 files changed, 188 insertions(+), 171 deletions(-) diff --git a/diff_diff/trop.py b/diff_diff/trop.py index a6e2fdca..d16eaa15 100644 --- a/diff_diff/trop.py +++ b/diff_diff/trop.py @@ -45,8 +45,10 @@ from diff_diff.utils import compute_confidence_interval, compute_p_value -# Sentinel value for "disabled" mode in LOOCV parameter search -# Following paper's footnote 2: λ=∞ disables the corresponding component +# Sentinel value for "disabled" λ_nn in LOOCV parameter search. +# Per paper's footnote 2: λ_nn=∞ disables the factor model (L=0). +# For λ_time and λ_unit, 0.0 means disabled (uniform weights) per Eq. 3: +# exp(-0 × dist) = 1 for all distances. _LAMBDA_INF: float = float('inf') @@ -116,15 +118,14 @@ class TROPResults: treatment_effects : dict Individual treatment effects for each treated (unit, time) pair. lambda_time : float - Selected time weight decay parameter from grid. Note: infinity values - are converted internally (∞ → 0.0 for uniform weights) for computation. + Selected time weight decay parameter from grid. 0.0 = uniform time + weights (disabled) per Eq. 3. lambda_unit : float - Selected unit weight decay parameter from grid. Note: infinity values - are converted internally (∞ → 0.0 for uniform weights) for computation. + Selected unit weight decay parameter from grid. 0.0 = uniform unit + weights (disabled) per Eq. 3. lambda_nn : float - Selected nuclear norm regularization parameter from grid. Note: infinity - values are converted internally (∞ → 1e10, factor model disabled) for - computation. + Selected nuclear norm regularization parameter from grid. inf = factor + model disabled (L=0); converted to 1e10 internally for computation. factor_matrix : np.ndarray Estimated low-rank factor matrix L (n_periods x n_units). effective_rank : float @@ -382,11 +383,14 @@ class TROP: penalty is finite. lambda_time_grid : list, optional - Grid of time weight decay parameters. Default: [0, 0.1, 0.5, 1, 2, 5]. + Grid of time weight decay parameters. 0.0 = uniform weights (disabled). + Must not contain inf. Default: [0, 0.1, 0.5, 1, 2, 5]. lambda_unit_grid : list, optional - Grid of unit weight decay parameters. Default: [0, 0.1, 0.5, 1, 2, 5]. + Grid of unit weight decay parameters. 0.0 = uniform weights (disabled). + Must not contain inf. Default: [0, 0.1, 0.5, 1, 2, 5]. lambda_nn_grid : list, optional - Grid of nuclear norm regularization parameters. Default: [0, 0.01, 0.1, 1]. + Grid of nuclear norm regularization parameters. inf = factor model + disabled (L=0). Default: [0, 0.01, 0.1, 1]. max_iter : int, default=100 Maximum iterations for nuclear norm optimization. tol : float, default=1e-6 @@ -491,6 +495,21 @@ def __init__( f"got '{variance_method}'" ) + # Validate that time/unit grids do not contain inf. + # Per Athey et al. (2025) Eq. 3, λ_time=0 and λ_unit=0 give uniform + # weights (exp(-0 × dist) = 1). Using inf is a misunderstanding of + # the paper's convention. Only λ_nn=∞ is valid (disables factor model). + for grid_name, grid_vals in [ + ("lambda_time_grid", self.lambda_time_grid), + ("lambda_unit_grid", self.lambda_unit_grid), + ]: + if any(np.isinf(v) for v in grid_vals): + raise ValueError( + f"{grid_name} must not contain inf. Use 0.0 for uniform " + f"weights (disabled) per Athey et al. (2025) Eq. 3: " + f"exp(-0 × dist) = 1 for all distances." + ) + # Internal state self.results_: Optional[TROPResults] = None self.is_fitted_: bool = False @@ -708,10 +727,11 @@ def _univariate_loocv_search( Following paper's footnote 2, this performs a univariate grid search for one tuning parameter while holding others fixed. The fixed_params - can include _LAMBDA_INF values to disable specific components: + use 0.0 for disabled time/unit weights and _LAMBDA_INF for disabled + factor model: - lambda_nn = inf: Skip nuclear norm regularization (L=0) - - lambda_time = inf: Uniform time weights (treated as 0) - - lambda_unit = inf: Uniform unit weights (treated as 0) + - lambda_time = 0.0: Uniform time weights (exp(-0×dist)=1) + - lambda_unit = 0.0: Uniform unit weights (exp(-0×dist)=1) Parameters ---------- @@ -732,7 +752,7 @@ def _univariate_loocv_search( grid : List[float] Grid of values to search over. fixed_params : Dict[str, float] - Fixed values for other parameters. May include _LAMBDA_INF. + Fixed values for other parameters. May include _LAMBDA_INF for lambda_nn. Returns ------- @@ -745,22 +765,14 @@ def _univariate_loocv_search( for value in grid: params = {**fixed_params, param_name: value} - # Convert inf values to 0 for computation (inf means "disabled" = uniform weights) lambda_time = params.get('lambda_time', 0.0) lambda_unit = params.get('lambda_unit', 0.0) lambda_nn = params.get('lambda_nn', 0.0) - # Handle infinity as "disabled" mode - # Per paper Equations 2-3: - # - λ_time/λ_unit=∞ → exp(-∞×dist)→0 for dist>0, uniform weights → use 0.0 - # - λ_nn=∞ → infinite penalty → L≈0 (factor model disabled) → use 1e10 - # Note: λ_nn=0 means NO regularization (full-rank L), opposite of "disabled" - if np.isinf(lambda_time): - lambda_time = 0.0 # Uniform time weights - if np.isinf(lambda_unit): - lambda_unit = 0.0 # Uniform unit weights + # Convert λ_nn=∞ → large finite value (factor model disabled, L≈0) + # λ_time and λ_unit use 0.0 for uniform weights per Eq. 3 (no inf conversion needed) if np.isinf(lambda_nn): - lambda_nn = 1e10 # Very large → L≈0 (factor model disabled) + lambda_nn = 1e10 try: score = self._loocv_score_obs_specific( @@ -1423,9 +1435,9 @@ def _fit_joint( for lambda_time_val in self.lambda_time_grid: for lambda_unit_val in self.lambda_unit_grid: for lambda_nn_val in self.lambda_nn_grid: - # Convert infinity values - lt = 0.0 if np.isinf(lambda_time_val) else lambda_time_val - lu = 0.0 if np.isinf(lambda_unit_val) else lambda_unit_val + # Convert λ_nn=∞ → large finite value (factor model disabled) + lt = lambda_time_val + lu = lambda_unit_val ln = 1e10 if np.isinf(lambda_nn_val) else lambda_nn_val try: @@ -1451,13 +1463,10 @@ def _fit_joint( # Final estimation with best parameters lambda_time, lambda_unit, lambda_nn = best_lambda - original_lambda_time, original_lambda_unit, original_lambda_nn = best_lambda + original_lambda_nn = lambda_nn - # Convert infinity values for computation - if np.isinf(lambda_time): - lambda_time = 0.0 - if np.isinf(lambda_unit): - lambda_unit = 0.0 + # Convert λ_nn=∞ → large finite value (factor model disabled, L≈0) + # λ_time and λ_unit use 0.0 for uniform weights directly (no conversion needed) if np.isinf(lambda_nn): lambda_nn = 1e10 @@ -1535,8 +1544,8 @@ def _fit_joint( unit_effects=unit_effects_dict, time_effects=time_effects_dict, treatment_effects=treatment_effects, - lambda_time=original_lambda_time, - lambda_unit=original_lambda_unit, + lambda_time=lambda_time, + lambda_unit=lambda_unit, lambda_nn=original_lambda_nn, factor_matrix=L, effective_rank=effective_rank, @@ -2053,11 +2062,11 @@ def fit( {'lambda_unit': 0.0, 'lambda_nn': _LAMBDA_INF} ) - # λ_nn search: fix λ_time=∞ (uniform time weights), λ_unit=0 + # λ_nn search: fix λ_time=0 (uniform time weights), λ_unit=0 lambda_nn_init, _ = self._univariate_loocv_search( Y, D, control_mask, control_unit_idx, n_units, n_periods, 'lambda_nn', self.lambda_nn_grid, - {'lambda_time': _LAMBDA_INF, 'lambda_unit': 0.0} + {'lambda_time': 0.0, 'lambda_unit': 0.0} ) # λ_unit search: fix λ_nn=∞, λ_time=0 @@ -2099,24 +2108,16 @@ def fit( self._optimal_lambda = best_lambda lambda_time, lambda_unit, lambda_nn = best_lambda - # Convert infinity values for final estimation (matching LOOCV conversion) - # This ensures final estimation uses the same effective parameters that LOOCV evaluated. - # See REGISTRY.md "λ=∞ implementation" for rationale. - # - # IMPORTANT: Store original grid values for results, use converted for computation. - # This lets users see what was selected from their grid, while ensuring consistent - # behavior between point estimation and variance estimation. - original_lambda_time, original_lambda_unit, original_lambda_nn = best_lambda - - if np.isinf(lambda_time): - lambda_time = 0.0 # Uniform time weights - if np.isinf(lambda_unit): - lambda_unit = 0.0 # Uniform unit weights + # Store original λ_nn for results (only λ_nn needs original→effective conversion). + # λ_time and λ_unit use 0.0 for uniform weights directly per Eq. 3. + original_lambda_nn = lambda_nn + + # Convert λ_nn=∞ → large finite value (factor model disabled, L≈0) if np.isinf(lambda_nn): - lambda_nn = 1e10 # Very large → L≈0 (factor model disabled) + lambda_nn = 1e10 - # Create effective_lambda with converted values for ALL downstream computation - # This ensures variance estimation uses the same parameters as point estimation + # effective_lambda with converted λ_nn for ALL downstream computation + # (variance estimation uses the same parameters as point estimation) effective_lambda = (lambda_time, lambda_unit, lambda_nn) # Step 2: Final estimation - per-observation model fitting following Algorithm 2 @@ -2214,10 +2215,8 @@ def fit( unit_effects=unit_effects_dict, time_effects=time_effects_dict, treatment_effects=treatment_effects, - # Store ORIGINAL grid values (possibly inf) so users see what was selected. - # Internally, infinity values are converted for computation (see effective_lambda). - lambda_time=original_lambda_time, - lambda_unit=original_lambda_unit, + lambda_time=lambda_time, + lambda_unit=lambda_unit, lambda_nn=original_lambda_nn, factor_matrix=L_hat, effective_rank=effective_rank, diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 637556ec..7b337a3c 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -488,30 +488,40 @@ has D=1), ATT will be incorrect - document this clearly. - No separate "post_periods" concept - D matrix is the sole input for treatment timing - Supports general assignment patterns including staggered adoption -*Estimator equation (as implemented):* +*Estimator equation (as implemented, Section 2.2):* -Factor model: +Working model (separating unit/time FE from regularized factor component): ``` -Y_it = L_it + τ D_it + ε_it +Y_it(0) = α_i + β_t + L_it + ε_it, E[ε_it | L] = 0 ``` -where L = UΣV' is low-rank factor structure. +where α_i are unit fixed effects, β_t are time fixed effects, and L = UΣV' is a low-rank +factor structure. The FE are estimated separately from L because L is regularized but +the fixed effects are not. -Factor estimation via nuclear norm regularization: +Optimization (Equation 2): ``` -L̂ = argmin_L ||Y_control - L||_F² + λ_nn ||L||_* +(α̂, β̂, L̂) = argmin_{α,β,L} Σ_j Σ_s θ_s^{i,t} ω_j^{i,t} (1-W_js)(Y_js - α_j - β_s - L_js)² + λ_nn ||L||_* ``` -Solved via soft-thresholding of singular values: +Solved via alternating minimization with soft-thresholding of singular values for L: ``` L̂ = U × soft_threshold(Σ, λ_nn) × V' ``` -Unit weights: +Per-observation weights (Equation 3): ``` -ω_j = exp(-λ_unit × d(j, treated)) / Σ_k exp(-λ_unit × d(k, treated)) +θ_s^{i,t}(λ) = exp(-λ_time × |t - s|) + +ω_j^{i,t}(λ) = exp(-λ_unit × dist^unit_{-t}(j, i)) + +dist^unit_{-t}(j, i) = (Σ_u 1{u≠t}(1-W_iu)(1-W_ju)(Y_iu - Y_ju)² / Σ_u 1{u≠t}(1-W_iu)(1-W_ju))^{1/2} ``` -where d(j, treated) is RMSE distance to treated units in pre-period. +Note: weights are per-(i,t) observation-specific. The distance formula excludes the +target period t and uses only periods where both units are untreated (W=0). -Time weights: analogous construction for periods. +*Special cases (Section 2.2):* +- λ_nn=∞, ω_j=θ_s=1 (uniform weights) → recovers DID/TWFE +- ω_j=θ_s=1, λ_nn<∞ → recovers Matrix Completion (Athey et al. 2021) +- λ_nn=∞ with specific ω_j, θ_s → recovers SC/SDID *LOOCV tuning parameter selection (Equation 5, Footnote 2):* ``` @@ -521,13 +531,15 @@ Q(λ) = Σ_{j,s: D_js=0} [τ̂_js^loocv(λ)]² - **Two-stage procedure** (per paper's footnote 2): - Stage 1: Univariate grid searches with extreme fixed values - λ_time search: fix λ_unit=0, λ_nn=∞ (disabled) - - λ_nn search: fix λ_time=∞ (uniform time weights), λ_unit=0 + - λ_nn search: fix λ_time=0 (uniform time weights), λ_unit=0 - λ_unit search: fix λ_nn=∞, λ_time=0 - Stage 2: Cycling (coordinate descent) until convergence -- **"Disabled" parameter semantics** (per paper Equations 2-3): - - `λ_time=∞` or `λ_unit=∞`: Converts to `0.0` internally → exp(-0×dist)=1 → uniform weights - - `λ_nn=∞`: Converts to `1e10` internally → very large penalty → L≈0 (factor model off, recovers DID/TWFE) +- **"Disabled" parameter semantics** (per paper Section 4.3, Table 5, Footnote 2): + - `λ_time=0`: Uniform time weights (disabled), because exp(-0 × dist) = 1 + - `λ_unit=0`: Uniform unit weights (disabled), because exp(-0 × dist) = 1 + - `λ_nn=∞`: Factor model disabled (L=0), because infinite penalty; converted to `1e10` internally - **Note**: `λ_nn=0` means NO regularization (full-rank L), which is the OPPOSITE of "disabled" + - **Validation**: `lambda_time_grid` and `lambda_unit_grid` must not contain inf. A `ValueError` is raised if they do, guiding users to use 0.0 for uniform weights per Eq. 3. - **Subsampling**: max_loocv_samples (default 100) for computational tractability - This subsamples control observations, NOT parameter combinations - Increases precision at cost of computation; increase for more precise tuning @@ -538,21 +550,20 @@ Q(λ) = Σ_{j,s: D_js=0} [τ̂_js^loocv(λ)]² - This ensures λ selection only considers fully estimable combinations *Standard errors:* -- Default: Block bootstrap preserving panel structure -- Alternative: Jackknife (leave-one-unit-out) +- Default: Block bootstrap preserving panel structure (Algorithm 3) +- Alternative: Jackknife (leave-one-unit-out) — **implementation addition** not described in the paper *Edge cases:* - Rank selection: automatic via cross-validation, information criterion, or elbow - Zero singular values: handled by soft-thresholding - Extreme distances: weights regularized to prevent degeneracy - LOOCV fit failures: returns Q(λ) = ∞ on first failure (per Equation 5 requirement that Q sums over ALL D==0 cells); if all parameter combinations fail, falls back to defaults (1.0, 1.0, 0.1) -- **λ=∞ implementation**: Infinity values are converted in both LOOCV search and final estimation: - - λ_time=∞ or λ_unit=∞ → 0.0 (uniform weights via exp(-0×d)=1) +- **λ_nn=∞ implementation**: Only λ_nn uses infinity (converted to 1e10 for computation): - λ_nn=∞ → 1e10 (large penalty → L≈0, factor model disabled) - Conversion applied to grid values during LOOCV (including Rust backend) - Conversion applied to selected values for point estimation - Conversion applied to selected values for variance estimation (ensures SE matches ATT) - - **Results storage**: `TROPResults` stores *original* grid values (e.g., inf), while computations use converted values. This lets users see what was selected from their grid. + - **Results storage**: `TROPResults` stores *original* λ_nn value (inf), while computations use 1e10. λ_time and λ_unit store their selected values directly (0.0 = uniform). - **Empty control observations**: If LOOCV control observations become empty (edge case during subsampling), returns Q(λ) = ∞ with warning. A score of 0.0 would incorrectly "win" over legitimate parameters. - **Infinite LOOCV score handling**: If best LOOCV score is infinite, `best_lambda` is set to None, triggering defaults fallback - Validation: requires at least 2 periods before first treatment @@ -572,7 +583,7 @@ Q(λ) = Σ_{j,s: D_js=0} [τ̂_js^loocv(λ)]² **Requirements checklist:** - [x] Factor matrix estimated via soft-threshold SVD -- [x] Unit weights: `exp(-λ_unit × distance)` with normalization +- [x] Unit weights: `exp(-λ_unit × distance)` (unnormalized, matching Eq. 2) - [x] LOOCV implemented for tuning parameter selection - [x] LOOCV uses SUM of squared errors per Equation 5 - [x] Multiple rank selection methods: cv, ic, elbow diff --git a/rust/src/trop.rs b/rust/src/trop.rs index 9465de7d..cc82b65f 100644 --- a/rust/src/trop.rs +++ b/rust/src/trop.rs @@ -188,34 +188,26 @@ fn univariate_loocv_search( let results: Vec<(f64, f64)> = grid .par_iter() .map(|&value| { - // Set parameters, converting inf for "disabled" mode - // Per paper Equations 2-3: - // - λ_time/λ_unit=∞ → uniform weights → use 0.0 - // - λ_nn=∞ → infinite penalty → L≈0 (factor model disabled) → use 1e10 - // Note: λ_nn=0 means NO regularization (full-rank L), opposite of "disabled" - // - // IMPORTANT: Convert the grid value BEFORE using it, matching Python behavior. - // This ensures Rust and Python evaluate the same objective for infinity grids. + // Convert λ_nn=∞ → 1e10 (factor model disabled, L≈0). + // λ_time and λ_unit use 0.0 for uniform weights per Eq. 3 (no inf conversion). let (lambda_time, lambda_unit, lambda_nn) = match param_type { 0 => { - // Searching λ_time: convert grid value if infinite - let value_converted = if value.is_infinite() { 0.0 } else { value }; - (value_converted, - if fixed_unit.is_infinite() { 0.0 } else { fixed_unit }, + // Searching λ_time: use grid value directly (no inf expected) + (value, + fixed_unit, if fixed_nn.is_infinite() { 1e10 } else { fixed_nn }) }, 1 => { - // Searching λ_unit: convert grid value if infinite - let value_converted = if value.is_infinite() { 0.0 } else { value }; - (if fixed_time.is_infinite() { 0.0 } else { fixed_time }, - value_converted, + // Searching λ_unit: use grid value directly (no inf expected) + (fixed_time, + value, if fixed_nn.is_infinite() { 1e10 } else { fixed_nn }) }, _ => { - // Searching λ_nn: convert grid value if infinite + // Searching λ_nn: convert inf → 1e10 (factor model disabled) let value_converted = if value.is_infinite() { 1e10 } else { value }; - (if fixed_time.is_infinite() { 0.0 } else { fixed_time }, - if fixed_unit.is_infinite() { 0.0 } else { fixed_unit }, + (fixed_time, + fixed_unit, value_converted) }, }; @@ -364,10 +356,10 @@ pub fn loocv_grid_search<'py>( &lambda_time_vec, 0.0, 0.0, f64::INFINITY, 0, max_iter, tol, ); - // λ_nn search: fix λ_time=∞ (disabled), λ_unit=0 + // λ_nn search: fix λ_time=0 (uniform time weights), λ_unit=0 let (lambda_nn_init, _) = univariate_loocv_search( &y_arr, &d_arr, &control_mask_arr, &time_dist_arr, &control_obs, - &lambda_nn_vec, f64::INFINITY, 0.0, 0.0, 2, max_iter, tol, + &lambda_nn_vec, 0.0, 0.0, 0.0, 2, max_iter, tol, ); // λ_unit search: fix λ_nn=∞, λ_time=0 @@ -384,13 +376,9 @@ pub fn loocv_grid_search<'py>( max_iter, tol, 10, ); - // Convert infinity values BEFORE computing final score (Issue 1 fix) - // Per paper Equations 2-3: - // - λ_time/λ_unit=∞ → uniform weights → use 0.0 - // - λ_nn=∞ → infinite penalty → L≈0 (factor model disabled) → use 1e10 - // This ensures final score computation matches what LOOCV evaluated. - let best_time_eff = if best_time.is_infinite() { 0.0 } else { best_time }; - let best_unit_eff = if best_unit.is_infinite() { 0.0 } else { best_unit }; + // Convert λ_nn=∞ → 1e10 for final score computation (factor model disabled) + let best_time_eff = best_time; + let best_unit_eff = best_unit; let best_nn_eff = if best_nn.is_infinite() { 1e10 } else { best_nn }; // Compute final score with converted values @@ -1522,9 +1510,9 @@ pub fn loocv_grid_search_joint<'py>( let results: Vec<(f64, f64, f64, f64, usize, Option<(usize, usize)>)> = grid_combinations .into_par_iter() .map(|(lt, lu, ln)| { - // Convert infinity values - let lt_eff = if lt.is_infinite() { 0.0 } else { lt }; - let lu_eff = if lu.is_infinite() { 0.0 } else { lu }; + // Convert λ_nn=∞ → 1e10 (factor model disabled) + let lt_eff = lt; + let lu_eff = lu; let ln_eff = if ln.is_infinite() { 1e10 } else { ln }; let (score, n_valid, first_failed) = loocv_score_joint( @@ -1629,9 +1617,9 @@ pub fn bootstrap_trop_variance_joint<'py>( } let treated_periods = n_periods.saturating_sub(first_treat_period); - // Convert infinity values for computation - let lt_eff = if lambda_time.is_infinite() { 0.0 } else { lambda_time }; - let lu_eff = if lambda_unit.is_infinite() { 0.0 } else { lambda_unit }; + // Convert λ_nn=∞ → 1e10 (factor model disabled) + let lt_eff = lambda_time; + let lu_eff = lambda_unit; let ln_eff = if lambda_nn.is_infinite() { 1e10 } else { lambda_nn }; // Run bootstrap iterations in parallel diff --git a/tests/test_trop.py b/tests/test_trop.py index eb38e3a0..111dfe5a 100644 --- a/tests/test_trop.py +++ b/tests/test_trop.py @@ -2121,25 +2121,20 @@ def always_infinity(*args, **kwargs): assert results.lambda_nn == 0.1, \ f"Expected default lambda_nn=0.1, got {results.lambda_nn}" - def test_infinity_grid_values_handled_consistently(self, simple_panel_data): + def test_uniform_weights_and_disabled_factor_handled_consistently(self, simple_panel_data): """ - Test that infinity in grids is handled consistently in LOOCV and final estimation. + Test that 0.0 (uniform weights) and inf (disabled factor) are handled + consistently in LOOCV and final estimation. - When infinity is in the parameter grid: - - LOOCV converts it for computation (λ_time=∞→0, λ_unit=∞→0, λ_nn=∞→1e10) - - LOOCV returns the original grid value (inf) if it was best - - Final estimation must also convert infinity to match LOOCV behavior - - This test ensures the conversion in final estimation matches LOOCV. + Per Athey et al. (2025) Eq. 3: + - λ_time=0.0 → uniform time weights (exp(-0×dist)=1) + - λ_unit=0.0 → uniform unit weights (exp(-0×dist)=1) + - λ_nn=∞ → factor model disabled (L=0), converted to 1e10 internally """ - # Create estimator with infinity in grids - # Use grids where infinity is likely to be selected: - # - lambda_time_grid: [inf] forces selection of inf - # - lambda_nn_grid: [inf] forces selection of inf - trop_est = TROP( - lambda_time_grid=[np.inf], # Only inf available → must be selected - lambda_unit_grid=[0.0], # Normal value - lambda_nn_grid=[np.inf], # Only inf available → must be selected + trop_est = TROP( + lambda_time_grid=[0.0], # Uniform time weights (disabled) + lambda_unit_grid=[0.0], # Uniform unit weights (disabled) + lambda_nn_grid=[np.inf], # Factor model disabled → converted to 1e10 n_bootstrap=5, seed=42 ) @@ -2152,9 +2147,9 @@ def test_infinity_grid_values_handled_consistently(self, simple_panel_data): time="period", ) - # ATT should be finite (no NaN/inf from unconverted infinity parameters) + # ATT should be finite assert np.isfinite(results.att), ( - f"ATT should be finite when infinity params are converted, got {results.att}" + f"ATT should be finite with uniform weights and no factor model, got {results.att}" ) # SE should be finite or at least non-negative @@ -2162,29 +2157,53 @@ def test_infinity_grid_values_handled_consistently(self, simple_panel_data): f"SE should be finite, got {results.se}" ) - # The stored lambda values should be the original grid values (inf) - # because we store what was selected, but conversion happens internally - # (This documents current behavior; the key is that ATT is finite) - assert np.isinf(results.lambda_time) or results.lambda_time == 0.0, ( - f"lambda_time should be inf (stored) or 0.0 (if converted for storage)" + # lambda_time and lambda_unit should be 0.0 (uniform weights) + assert results.lambda_time == 0.0, ( + f"lambda_time should be 0.0 (uniform weights), got {results.lambda_time}" ) - assert np.isinf(results.lambda_nn) or results.lambda_nn == 1e10, ( - f"lambda_nn should be inf (stored) or 1e10 (if converted for storage)" + # lambda_nn should store the original inf value + assert np.isinf(results.lambda_nn), ( + f"lambda_nn should be inf (original grid value), got {results.lambda_nn}" ) + def test_inf_in_time_unit_grids_raises_valueerror(self): + """ + Test that inf in lambda_time_grid or lambda_unit_grid raises ValueError. + + Per Athey et al. (2025) Eq. 3, λ_time=0 and λ_unit=0 give uniform + weights. Using inf is a misunderstanding; only λ_nn=∞ is valid. + """ + import pytest + + # inf in lambda_time_grid should raise + with pytest.raises(ValueError, match="lambda_time_grid must not contain inf"): + TROP(lambda_time_grid=[np.inf]) + + with pytest.raises(ValueError, match="lambda_time_grid must not contain inf"): + TROP(lambda_time_grid=[0.0, np.inf, 1.0]) + + # inf in lambda_unit_grid should raise + with pytest.raises(ValueError, match="lambda_unit_grid must not contain inf"): + TROP(lambda_unit_grid=[np.inf]) + + with pytest.raises(ValueError, match="lambda_unit_grid must not contain inf"): + TROP(lambda_unit_grid=[0.5, np.inf]) + + # inf in lambda_nn_grid should still be valid + trop_est = TROP(lambda_nn_grid=[np.inf]) + assert np.inf in trop_est.lambda_nn_grid + def test_variance_estimation_uses_converted_params(self, simple_panel_data): """ Test that variance estimation uses the same converted parameters as point estimation. - When infinity is in the grid and gets selected, both ATT and SE should be - computed with the same effective parameters (e.g., λ_time=∞ converted to 0.0). - This test verifies the fix for variance estimation inconsistency (PR #110 Round 7). + λ_nn=∞ is converted to 1e10 for computation. λ_time and λ_unit use 0.0 + directly for uniform weights (no conversion needed). """ from unittest.mock import patch - # Use grids with only infinity values to force selection trop_est = TROP( - lambda_time_grid=[np.inf], # Will be converted to 0.0 internally + lambda_time_grid=[0.0], # Uniform time weights (paper convention) lambda_unit_grid=[0.0], lambda_nn_grid=[np.inf], # Will be converted to 1e10 internally n_bootstrap=5, @@ -2210,19 +2229,20 @@ def tracking_fit(self, data, outcome, treatment, unit, time, fixed_lambda): time="period", ) - # Results should store original grid values - assert np.isinf(results.lambda_time), "Results should store original infinity value" - assert np.isinf(results.lambda_nn), "Results should store original infinity value" + # Results should store 0.0 for time (direct value, no conversion) + assert results.lambda_time == 0.0, "lambda_time should be 0.0" + # Results should store original inf for lambda_nn + assert np.isinf(results.lambda_nn), "Results should store original infinity value for lambda_nn" # ATT should be finite (computed with converted params) assert np.isfinite(results.att), "ATT should be finite" # Variance estimation should have received converted parameters - # Check that bootstrap iterations used converted (non-infinite) values + # Check that bootstrap iterations used converted (non-infinite) λ_nn values for captured in captured_lambda: lambda_time, lambda_unit, lambda_nn = captured - assert not np.isinf(lambda_time), ( - f"Bootstrap should receive converted λ_time=0.0, not {lambda_time}" + assert lambda_time == 0.0, ( + f"Bootstrap should receive λ_time=0.0, got {lambda_time}" ) assert not np.isinf(lambda_nn), ( f"Bootstrap should receive converted λ_nn=1e10, not {lambda_nn}" @@ -2294,16 +2314,15 @@ def test_empty_control_obs_returns_infinity(self, simple_panel_data): def test_original_grid_values_stored_in_results(self, simple_panel_data): """ - Test that TROPResults stores the original grid values, not converted ones. + Test that TROPResults stores the selected grid values correctly. - Per the design decision in PR #110 Round 7, results should store the - original grid values (possibly inf) so users can see what was selected, - while internally using converted values for computation. + λ_time and λ_unit store values directly (0.0 = uniform weights). + λ_nn stores the original inf value when factor model is disabled. """ trop_est = TROP( - lambda_time_grid=[np.inf], # Original value: inf + lambda_time_grid=[0.0], # Uniform time weights lambda_unit_grid=[0.5], - lambda_nn_grid=[np.inf], # Original value: inf + lambda_nn_grid=[np.inf], # Factor model disabled (original: inf) n_bootstrap=5, seed=42 ) @@ -2316,13 +2335,14 @@ def test_original_grid_values_stored_in_results(self, simple_panel_data): time="period", ) - # Results should store original grid values (inf) - assert np.isinf(results.lambda_time), ( - f"results.lambda_time should be inf (original), got {results.lambda_time}" + # lambda_time stores selected value directly (0.0 = uniform) + assert results.lambda_time == 0.0, ( + f"results.lambda_time should be 0.0, got {results.lambda_time}" ) assert results.lambda_unit == 0.5, ( f"results.lambda_unit should be 0.5, got {results.lambda_unit}" ) + # lambda_nn stores original inf (converted to 1e10 only for computation) assert np.isinf(results.lambda_nn), ( f"results.lambda_nn should be inf (original), got {results.lambda_nn}" ) @@ -2527,22 +2547,22 @@ def test_unbalanced_panel_multiple_missing_periods(self): assert results is not None assert np.isfinite(results.att) - def test_infinity_grid_values_with_final_score_computation(self, simple_panel_data): - """Test that infinity grid values are properly converted for final score. + def test_mixed_grid_values_with_final_score_computation(self, simple_panel_data): + """Test that grid values including 0.0 (uniform) and inf (λ_nn) work for final score. - Issue 1 fix: When LOOCV selects infinity values from the grid, the - final score computation should use converted values (0.0 or 1e10), - not raw infinity. + When LOOCV selects λ_nn=∞, the final score computation should use + converted value (1e10), not raw infinity. λ_time and λ_unit grids + use finite values only (0.0 = uniform weights per Eq. 3). """ trop_est = TROP( - lambda_time_grid=[np.inf, 0.5], # inf should convert to 0.0 - lambda_unit_grid=[np.inf, 0.5], # inf should convert to 0.0 + lambda_time_grid=[0.0, 0.5], # 0.0 = uniform time weights + lambda_unit_grid=[0.0, 0.5], # 0.0 = uniform unit weights lambda_nn_grid=[np.inf, 0.1], # inf should convert to 1e10 n_bootstrap=5, seed=42 ) - # This should complete without error, even if inf values are selected + # This should complete without error results = trop_est.fit( simple_panel_data, outcome="outcome", @@ -2552,11 +2572,10 @@ def test_infinity_grid_values_with_final_score_computation(self, simple_panel_da ) # ATT should be finite regardless of which grid values were selected - assert np.isfinite(results.att), "ATT should be finite with inf grid values" + assert np.isfinite(results.att), "ATT should be finite with mixed grid values" assert results.se >= 0, "SE should be non-negative" - # If inf values were selected, LOOCV score should still be computed correctly - # (using converted values, not raw inf) + # If inf was selected for λ_nn, LOOCV score should still be computed correctly if np.isinf(results.loocv_score): # Infinite LOOCV score is acceptable (means fits failed) # but ATT should still be finite (falls back to defaults) From b350d659d4d391fb6942075be01e7a3238074e3b Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 31 Jan 2026 18:32:04 -0500 Subject: [PATCH 2/3] Add inf validation to Rust TROP entrypoints and fix docstring Defense-in-depth: validate lambda_time_grid and lambda_unit_grid don't contain inf in loocv_grid_search and loocv_grid_search_joint Rust entrypoints, matching the Python-side validation. Fix univariate_loocv_search docstring to use paper conventions (0.0 for uniform weights, inf only for disabling factor model). Co-Authored-By: Claude Opus 4.5 --- rust/src/trop.rs | 40 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 37 insertions(+), 3 deletions(-) diff --git a/rust/src/trop.rs b/rust/src/trop.rs index cc82b65f..b7be5f15 100644 --- a/rust/src/trop.rs +++ b/rust/src/trop.rs @@ -158,9 +158,9 @@ fn compute_pair_distance( /// * `time_dist` - Time distance matrix /// * `control_obs` - List of control observations for LOOCV /// * `grid` - Grid of values to search -/// * `fixed_time` - Fixed lambda_time (inf for disabled) -/// * `fixed_unit` - Fixed lambda_unit (inf for disabled) -/// * `fixed_nn` - Fixed lambda_nn (inf for disabled) +/// * `fixed_time` - Fixed lambda_time (0.0 for uniform weights) +/// * `fixed_unit` - Fixed lambda_unit (0.0 for uniform weights) +/// * `fixed_nn` - Fixed lambda_nn (inf to disable factor model) /// * `param_type` - Which parameter to search: 0=time, 1=unit, 2=nn /// * `max_iter` - Maximum iterations /// * `tol` - Convergence tolerance @@ -339,6 +339,23 @@ pub fn loocv_grid_search<'py>( let lambda_unit_vec: Vec = lambda_unit_grid.as_array().to_vec(); let lambda_nn_vec: Vec = lambda_nn_grid.as_array().to_vec(); + // Validate: lambda_time_grid and lambda_unit_grid must not contain inf. + // Per Athey et al. (2025) Eq. 3: use 0.0 for uniform weights, not inf. + for &v in &lambda_time_vec { + if v.is_infinite() { + return Err(pyo3::exceptions::PyValueError::new_err( + "lambda_time_grid must not contain inf. Use 0.0 for uniform weights (disabled) per Athey et al. (2025) Eq. 3." + )); + } + } + for &v in &lambda_unit_vec { + if v.is_infinite() { + return Err(pyo3::exceptions::PyValueError::new_err( + "lambda_unit_grid must not contain inf. Use 0.0 for uniform weights (disabled) per Athey et al. (2025) Eq. 3." + )); + } + } + // Get control observations for LOOCV let control_obs = get_control_observations( &y_arr, @@ -1477,6 +1494,23 @@ pub fn loocv_grid_search_joint<'py>( let lambda_unit_vec: Vec = lambda_unit_grid.as_array().to_vec(); let lambda_nn_vec: Vec = lambda_nn_grid.as_array().to_vec(); + // Validate: lambda_time_grid and lambda_unit_grid must not contain inf. + // Per Athey et al. (2025) Eq. 3: use 0.0 for uniform weights, not inf. + for &v in &lambda_time_vec { + if v.is_infinite() { + return Err(pyo3::exceptions::PyValueError::new_err( + "lambda_time_grid must not contain inf. Use 0.0 for uniform weights (disabled) per Athey et al. (2025) Eq. 3." + )); + } + } + for &v in &lambda_unit_vec { + if v.is_infinite() { + return Err(pyo3::exceptions::PyValueError::new_err( + "lambda_unit_grid must not contain inf. Use 0.0 for uniform weights (disabled) per Athey et al. (2025) Eq. 3." + )); + } + } + let n_periods = y_arr.nrows(); let n_units = y_arr.ncols(); From 5e7f571e0a9558c7c7e960add3661aee2e14144c Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 31 Jan 2026 18:44:22 -0500 Subject: [PATCH 3/3] Fix TROP.fit() docstring to reflect inf validation for lambda grids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Returns docstring still claimed λ_time/λ_unit=∞ were converted to 0.0 internally. This contradicts the constructor validation added in this PR which now raises on inf for those grids. Updated to accurately state that inf is not accepted for λ_time/λ_unit while λ_nn=∞ is still converted to 1e10. Co-Authored-By: Claude Opus 4.5 --- diff_diff/trop.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/diff_diff/trop.py b/diff_diff/trop.py index d16eaa15..8c7ee111 100644 --- a/diff_diff/trop.py +++ b/diff_diff/trop.py @@ -1875,9 +1875,9 @@ def fit( TROPResults Object containing the ATT estimate, standard error, factor estimates, and tuning parameters. The lambda_* - attributes show the selected grid values. Infinity values - (∞) are converted internally: λ_time/λ_unit=∞ → 0.0 (uniform - weights), λ_nn=∞ → 1e10 (factor model disabled). + attributes show the selected grid values. For λ_time and + λ_unit, 0.0 means uniform weights; inf is not accepted. + For λ_nn, ∞ is converted to 1e10 (factor model disabled). """ # Validate inputs required_cols = [outcome, treatment, unit, time]