diff --git a/diff_diff/trop.py b/diff_diff/trop.py index a6e2fdca..8c7ee111 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, @@ -1866,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] @@ -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..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 @@ -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) }, }; @@ -347,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, @@ -364,10 +373,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 +393,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 @@ -1489,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(); @@ -1522,9 +1544,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 +1651,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)