From 95cf1ef1c36035acabb01347ccfdaa1e22ebac00 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 13:24:28 -0400 Subject: [PATCH 1/6] Add UserWarning emissions for 8 silent operations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add warnings for: lstsq→pinv fallback, NaN y_tilde masking, always-treated survey weight note, consolidated (g,t) cell skip, missing treatment fill, Rust→Python fallback (4 sites), weight normalization, and inf→0 never-treated conversion. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/staggered.py | 156 ++++-- diff_diff/survey.py | 10 +- diff_diff/trop.py | 16 + diff_diff/trop_global.py | 30 ++ diff_diff/trop_local.py | 7 + diff_diff/two_stage.py | 13 +- tests/test_staggered.py | 244 +++++++-- tests/test_survey.py | 38 ++ tests/test_trop.py | 1092 ++++++++++++++++++++++---------------- tests/test_two_stage.py | 121 ++++- 10 files changed, 1141 insertions(+), 586 deletions(-) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 3e6b04d60..c9e8598aa 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -320,8 +320,7 @@ def __init__( raise ValueError(f"epv_threshold must be > 0, got {epv_threshold}") if pscore_fallback not in ["error", "unconditional"]: raise ValueError( - f"pscore_fallback must be 'error' or 'unconditional', " - f"got '{pscore_fallback}'" + f"pscore_fallback must be 'error' or 'unconditional', " f"got '{pscore_fallback}'" ) # Handle bootstrap_weight_type deprecation @@ -429,30 +428,53 @@ def diagnose_propensity( if self.estimation_method == "reg": return pd.DataFrame( columns=[ - "group", "n_treated", "n_control", - "n_covariates", "n_params", "epv", "status", + "group", + "n_treated", + "n_control", + "n_covariates", + "n_params", + "epv", + "status", ] ) if not covariates: return pd.DataFrame( columns=[ - "group", "n_treated", "n_control", - "n_covariates", "n_params", "epv", "status", + "group", + "n_treated", + "n_control", + "n_covariates", + "n_params", + "epv", + "status", ] ) # Normalize np.inf → 0 for never-treated encoding (same as fit()) df = df.copy() + _inf_mask_diag = df[first_treat].isin([np.inf, float("inf")]) + if _inf_mask_diag.any(): + n_inf_units = df.loc[_inf_mask_diag, unit].nunique() + warnings.warn( + f"{n_inf_units} unit(s) have first_treat=inf; recoding to 0 " + f"(never-treated). Use first_treat=0 to suppress this warning.", + UserWarning, + stacklevel=2, + ) df[first_treat] = df[first_treat].replace([np.inf, float("inf")], 0) # Compute time_periods and treatment_groups (same logic as fit()) time_periods = sorted(df[time].unique()) - treatment_groups = sorted( - [g for g in df[first_treat].unique() if g > 0] - ) + treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0]) precomputed = self._precompute_structures( - df, outcome, unit, time, first_treat, covariates, - time_periods=time_periods, treatment_groups=treatment_groups, + df, + outcome, + unit, + time, + first_treat, + covariates, + time_periods=time_periods, + treatment_groups=treatment_groups, ) cohort_masks = precomputed["cohort_masks"] never_treated_mask = precomputed["never_treated_mask"] @@ -484,15 +506,17 @@ def diagnose_propensity( else: status = "critical" - rows.append({ - "group": g, - "n_treated": n_treated, - "n_control": n_control, - "n_covariates": n_covariates, - "n_params": n_params, - "epv": round(epv, 1), - "status": status, - }) + rows.append( + { + "group": g, + "n_treated": n_treated, + "n_control": n_control, + "n_covariates": n_covariates, + "n_params": n_params, + "epv": round(epv, 1), + "status": status, + } + ) return pd.DataFrame(rows) @@ -847,6 +871,8 @@ def _compute_all_att_gt_vectorized( group_time_effects = {} influence_func_info = {} + skipped_missing_period: List[Tuple] = [] + skipped_empty_cell: List[Tuple] = [] # Collect all valid (g, t, base_col, post_col) tuples tasks = [] @@ -870,6 +896,7 @@ def _compute_all_att_gt_vectorized( base_period_val = g - 1 - self.anticipation if base_period_val not in period_to_col or t not in period_to_col: + skipped_missing_period.append((g, t)) continue tasks.append( @@ -905,6 +932,7 @@ def _compute_all_att_gt_vectorized( n_control = np.sum(control_valid) if n_treated == 0 or n_control == 0: + skipped_empty_cell.append((g, t)) continue treated_change = outcome_change[treated_valid] @@ -1000,7 +1028,11 @@ def _compute_all_att_gt_vectorized( group_time_effects[key]["p_value"] = float(p_values[idx]) group_time_effects[key]["conf_int"] = (float(ci_lowers[idx]), float(ci_uppers[idx])) - return group_time_effects, influence_func_info + skip_info = { + "missing_period": skipped_missing_period, + "empty_cell": skipped_empty_cell, + } + return group_time_effects, influence_func_info, skip_info def _compute_all_att_gt_covariate_reg( self, @@ -1037,6 +1069,8 @@ def _compute_all_att_gt_covariate_reg( ses = [] task_keys = [] n_nan_cells = 0 + skipped_missing_period: List[Tuple] = [] + skipped_empty_cell: List[Tuple] = [] # Collect all valid (g, t) tasks with their base periods tasks_by_group = {} # control_key -> list of (g, t, base_period_val, base_col, post_col) @@ -1059,6 +1093,7 @@ def _compute_all_att_gt_covariate_reg( base_period_val = g - 1 - self.anticipation if base_period_val not in period_to_col or t not in period_to_col: + skipped_missing_period.append((g, t)) continue # Determine control regression grouping key. @@ -1108,6 +1143,7 @@ def _compute_all_att_gt_covariate_reg( # Build X_ctrl with intercept n_c_base = int(np.sum(control_valid_base)) if n_c_base == 0: + skipped_empty_cell.extend((g, t) for g, t, *_ in tasks) continue X_ctrl = None @@ -1176,6 +1212,7 @@ def _compute_all_att_gt_covariate_reg( n_c = int(np.sum(control_valid)) if n_t == 0 or n_c == 0: + skipped_empty_cell.append((g, t)) continue treated_change = outcome_change[treated_valid] @@ -1342,7 +1379,11 @@ def _compute_all_att_gt_covariate_reg( group_time_effects[key]["p_value"] = float(p_values[idx]) group_time_effects[key]["conf_int"] = (float(ci_lowers[idx]), float(ci_uppers[idx])) - return group_time_effects, influence_func_info + skip_info = { + "missing_period": skipped_missing_period, + "empty_cell": skipped_empty_cell, + } + return group_time_effects, influence_func_info, skip_info def fit( self, @@ -1482,7 +1523,16 @@ def fit( # Never-treated indicator (must precede treatment_groups to exclude np.inf) df["_never_treated"] = (df[first_treat] == 0) | (df[first_treat] == np.inf) # Normalize np.inf → 0 so all downstream `> 0` checks exclude never-treated - df.loc[df[first_treat] == np.inf, first_treat] = 0 + _inf_mask = df[first_treat] == np.inf + if _inf_mask.any(): + n_inf_units = df.loc[_inf_mask, unit].nunique() + warnings.warn( + f"{n_inf_units} unit(s) have first_treat=inf; recoding to 0 " + f"(never-treated). Use first_treat=0 to suppress this warning.", + UserWarning, + stacklevel=2, + ) + df.loc[_inf_mask, first_treat] = 0 # Identify groups and time periods time_periods = sorted(df[time].unique()) @@ -1572,6 +1622,8 @@ def fit( min_period = min(time_periods) has_survey = resolved_survey is not None + _skip_info = {"missing_period": [], "empty_cell": []} + if not self.panel: # --- Repeated cross-section path --- # No vectorized/Cholesky fast paths (panel-only optimizations). @@ -1629,8 +1681,10 @@ def fit( elif covariates is None and self.estimation_method == "reg": # Fast vectorized path for the common no-covariates regression case - group_time_effects, influence_func_info = self._compute_all_att_gt_vectorized( - precomputed, treatment_groups, time_periods, min_period + group_time_effects, influence_func_info, _skip_info = ( + self._compute_all_att_gt_vectorized( + precomputed, treatment_groups, time_periods, min_period + ) ) epv_diagnostics = None # No logit in this path elif ( @@ -1640,8 +1694,10 @@ def fit( and not has_survey # Cholesky cache uses X'X; survey needs X'WX ): # Optimized covariate regression path with Cholesky caching - group_time_effects, influence_func_info = self._compute_all_att_gt_covariate_reg( - precomputed, treatment_groups, time_periods, min_period + group_time_effects, influence_func_info, _skip_info = ( + self._compute_all_att_gt_covariate_reg( + precomputed, treatment_groups, time_periods, min_period + ) ) epv_diagnostics = None # No logit in this path else: @@ -1738,6 +1794,24 @@ def fit( stacklevel=2, ) + # Consolidated (g,t) cell skip warning + _n_missing = len(_skip_info.get("missing_period", [])) + _n_empty = len(_skip_info.get("empty_cell", [])) + _n_total_skipped = _n_missing + _n_empty + if _n_total_skipped > 0: + _parts = [] + if _n_missing: + _parts.append( + f"{_n_missing} due to missing base/post period " f"in panel structure" + ) + if _n_empty: + _parts.append(f"{_n_empty} due to zero treated or control " f"observations") + warnings.warn( + f"{_n_total_skipped} (group, time) cell(s) skipped: " f"{'; '.join(_parts)}.", + UserWarning, + stacklevel=2, + ) + # Compute overall ATT (simple aggregation) overall_att, overall_se, overall_effective_df = self._aggregate_simple( group_time_effects, influence_func_info, df, unit, precomputed @@ -2143,15 +2217,10 @@ def _ipw_estimation( # dropped rank-deficient columns to prevent NaN # propagation on cache reuse) alongside EPV diagnostics if pscore_cache is not None and pscore_key is not None: - beta_clean = np.where( - np.isfinite(beta_logistic), beta_logistic, 0.0 - ) + beta_clean = np.where(np.isfinite(beta_logistic), beta_logistic, 0.0) pscore_cache[pscore_key] = (beta_clean, diag) except (np.linalg.LinAlgError, ValueError): - if ( - self.pscore_fallback == "error" - or self.rank_deficient_action == "error" - ): + if self.pscore_fallback == "error" or self.rank_deficient_action == "error": raise # Fallback to unconditional if logistic regression fails ctx = f" for {context_label}" if context_label else "" @@ -2430,15 +2499,10 @@ def _doubly_robust( ) _check_propensity_diagnostics(pscore, self.pscore_trim) if pscore_cache is not None and pscore_key is not None: - beta_clean = np.where( - np.isfinite(beta_logistic), beta_logistic, 0.0 - ) + beta_clean = np.where(np.isfinite(beta_logistic), beta_logistic, 0.0) pscore_cache[pscore_key] = (beta_clean, diag) except (np.linalg.LinAlgError, ValueError): - if ( - self.pscore_fallback == "error" - or self.rank_deficient_action == "error" - ): + if self.pscore_fallback == "error" or self.rank_deficient_action == "error": raise # Fallback to unconditional if logistic regression fails ctx = f" for {context_label}" if context_label else "" @@ -3245,10 +3309,7 @@ def _ipw_estimation_rc( ) _check_propensity_diagnostics(pscore, self.pscore_trim) except (np.linalg.LinAlgError, ValueError): - if ( - self.pscore_fallback == "error" - or self.rank_deficient_action == "error" - ): + if self.pscore_fallback == "error" or self.rank_deficient_action == "error": raise ctx = f" for {context_label}" if context_label else "" warnings.warn( @@ -3513,10 +3574,7 @@ def _doubly_robust_rc( ) _check_propensity_diagnostics(pscore, self.pscore_trim) except (np.linalg.LinAlgError, ValueError): - if ( - self.pscore_fallback == "error" - or self.rank_deficient_action == "error" - ): + if self.pscore_fallback == "error" or self.rank_deficient_action == "error": raise ctx = f" for {context_label}" if context_label else "" warnings.warn( diff --git a/diff_diff/survey.py b/diff_diff/survey.py index bd18700e2..2dc78f69b 100644 --- a/diff_diff/survey.py +++ b/diff_diff/survey.py @@ -192,7 +192,15 @@ def resolve(self, data: pd.DataFrame) -> "ResolvedSurveyDesign": if self.replicate_weights is not None: weights = raw_weights.copy() elif self.weight_type in ("pweight", "aweight"): - weights = raw_weights * (n / np.sum(raw_weights)) + raw_sum = float(np.sum(raw_weights)) + weights = raw_weights * (n / raw_sum) + if not np.isclose(raw_sum, n): + warnings.warn( + f"{self.weight_type} weights normalized to mean=1 " + f"(sum={n}). Original sum was {raw_sum:.4g}.", + UserWarning, + stacklevel=2, + ) else: weights = raw_weights.copy() else: diff --git a/diff_diff/trop.py b/diff_diff/trop.py index 4d26a69f5..a5976a13c 100644 --- a/diff_diff/trop.py +++ b/diff_diff/trop.py @@ -524,6 +524,15 @@ def fit( index=all_periods, columns=all_units ) missing_mask = pd.isna(D_raw).values # True where originally missing + n_missing_treatment = int(pd.isna(D_raw).sum().sum()) + if n_missing_treatment > 0: + warnings.warn( + f"{n_missing_treatment} missing treatment indicator(s) in the " + f"(time x unit) panel matrix filled with 0 (assumed " + f"untreated). This typically occurs in unbalanced panels.", + UserWarning, + stacklevel=2, + ) D = D_raw.fillna(0).astype(int).values # Validate D is monotonic non-decreasing per unit (absorbing state) @@ -652,6 +661,13 @@ def fit( except Exception as e: # Fall back to Python implementation on error logger.debug("Rust LOOCV grid search failed, falling back to Python: %s", e) + warnings.warn( + f"Rust backend failed for LOOCV grid search; " + f"falling back to Python. Performance may be reduced. " + f"Error: {e}", + UserWarning, + stacklevel=2, + ) best_lambda = None best_score = np.inf diff --git a/diff_diff/trop_global.py b/diff_diff/trop_global.py index 7042ea322..f2a12aa1e 100644 --- a/diff_diff/trop_global.py +++ b/diff_diff/trop_global.py @@ -370,6 +370,13 @@ def _solve_global_no_lowrank( coeffs, _, _, _ = np.linalg.lstsq(X_weighted, y_weighted, rcond=None) except np.linalg.LinAlgError: # Fallback: use pseudo-inverse + warnings.warn( + "Least-squares solver failed in TROP global estimation; " + "falling back to pseudo-inverse. Results may be less " + "numerically stable.", + UserWarning, + stacklevel=2, + ) coeffs = np.dot(np.linalg.pinv(X_weighted), y_weighted) # Extract parameters @@ -560,6 +567,15 @@ def _fit_global( index=all_periods, columns=all_units ) missing_mask = pd.isna(D_raw).values + n_missing_treatment = int(pd.isna(D_raw).sum().sum()) + if n_missing_treatment > 0: + warnings.warn( + f"{n_missing_treatment} missing treatment indicator(s) in the " + f"(time x unit) panel matrix filled with 0 (assumed " + f"untreated). This typically occurs in unbalanced panels.", + UserWarning, + stacklevel=2, + ) D = D_raw.fillna(0).astype(int).values # Validate absorbing state @@ -692,6 +708,13 @@ def _fit_global( logger.debug( "Rust LOOCV grid search (global) failed, falling back to Python: %s", e ) + warnings.warn( + f"Rust backend failed for LOOCV grid search (global); " + f"falling back to Python. Performance may be reduced. " + f"Error: {e}", + UserWarning, + stacklevel=2, + ) best_lambda = None best_score = np.inf @@ -952,6 +975,13 @@ def _bootstrap_variance_global( except Exception as e: logger.debug("Rust bootstrap (global) failed, falling back to Python: %s", e) + warnings.warn( + f"Rust backend failed for bootstrap variance (global); " + f"falling back to Python. Performance may be reduced. " + f"Error: {e}", + UserWarning, + stacklevel=2, + ) # Python fallback implementation rng = np.random.default_rng(self.seed) diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index dd78fbb60..4ca780c8c 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -928,6 +928,13 @@ def _bootstrap_variance( ) except Exception as e: logger.debug("Rust bootstrap variance failed, falling back to Python: %s", e) + warnings.warn( + f"Rust backend failed for bootstrap variance; " + f"falling back to Python. Performance may be reduced. " + f"Error: {e}", + UserWarning, + stacklevel=2, + ) # Python implementation (fallback) rng = np.random.default_rng(self.seed) diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 1b3b301cf..5b0a892bc 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -305,11 +305,14 @@ def fit( if n_always_treated > 0: unit_list = ", ".join(str(u) for u in always_treated_units[:10]) suffix = f" (and {n_always_treated - 10} more)" if n_always_treated > 10 else "" + survey_note = "" + if survey_weights is not None or resolved_survey is not None: + survey_note = " Associated survey weights and design arrays " "adjusted to match." warnings.warn( f"{n_always_treated} unit(s) are treated in all observed periods " f"(first_treat <= {min_time}): [{unit_list}{suffix}]. " "These units have no untreated observations and cannot contribute " - "to the counterfactual model. Excluding from estimation.", + f"to the counterfactual model. Excluding from estimation.{survey_note}", UserWarning, stacklevel=2, ) @@ -1051,6 +1054,14 @@ def _stage2_static( # or contribute NaN treatment effects (excluded from point estimate). nan_mask = ~np.isfinite(y_tilde) if nan_mask.any(): + n_nan = int(nan_mask.sum()) + warnings.warn( + f"{n_nan} observation(s) have non-finite imputed outcomes " + f"(y_tilde) from unidentified fixed effects. These " + f"observations are excluded from ATT estimation.", + UserWarning, + stacklevel=2, + ) y_tilde[nan_mask] = 0.0 D = omega_1_mask.values.astype(float) diff --git a/tests/test_staggered.py b/tests/test_staggered.py index f5ff79380..e64573c31 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -2,6 +2,8 @@ Tests for Callaway-Sant'Anna staggered DiD estimator. """ +import warnings + import numpy as np import pandas as pd import pytest @@ -3492,9 +3494,7 @@ def test_dr_fallback_warning(self): data = generate_staggered_data_with_covariates(seed=42) - cs = CallawaySantAnna( - estimation_method="dr", pscore_fallback="unconditional" - ) + cs = CallawaySantAnna(estimation_method="dr", pscore_fallback="unconditional") with patch("diff_diff.staggered.solve_logit", side_effect=ValueError("test")): import warnings @@ -3510,9 +3510,7 @@ def test_dr_fallback_warning(self): covariates=["x1"], ) - fallback_warns = [ - x for x in w if "unconditional propensity" in str(x.message) - ] + fallback_warns = [x for x in w if "unconditional propensity" in str(x.message)] assert len(fallback_warns) > 0, "Expected fallback warning in DR path" assert results.overall_att is not None @@ -3690,9 +3688,7 @@ def test_cs_pscore_fallback_unconditional_opt_in(self): from unittest.mock import patch data = generate_staggered_data_with_covariates(seed=42) - cs = CallawaySantAnna( - estimation_method="dr", pscore_fallback="unconditional" - ) + cs = CallawaySantAnna(estimation_method="dr", pscore_fallback="unconditional") with patch("diff_diff.staggered.solve_logit", side_effect=ValueError("test")): import warnings @@ -3707,9 +3703,7 @@ def test_cs_pscore_fallback_unconditional_opt_in(self): first_treat="first_treat", covariates=["x1"], ) - fallback_warns = [ - x for x in w if "unconditional propensity" in str(x.message) - ] + fallback_warns = [x for x in w if "unconditional propensity" in str(x.message)] assert len(fallback_warns) > 0 assert results.overall_att is not None @@ -3750,15 +3744,26 @@ def test_cs_diagnose_propensity_identifies_critical(self): x2 = np.repeat(np.random.randn(n_units), n_periods) x3 = np.repeat(np.random.randn(n_units), n_periods) - data = pd.DataFrame({ - "unit": units, "time": times, "first_treat": first_treat_exp, - "outcome": outcome, "x1": x1, "x2": x2, "x3": x3, - }) + data = pd.DataFrame( + { + "unit": units, + "time": times, + "first_treat": first_treat_exp, + "outcome": outcome, + "x1": x1, + "x2": x2, + "x3": x3, + } + ) cs = CallawaySantAnna(estimation_method="ipw") df = cs.diagnose_propensity( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", covariates=["x1", "x2", "x3"], + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2", "x3"], ) # With 1 treated unit and 3 predictor variables: EPV = 1/3 ≈ 0.33 → critical assert any(df["status"] == "critical") @@ -3782,9 +3787,7 @@ def test_cs_epv_in_summary_output(self): covariates=["x1", "x2"], ) if results.epv_diagnostics: - low_epv = { - k: v for k, v in results.epv_diagnostics.items() if v.get("is_low") - } + low_epv = {k: v for k, v in results.epv_diagnostics.items() if v.get("is_low")} if low_epv: summary = results.summary() assert "EPV" in summary @@ -3834,10 +3837,16 @@ def test_cs_cached_rank_deficient_pscore_no_nan(self): # x2 is a duplicate of x1 — will cause rank deficiency x2 = x1.copy() - data = pd.DataFrame({ - "unit": units, "time": times, "first_treat": first_treat_exp, - "outcome": outcome, "x1": x1, "x2": x2, - }) + data = pd.DataFrame( + { + "unit": units, + "time": times, + "first_treat": first_treat_exp, + "outcome": outcome, + "x1": x1, + "x2": x2, + } + ) cs = CallawaySantAnna( estimation_method="ipw", @@ -3849,15 +3858,19 @@ def test_cs_cached_rank_deficient_pscore_no_nan(self): with warnings.catch_warnings(): warnings.simplefilter("ignore") results = cs.fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", covariates=["x1", "x2"], + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], ) # All ATTs should be finite (no NaN from cache poisoning) for (g, t), eff in results.group_time_effects.items(): - assert np.isfinite(eff["effect"]), ( - f"ATT({g},{t}) is {eff['effect']} — NaN cache poisoning" - ) + assert np.isfinite( + eff["effect"] + ), f"ATT({g},{t}) is {eff['effect']} — NaN cache poisoning" assert np.isfinite(results.overall_att) def test_cs_strict_mode_not_swallowed_by_unconditional_fallback(self): @@ -3882,8 +3895,12 @@ def test_cs_strict_mode_not_swallowed_by_unconditional_fallback(self): ): with pytest.raises(ValueError, match="Rank-deficient"): cs.fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", covariates=["x1"], + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1"], ) def test_cs_rc_strict_mode_not_swallowed(self): @@ -3893,13 +3910,15 @@ def test_cs_rc_strict_mode_not_swallowed(self): # RCS data: unique unit IDs per observation np.random.seed(99) n = 300 - data = pd.DataFrame({ - "unit": np.arange(n), - "time": np.random.choice([0, 1, 2, 3, 4], n), - "outcome": np.random.randn(n), - "first_treat": np.where(np.arange(n) < 100, 3, 0), - "x1": np.random.randn(n), - }) + data = pd.DataFrame( + { + "unit": np.arange(n), + "time": np.random.choice([0, 1, 2, 3, 4], n), + "outcome": np.random.randn(n), + "first_treat": np.where(np.arange(n) < 100, 3, 0), + "x1": np.random.randn(n), + } + ) cs = CallawaySantAnna( estimation_method="ipw", rank_deficient_action="error", @@ -3912,18 +3931,151 @@ def test_cs_rc_strict_mode_not_swallowed(self): ): with pytest.raises(ValueError, match="Rank-deficient"): cs.fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", covariates=["x1"], + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1"], ) def test_cs_diagnose_propensity_rejects_not_yet_treated(self): """diagnose_propensity() raises for control_group='not_yet_treated'.""" data = generate_staggered_data_with_covariates(seed=42) - cs = CallawaySantAnna( - estimation_method="ipw", control_group="not_yet_treated" - ) + cs = CallawaySantAnna(estimation_method="ipw", control_group="not_yet_treated") with pytest.raises(NotImplementedError, match="not_yet_treated"): cs.diagnose_propensity( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", covariates=["x1"], + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1"], + ) + + +class TestSilentWarningAudit: + """Tests for UserWarning emissions added by the silent warning audit.""" + + def test_item8_inf_to_zero_warning_in_fit(self): + """Item 8: Warn when first_treat=inf is recoded to 0 in fit().""" + import warnings + + data = generate_staggered_data(seed=42) + # Set some units to inf (never-treated encoding) + never_units = data.loc[data["first_treat"] == 0, "unit"].unique()[:5] + data.loc[data["unit"].isin(never_units), "first_treat"] = np.inf + + cs = CallawaySantAnna() + with pytest.warns(UserWarning, match="first_treat=inf"): + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + + def test_item8_inf_to_zero_warning_in_diagnose_propensity(self): + """Item 8: Warn when first_treat=inf is recoded in diagnose_propensity().""" + import warnings + + data = generate_staggered_data_with_covariates(seed=42) + # Set some units to inf + never_units = data.loc[data["first_treat"] == 0, "unit"].unique()[:5] + data.loc[data["unit"].isin(never_units), "first_treat"] = np.inf + + cs = CallawaySantAnna(estimation_method="ipw") + with pytest.warns(UserWarning, match="first_treat=inf"): + cs.diagnose_propensity( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1"], + ) + + def test_item8_no_warning_when_first_treat_zero(self): + """Item 8 negative: No warning when never-treated encoded as 0.""" + import warnings + + data = generate_staggered_data(seed=42) + cs = CallawaySantAnna() + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + inf_warnings = [x for x in w if "first_treat=inf" in str(x.message)] + assert len(inf_warnings) == 0 + + def test_item4_consolidated_skip_warning(self): + """Item 4: Consolidated warning when (g,t) cells are skipped.""" + import warnings + + # Create data with a cohort whose base period is outside the panel range. + # Cohort with first_treat=1 needs base_period=0 (universal) which exists, + # but cohort with first_treat=2 needs base_period=1 — which exists. + # To trigger missing period skips, create a very sparse panel. + rng = np.random.default_rng(42) + n_units = 30 + rows = [] + for u in range(n_units): + # Only include time periods 3, 5, 7 (skip 0, 1, 2, 4, 6) + for t in [3, 5, 7]: + ft = 0 if u < 10 else 5 # Cohort at t=5, or never-treated + outcome = rng.standard_normal() + (2.0 if (ft > 0 and t >= ft) else 0.0) + rows.append( + { + "unit": u, + "time": t, + "outcome": outcome, + "first_treat": ft, + } + ) + data = pd.DataFrame(rows) + + cs = CallawaySantAnna(base_period="universal") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + try: + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + except ValueError: + pass # May fail if all cells skipped + + skip_warnings = [x for x in w if "cell(s) skipped" in str(x.message)] + # Should have at least one skip warning if cells were skipped + # (The sparse panel structure should cause missing period skips) + if skip_warnings: + assert "missing base/post period" in str( + skip_warnings[0].message + ) or "zero treated or control" in str(skip_warnings[0].message) + + def test_item4_no_skip_warning_normal_data(self): + """Item 4 negative: No skip warning on well-formed balanced data.""" + import warnings + + data = generate_staggered_data(seed=42) + cs = CallawaySantAnna() + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", ) + skip_warnings = [x for x in w if "cell(s) skipped" in str(x.message)] + assert len(skip_warnings) == 0, f"Unexpected skip warning: {skip_warnings}" diff --git a/tests/test_survey.py b/tests/test_survey.py index 534379e07..4534ce9ac 100644 --- a/tests/test_survey.py +++ b/tests/test_survey.py @@ -3294,3 +3294,41 @@ def test_twfe_non_survey_default_clustering_unaffected(self, twfe_panel_data): assert result is not None assert np.isfinite(result.se) assert result.se > 0 + + +class TestSilentWarningAudit: + """Tests for UserWarning emissions added by the silent warning audit.""" + + def test_item7_weight_normalization_warning(self): + """Item 7: Warn when pweight/aweight are normalized.""" + n = 100 + raw_weights = np.random.default_rng(42).uniform(1.0, 10.0, n) + assert not np.isclose(np.sum(raw_weights), n) + + df = pd.DataFrame({"x": np.arange(n), "w": raw_weights}) + sd = SurveyDesign(weights="w", weight_type="pweight") + with pytest.warns(UserWarning, match="pweight weights normalized"): + sd.resolve(df) + + def test_item7_no_warning_when_already_normalized(self): + """Item 7 negative: No warning when weights already sum to n.""" + n = 100 + raw_weights = np.ones(n) # sum = n, mean = 1 + assert np.isclose(np.sum(raw_weights), n) + + df = pd.DataFrame({"x": np.arange(n), "w": raw_weights}) + sd = SurveyDesign(weights="w", weight_type="pweight") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + sd.resolve(df) + norm_warnings = [x for x in w if "normalized" in str(x.message)] + assert len(norm_warnings) == 0 + + def test_item7_aweight_normalization_warning(self): + """Item 7: aweight also triggers normalization warning.""" + n = 50 + raw_weights = np.random.default_rng(42).uniform(1.0, 5.0, n) + df = pd.DataFrame({"x": np.arange(n), "w": raw_weights}) + sd = SurveyDesign(weights="w", weight_type="aweight") + with pytest.warns(UserWarning, match="aweight weights normalized"): + sd.resolve(df) diff --git a/tests/test_trop.py b/tests/test_trop.py index fd3762e1a..d39cfd6be 100644 --- a/tests/test_trop.py +++ b/tests/test_trop.py @@ -82,12 +82,14 @@ def simple_panel_data(): if treatment_indicator: y += true_att y += rng.normal(0, 0.5) - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) return pd.DataFrame(data) @@ -109,7 +111,7 @@ def test_basic_fit(self, simple_panel_data): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( simple_panel_data, @@ -133,7 +135,7 @@ def test_fit_with_factors(self, factor_dgp_data, ci_params): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1, 1.0], n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( factor_dgp_data, @@ -157,7 +159,7 @@ def test_treatment_effect_recovery(self, factor_dgp_data, ci_params): lambda_unit_grid=[0.0, 0.5, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( factor_dgp_data, @@ -180,7 +182,7 @@ def test_tuning_parameter_selection(self, simple_panel_data, ci_params): lambda_unit_grid=[0.0, 0.5, 1.0], lambda_nn_grid=[0.0, 0.1, 1.0], n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( simple_panel_data, @@ -203,7 +205,7 @@ def test_bootstrap_variance(self, simple_panel_data, ci_params): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( simple_panel_data, @@ -226,7 +228,7 @@ def test_confidence_interval(self, simple_panel_data, ci_params): lambda_nn_grid=[0.0, 0.1], alpha=0.05, n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( simple_panel_data, @@ -253,10 +255,7 @@ def test_get_set_params(self): def test_missing_columns(self, simple_panel_data): """Test error when column is missing.""" trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) with pytest.raises(ValueError, match="Missing columns"): trop_est.fit( @@ -269,18 +268,17 @@ def test_missing_columns(self, simple_panel_data): def test_no_treated_observations(self): """Test error when no treated observations.""" - data = pd.DataFrame({ - "unit": [0, 0, 1, 1], - "period": [0, 1, 0, 1], - "outcome": [1, 2, 3, 4], - "treated": [0, 0, 0, 0], - }) + data = pd.DataFrame( + { + "unit": [0, 0, 1, 1], + "period": [0, 1, 0, 1], + "outcome": [1, 2, 3, 4], + "treated": [0, 0, 0, 0], + } + ) trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) with pytest.raises(ValueError, match="No treated observations"): trop_est.fit( @@ -293,18 +291,17 @@ def test_no_treated_observations(self): def test_no_control_units(self): """Test error when no control units.""" - data = pd.DataFrame({ - "unit": [0, 0, 1, 1], - "period": [0, 1, 0, 1], - "outcome": [1, 2, 3, 4], - "treated": [0, 1, 0, 1], # Both units become treated - }) + data = pd.DataFrame( + { + "unit": [0, 0, 1, 1], + "period": [0, 1, 0, 1], + "outcome": [1, 2, 3, 4], + "treated": [0, 1, 0, 1], # Both units become treated + } + ) trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) with pytest.raises(ValueError, match="No control units"): trop_est.fit( @@ -334,10 +331,14 @@ def fitted_results(self): if is_treated and post: y += true_att y += rng.normal(0, 0.5) - data.append({ - "unit": i, "period": t, "outcome": y, - "treated": 1 if (is_treated and post) else 0, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": 1 if (is_treated and post) else 0, + } + ) panel = pd.DataFrame(data) trop_est = TROP( @@ -348,8 +349,11 @@ def fitted_results(self): seed=42, ) return trop_est.fit( - panel, outcome="outcome", treatment="treated", - unit="unit", time="period", + panel, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", ) def test_summary(self, fitted_results): @@ -410,7 +414,7 @@ def test_significance_properties(self, simple_panel_data, ci_params): lambda_nn_grid=[0.0, 0.1], alpha=0.05, n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( simple_panel_data, @@ -488,7 +492,7 @@ def test_trop_handles_factor_dgp(self, ci_params): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1, 1.0], n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( data, @@ -579,12 +583,14 @@ def test_limiting_case_uniform_weights(self): if treatment_indicator: y += true_att y += rng.normal(0, 0.3) - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -594,7 +600,7 @@ def test_limiting_case_uniform_weights(self): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -605,8 +611,9 @@ def test_limiting_case_uniform_weights(self): ) # Should recover treatment effect within reasonable tolerance - assert abs(results.att - true_att) < 1.0, \ - f"ATT={results.att:.3f} should be close to true={true_att}" + assert ( + abs(results.att - true_att) < 1.0 + ), f"ATT={results.att:.3f} should be close to true={true_att}" # Check that uniform weights were selected assert results.lambda_time == 0.0 assert results.lambda_unit == 0.0 @@ -645,12 +652,14 @@ def test_unit_weights_reduce_bias(self): if treatment_indicator: y += true_att y += rng.normal(0, 0.3) - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -660,7 +669,7 @@ def test_unit_weights_reduce_bias(self): lambda_unit_grid=[0.0, 1.0, 2.0], lambda_nn_grid=[0.0], n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -671,8 +680,9 @@ def test_unit_weights_reduce_bias(self): ) # Should recover treatment effect reasonably well - assert abs(results.att - true_att) < 1.5, \ - f"ATT={results.att:.3f} should be close to true={true_att}" + assert ( + abs(results.att - true_att) < 1.5 + ), f"ATT={results.att:.3f} should be close to true={true_att}" def test_time_weights_reduce_bias(self): """ @@ -696,18 +706,20 @@ def test_time_weights_reduce_bias(self): for t in range(n_pre + n_post): post = t >= n_pre # Time trend that accelerates near treatment - time_fe = 0.1 * t + 0.05 * (t ** 2 / n_pre) + time_fe = 0.1 * t + 0.05 * (t**2 / n_pre) y = 10.0 + unit_fe + time_fe treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_att y += rng.normal(0, 0.3) - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -717,7 +729,7 @@ def test_time_weights_reduce_bias(self): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -760,7 +772,7 @@ def test_factor_model_reduces_bias(self, ci_params): lambda_unit_grid=[0.0, 0.5], lambda_nn_grid=nn_grid, n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( data, @@ -772,8 +784,9 @@ def test_factor_model_reduces_bias(self, ci_params): true_att = 2.0 # With factor adjustment, should recover treatment effect - assert abs(results.att - true_att) < 2.0, \ - f"ATT={results.att:.3f} should be within 2.0 of true={true_att}" + assert ( + abs(results.att - true_att) < 2.0 + ), f"ATT={results.att:.3f} should be within 2.0 of true={true_att}" # Factor matrix should capture some structure assert results.effective_rank > 0, "Factor matrix should have positive rank" @@ -825,12 +838,14 @@ def test_paper_dgp_recovery(self, ci_params): y += true_tau y += rng.normal(0, 0.5) # Idiosyncratic noise - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -841,7 +856,7 @@ def test_paper_dgp_recovery(self, ci_params): lambda_unit_grid=[0.0, 0.5, 1.0], lambda_nn_grid=[0.0, 0.1, 1.0], n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -853,8 +868,9 @@ def test_paper_dgp_recovery(self, ci_params): # Under null hypothesis, ATT should be close to zero # Allow for estimation error (this is a finite sample) - assert abs(results.att) < 2.0, \ - f"ATT={results.att:.3f} should be close to true={true_tau} under null" + assert ( + abs(results.att) < 2.0 + ), f"ATT={results.att:.3f} should be close to true={true_tau} under null" # Check that factor model was used assert results.effective_rank >= 0 @@ -880,7 +896,7 @@ def test_precomputed_structures_consistency(self, simple_panel_data): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # Fit to populate precomputed structures @@ -947,8 +963,7 @@ def test_vectorized_alternating_minimization(self): # Run the estimation alpha_est, beta_est, L_est = trop_est._estimate_model( - Y, control_mask, W, lambda_nn=0.0, - n_units=n_units, n_periods=n_periods + Y, control_mask, W, lambda_nn=0.0, n_units=n_units, n_periods=n_periods ) # Check that we recovered the fixed effects structure @@ -973,7 +988,7 @@ def test_vectorized_weights_computation(self, simple_panel_data): lambda_unit_grid=[0.5], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # Fit to populate precomputed structures @@ -1013,8 +1028,7 @@ def test_vectorized_weights_computation(self, simple_panel_data): lambda_unit = 0.5 weights = trop_est._compute_observation_weights( - Y, D, i, t, lambda_time, lambda_unit, control_unit_idx, - n_units, n_periods + Y, D, i, t, lambda_time, lambda_unit, control_unit_idx, n_units, n_periods ) # Verify shape @@ -1025,8 +1039,9 @@ def test_vectorized_weights_computation(self, simple_panel_data): for s in range(n_periods): expected = np.exp(-lambda_time * abs(t - s)) # Time weight should be proportional to expected - assert np.isclose(time_weights[s], expected, rtol=1e-5) or \ - np.isclose(time_weights[s] / weights[t, i], expected / weights[t, i], rtol=1e-5) + assert np.isclose(time_weights[s], expected, rtol=1e-5) or np.isclose( + time_weights[s] / weights[t, i], expected / weights[t, i], rtol=1e-5 + ) def test_pivot_vs_iterrows_equivalence(self): """ @@ -1042,12 +1057,14 @@ def test_pivot_vs_iterrows_equivalence(self): data = [] for i in range(n_units): for t in range(n_periods): - data.append({ - "unit": i, - "period": t, - "outcome": rng.normal(0, 1), - "treated": 1 if (i < 3 and t >= 3) else 0, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": rng.normal(0, 1), + "treated": 1 if (i < 3 and t >= 3) else 0, + } + ) df = pd.DataFrame(data) all_units = sorted(df["unit"].unique()) @@ -1148,12 +1165,14 @@ def test_d_matrix_absorbing_state_validation_valid(self): y = 10.0 + rng.normal(0, 0.5) if is_treated: y += 2.0 - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": 1 if is_treated else 0, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": 1 if is_treated else 0, + } + ) df = pd.DataFrame(data) @@ -1163,7 +1182,7 @@ def test_d_matrix_absorbing_state_validation_valid(self): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -1190,20 +1209,19 @@ def test_d_matrix_absorbing_state_validation_invalid(self): treated = 1 else: treated = 0 - data.append({ - "unit": i, - "period": t, - "outcome": float(i + t), - "treated": treated, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": float(i + t), + "treated": treated, + } + ) df = pd.DataFrame(data) trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) with pytest.raises(ValueError, match="not an absorbing state"): @@ -1227,20 +1245,19 @@ def test_d_matrix_validation_error_message_helpful(self): else: # Other units: proper absorbing state treated = 1 if (i < 3 and t >= 3) else 0 - data.append({ - "unit": i, - "period": t, - "outcome": float(i + t), - "treated": treated, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": float(i + t), + "treated": treated, + } + ) df = pd.DataFrame(data) trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) with pytest.raises(ValueError) as exc_info: @@ -1271,7 +1288,7 @@ def test_cycling_search_converges(self, simple_panel_data): lambda_unit_grid=[0.0, 0.5, 1.0], lambda_nn_grid=[0.0, 0.1, 1.0], n_bootstrap=5, - seed=42 + seed=42, ) results = trop_est.fit( @@ -1330,9 +1347,9 @@ def test_cycling_search_single_value_grids(self, simple_panel_data): trop_est = TROP( lambda_time_grid=[0.5], # Single value lambda_unit_grid=[0.5], # Single value - lambda_nn_grid=[0.1], # Single value + lambda_nn_grid=[0.1], # Single value n_bootstrap=5, - seed=42 + seed=42, ) results = trop_est.fit( @@ -1374,8 +1391,8 @@ def test_issue_a_control_includes_pretreatment_obs(self): rng = np.random.default_rng(42) n_units = 20 n_early_treat = 5 # Units treated at period 3 - n_late_treat = 5 # Units treated at period 5 - n_control = 10 # Never-treated units + n_late_treat = 5 # Units treated at period 5 + n_control = 10 # Never-treated units n_periods = 8 true_att = 2.0 @@ -1399,12 +1416,14 @@ def test_issue_a_control_includes_pretreatment_obs(self): if treatment_indicator: y += true_att y += rng.normal(0, 0.3) - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1415,7 +1434,7 @@ def test_issue_a_control_includes_pretreatment_obs(self): lambda_unit_grid=[1.0], # Use unit weights so distance matters lambda_nn_grid=[0.0], n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -1453,12 +1472,14 @@ def test_issue_b_distance_excludes_target_period(self): y = 5.0 + rng.normal(0, 0.1) treatment_indicator = 1 if (is_treated and t >= 3) else 0 - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1467,7 +1488,7 @@ def test_issue_b_distance_excludes_target_period(self): lambda_unit_grid=[1.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # With Issue B fix (target period excluded), this should complete @@ -1515,12 +1536,14 @@ def test_issue_c_weighted_nuclear_norm(self): if treatment_indicator: y += true_att y += rng.normal(0, 0.3) - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1530,7 +1553,7 @@ def test_issue_c_weighted_nuclear_norm(self): lambda_unit_grid=[0.0], lambda_nn_grid=[0.1, 1.0], # Use regularization n_bootstrap=10, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -1571,12 +1594,14 @@ def test_issue_d_stratified_bootstrap(self, ci_params): treatment_indicator = 1 if (is_treated and post) else 0 if treatment_indicator: y += true_att - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1587,7 +1612,7 @@ def test_issue_d_stratified_bootstrap(self, ci_params): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=n_boot, - seed=42 + seed=42, ) results = trop_est.fit( df, @@ -1639,8 +1664,9 @@ def test_weighted_nuclear_norm_solver_convergence(self): _, s, _ = np.linalg.svd(L, full_matrices=False) _, s_orig, _ = np.linalg.svd(Y, full_matrices=False) # Regularized singular values should be smaller than original - assert np.sum(s) < np.sum(s_orig), \ - "Nuclear norm regularization should reduce total singular value mass" + assert np.sum(s) < np.sum( + s_orig + ), "Nuclear norm regularization should reduce total singular value mass" class TestAPIChangesV2_1_8: @@ -1660,7 +1686,7 @@ def test_fit_no_post_periods_parameter(self, simple_panel_data): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # This should work - no post_periods parameter @@ -1725,7 +1751,7 @@ def test_results_has_period_counts_not_lists(self, simple_panel_data): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) results = trop_est.fit( simple_panel_data, @@ -1752,18 +1778,17 @@ def test_results_has_period_counts_not_lists(self, simple_panel_data): def test_validation_still_checks_pre_periods(self): """Test that validation still requires at least 2 pre-treatment periods.""" # Create data with only 1 pre-treatment period - data = pd.DataFrame({ - "unit": [0, 0, 1, 1], - "period": [0, 1, 0, 1], - "outcome": [1.0, 2.0, 1.5, 2.5], - "treated": [0, 1, 0, 0], # Treatment at period 1 - }) + data = pd.DataFrame( + { + "unit": [0, 0, 1, 1], + "period": [0, 1, 0, 1], + "outcome": [1.0, 2.0, 1.5, 2.5], + "treated": [0, 1, 0, 0], # Treatment at period 1 + } + ) trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) with pytest.raises(ValueError, match="at least 2 pre-treatment periods"): @@ -1792,12 +1817,14 @@ def test_loocv_warning_on_many_failures(self): # Add some extreme values that might cause numerical issues y = rng.normal(0, 1) if not (is_treated and post) else 1e10 treatment_indicator = 1 if (is_treated and post) else 0 - data.append({ - "unit": i, - "period": t, - "outcome": y, - "treated": treatment_indicator, - }) + data.append( + { + "unit": i, + "period": t, + "outcome": y, + "treated": treatment_indicator, + } + ) df = pd.DataFrame(data) @@ -1806,7 +1833,7 @@ def test_loocv_warning_on_many_failures(self): lambda_unit_grid=[100.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # Capture warnings and verify the warning code path @@ -1828,9 +1855,7 @@ def test_loocv_warning_on_many_failures(self): # Check for LOOCV-related warnings loocv_warnings = [ - x for x in w - if issubclass(x.category, UserWarning) - and "LOOCV" in str(x.message) + x for x in w if issubclass(x.category, UserWarning) and "LOOCV" in str(x.message) ] # If fit succeeded, check that we can capture warnings properly @@ -1860,7 +1885,7 @@ def test_loocv_warning_deterministic_with_mock(self, simple_panel_data): lambda_unit_grid=[1.0], lambda_nn_grid=[0.1], n_bootstrap=5, - seed=42 + seed=42, ) # Mock _estimate_model to fail on the first LOOCV call @@ -1881,10 +1906,13 @@ def mock_estimate_with_failure(*args, **kwargs): # Disable Rust backend for this test by patching the module-level variables import sys - trop_module = sys.modules['diff_diff.trop'] - with patch.object(trop_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_module, '_rust_loocv_grid_search', None), \ - patch.object(trop_est, '_estimate_model', mock_estimate_with_failure): + + trop_module = sys.modules["diff_diff.trop"] + with ( + patch.object(trop_module, "HAS_RUST_BACKEND", False), + patch.object(trop_module, "_rust_loocv_grid_search", None), + patch.object(trop_est, "_estimate_model", mock_estimate_with_failure), + ): try: trop_est.fit( simple_panel_data, @@ -1899,9 +1927,7 @@ def mock_estimate_with_failure(*args, **kwargs): # Check that LOOCV warning was raised on first failure loocv_warnings = [ - x for x in w - if issubclass(x.category, UserWarning) - and "LOOCV" in str(x.message) + x for x in w if issubclass(x.category, UserWarning) and "LOOCV" in str(x.message) ] # With any failure, we should get a warning about returning infinity @@ -1938,7 +1964,7 @@ def test_infinite_score_triggers_fallback(self, simple_panel_data): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=5, - seed=42 + seed=42, ) # Mock LOOCV to always return infinity @@ -1949,10 +1975,12 @@ def always_infinity(*args, **kwargs): warnings.simplefilter("always") # Disable Rust backend and mock LOOCV score to always return infinity - trop_module = sys.modules['diff_diff.trop'] - with patch.object(trop_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_module, '_rust_loocv_grid_search', None), \ - patch.object(trop_est, '_loocv_score_obs_specific', always_infinity): + trop_module = sys.modules["diff_diff.trop"] + with ( + patch.object(trop_module, "HAS_RUST_BACKEND", False), + patch.object(trop_module, "_rust_loocv_grid_search", None), + patch.object(trop_est, "_loocv_score_obs_specific", always_infinity), + ): results = trop_est.fit( simple_panel_data, outcome="outcome", @@ -1963,21 +1991,24 @@ def always_infinity(*args, **kwargs): # Verify warning emitted about fallback to defaults fallback_warnings = [ - x for x in w - if issubclass(x.category, UserWarning) - and "defaults" in str(x.message).lower() + x + for x in w + if issubclass(x.category, UserWarning) and "defaults" in str(x.message).lower() ] - assert len(fallback_warnings) > 0, ( - f"Expected fallback warning, got: {[str(x.message) for x in w]}" - ) + assert ( + len(fallback_warnings) > 0 + ), f"Expected fallback warning, got: {[str(x.message) for x in w]}" # Verify defaults used (per REGISTRY.md: 1.0, 1.0, 0.1) - assert results.lambda_time == 1.0, \ - f"Expected default lambda_time=1.0, got {results.lambda_time}" - assert results.lambda_unit == 1.0, \ - f"Expected default lambda_unit=1.0, got {results.lambda_unit}" - assert results.lambda_nn == 0.1, \ - f"Expected default lambda_nn=0.1, got {results.lambda_nn}" + assert ( + results.lambda_time == 1.0 + ), f"Expected default lambda_time=1.0, got {results.lambda_time}" + assert ( + results.lambda_unit == 1.0 + ), f"Expected default lambda_unit=1.0, got {results.lambda_unit}" + assert ( + results.lambda_nn == 0.1 + ), f"Expected default lambda_nn=0.1, got {results.lambda_nn}" # Verify estimation still completed assert np.isfinite(results.att), "ATT should be finite even with default params" @@ -1999,7 +2030,7 @@ def test_rust_infinite_score_triggers_fallback(self, simple_panel_data): lambda_unit_grid=[0.0, 1.0], lambda_nn_grid=[0.0, 0.1], n_bootstrap=5, - seed=42 + seed=42, ) # Mock Rust function to return infinite score @@ -2013,10 +2044,12 @@ def always_infinity(*args, **kwargs): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - trop_module = sys.modules['diff_diff.trop'] - with patch.object(trop_module, 'HAS_RUST_BACKEND', True), \ - patch.object(trop_module, '_rust_loocv_grid_search', mock_rust_loocv), \ - patch.object(trop_est, '_loocv_score_obs_specific', always_infinity): + trop_module = sys.modules["diff_diff.trop"] + with ( + patch.object(trop_module, "HAS_RUST_BACKEND", True), + patch.object(trop_module, "_rust_loocv_grid_search", mock_rust_loocv), + patch.object(trop_est, "_loocv_score_obs_specific", always_infinity), + ): results = trop_est.fit( simple_panel_data, outcome="outcome", @@ -2027,21 +2060,24 @@ def always_infinity(*args, **kwargs): # Verify warning emitted about fallback to defaults fallback_warnings = [ - x for x in w - if issubclass(x.category, UserWarning) - and "defaults" in str(x.message).lower() + x + for x in w + if issubclass(x.category, UserWarning) and "defaults" in str(x.message).lower() ] - assert len(fallback_warnings) > 0, ( - f"Expected fallback warning with Rust backend, got: {[str(x.message) for x in w]}" - ) + assert ( + len(fallback_warnings) > 0 + ), f"Expected fallback warning with Rust backend, got: {[str(x.message) for x in w]}" # Verify defaults used (NOT the Rust-returned values) - assert results.lambda_time == 1.0, \ - f"Expected default lambda_time=1.0, got {results.lambda_time}" - assert results.lambda_unit == 1.0, \ - f"Expected default lambda_unit=1.0, got {results.lambda_unit}" - assert results.lambda_nn == 0.1, \ - f"Expected default lambda_nn=0.1, got {results.lambda_nn}" + assert ( + results.lambda_time == 1.0 + ), f"Expected default lambda_time=1.0, got {results.lambda_time}" + assert ( + results.lambda_unit == 1.0 + ), f"Expected default lambda_unit=1.0, got {results.lambda_unit}" + assert ( + results.lambda_nn == 0.1 + ), f"Expected default lambda_nn=0.1, got {results.lambda_nn}" def test_uniform_weights_and_disabled_factor_handled_consistently(self, simple_panel_data): """ @@ -2054,11 +2090,11 @@ def test_uniform_weights_and_disabled_factor_handled_consistently(self, simple_p - λ_nn=∞ → factor model disabled (L=0), converted to 1e10 internally """ 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 + 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 + seed=42, ) results = trop_est.fit( @@ -2070,23 +2106,21 @@ def test_uniform_weights_and_disabled_factor_handled_consistently(self, simple_p ) # ATT should be finite - assert np.isfinite(results.att), ( - f"ATT should be finite with uniform weights and no factor model, got {results.att}" - ) + assert np.isfinite( + 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 - assert np.isfinite(results.se) or results.se >= 0, ( - f"SE should be finite, got {results.se}" - ) + assert np.isfinite(results.se) or results.se >= 0, f"SE should be finite, got {results.se}" # 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 ( + results.lambda_time == 0.0 + ), f"lambda_time should be 0.0 (uniform weights), got {results.lambda_time}" # 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}" - ) + 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): """ @@ -2125,11 +2159,11 @@ def test_variance_estimation_uses_converted_params(self, simple_panel_data): from unittest.mock import patch trop_est = TROP( - lambda_time_grid=[0.0], # Uniform time weights (paper convention) + 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 + lambda_nn_grid=[np.inf], # Will be converted to 1e10 internally n_bootstrap=5, - seed=42 + seed=42, ) # Track what parameters are passed to _fit_with_fixed_lambda @@ -2139,9 +2173,11 @@ def test_variance_estimation_uses_converted_params(self, simple_panel_data): def tracking_fit(self, data, outcome, treatment, unit, time, fixed_lambda, **kwargs): captured_lambda.append(fixed_lambda) - return original_fit_with_fixed(self, data, outcome, treatment, unit, time, fixed_lambda, **kwargs) + return original_fit_with_fixed( + self, data, outcome, treatment, unit, time, fixed_lambda, **kwargs + ) - with patch.object(TROP, '_fit_with_fixed_lambda', tracking_fit): + with patch.object(TROP, "_fit_with_fixed_lambda", tracking_fit): results = trop_est.fit( simple_panel_data, outcome="outcome", @@ -2153,7 +2189,9 @@ def tracking_fit(self, data, outcome, treatment, unit, time, fixed_lambda, **kwa # 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" + 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" @@ -2162,12 +2200,10 @@ def tracking_fit(self, data, outcome, treatment, unit, time, fixed_lambda, **kwa # Check that bootstrap iterations used converted (non-infinite) λ_nn values for captured in captured_lambda: lambda_time, lambda_unit, lambda_nn = captured - 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}" - ) + 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}" def test_empty_control_obs_returns_infinity(self, simple_panel_data): """ @@ -2179,26 +2215,23 @@ def test_empty_control_obs_returns_infinity(self, simple_panel_data): import warnings trop_est = TROP( - lambda_time_grid=[1.0], - lambda_unit_grid=[1.0], - lambda_nn_grid=[1.0], - seed=42 + lambda_time_grid=[1.0], lambda_unit_grid=[1.0], lambda_nn_grid=[1.0], seed=42 ) # Setup matrices from data data = simple_panel_data - all_units = sorted(data['unit'].unique()) - all_periods = sorted(data['period'].unique()) + all_units = sorted(data["unit"].unique()) + all_periods = sorted(data["period"].unique()) n_units = len(all_units) n_periods = len(all_periods) Y = ( - data.pivot(index='period', columns='unit', values='outcome') + data.pivot(index="period", columns="unit", values="outcome") .reindex(index=all_periods, columns=all_units) .values ) D = ( - data.pivot(index='period', columns='unit', values='treated') + data.pivot(index="period", columns="unit", values="treated") .reindex(index=all_periods, columns=all_units) .fillna(0) .astype(int) @@ -2211,16 +2244,15 @@ def test_empty_control_obs_returns_infinity(self, simple_panel_data): # Force empty control_obs by setting precomputed with empty list trop_est._precomputed = { "control_obs": [], # Empty! - "time_dist_matrix": np.abs(np.subtract.outer( - np.arange(n_periods), np.arange(n_periods) - )), + "time_dist_matrix": np.abs( + np.subtract.outer(np.arange(n_periods), np.arange(n_periods)) + ), } with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") score = trop_est._loocv_score_obs_specific( - Y, D, control_mask, control_unit_idx, - 1.0, 1.0, 1.0, n_units, n_periods + Y, D, control_mask, control_unit_idx, 1.0, 1.0, 1.0, n_units, n_periods ) # Should return infinity, not 0.0 @@ -2228,9 +2260,9 @@ def test_empty_control_obs_returns_infinity(self, simple_panel_data): # Should emit warning warning_msgs = [str(warning.message) for warning in w] - assert any("No valid control observations" in msg for msg in warning_msgs), ( - f"Should warn about empty control obs. Warnings: {warning_msgs}" - ) + assert any( + "No valid control observations" in msg for msg in warning_msgs + ), f"Should warn about empty control obs. Warnings: {warning_msgs}" def test_original_grid_values_stored_in_results(self, simple_panel_data): """ @@ -2240,11 +2272,11 @@ def test_original_grid_values_stored_in_results(self, simple_panel_data): λ_nn stores the original inf value when factor model is disabled. """ trop_est = TROP( - lambda_time_grid=[0.0], # Uniform time weights + lambda_time_grid=[0.0], # Uniform time weights lambda_unit_grid=[0.5], - lambda_nn_grid=[np.inf], # Factor model disabled (original: inf) + lambda_nn_grid=[np.inf], # Factor model disabled (original: inf) n_bootstrap=5, - seed=42 + seed=42, ) results = trop_est.fit( @@ -2256,16 +2288,16 @@ def test_original_grid_values_stored_in_results(self, simple_panel_data): ) # 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}" - ) + 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}" - ) + assert np.isinf( + results.lambda_nn + ), f"results.lambda_nn should be inf (original), got {results.lambda_nn}" # But ATT should still be finite (computed with converted values) assert np.isfinite(results.att), "ATT should be finite" @@ -2293,33 +2325,39 @@ def test_unbalanced_panel_d_matrix_validation(self): # Unit 0: control, complete panel for t in range(6): - data.append({ - "unit": 0, - "period": t, - "outcome": 10.0 + t, - "treated": 0, - }) + data.append( + { + "unit": 0, + "period": t, + "outcome": 10.0 + t, + "treated": 0, + } + ) # Unit 1: treated from t=3, missing t=5 (unbalanced) for t in range(6): if t == 5: continue # Skip period 5 - creates unbalanced panel treated = 1 if t >= 3 else 0 - data.append({ - "unit": 1, - "period": t, - "outcome": 10.0 + t + (2.0 if treated else 0), - "treated": treated, - }) + data.append( + { + "unit": 1, + "period": t, + "outcome": 10.0 + t + (2.0 if treated else 0), + "treated": treated, + } + ) # Unit 2: control, complete panel for t in range(6): - data.append({ - "unit": 2, - "period": t, - "outcome": 10.0 + t, - "treated": 0, - }) + data.append( + { + "unit": 2, + "period": t, + "outcome": 10.0 + t, + "treated": 0, + } + ) df = pd.DataFrame(data) @@ -2329,7 +2367,7 @@ def test_unbalanced_panel_d_matrix_validation(self): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # Should not raise ValueError - missing data is not a violation @@ -2361,12 +2399,14 @@ def test_unbalanced_panel_real_violation_still_caught(self): # Unit 0: control, complete for t in range(5): - data.append({ - "unit": 0, - "period": t, - "outcome": 10.0 + t, - "treated": 0, - }) + data.append( + { + "unit": 0, + "period": t, + "outcome": 10.0 + t, + "treated": 0, + } + ) # Unit 1: REAL violation - D goes 0→1→0 on observed periods (t=2: D=1, t=3: D=0) # This is a real violation, not a missing data artifact @@ -2375,29 +2415,30 @@ def test_unbalanced_panel_real_violation_still_caught(self): treated = 1 else: treated = 0 - data.append({ - "unit": 1, - "period": t, - "outcome": 10.0 + t, - "treated": treated, - }) + data.append( + { + "unit": 1, + "period": t, + "outcome": 10.0 + t, + "treated": treated, + } + ) # Unit 2: control for t in range(5): - data.append({ - "unit": 2, - "period": t, - "outcome": 10.0 + t, - "treated": 0, - }) + data.append( + { + "unit": 2, + "period": t, + "outcome": 10.0 + t, + "treated": 0, + } + ) df = pd.DataFrame(data) trop_est = TROP( - lambda_time_grid=[0.0], - lambda_unit_grid=[0.0], - lambda_nn_grid=[0.0], - n_bootstrap=5 + lambda_time_grid=[0.0], lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5 ) # This SHOULD raise an error - real violation @@ -2416,35 +2457,41 @@ def test_unbalanced_panel_multiple_missing_periods(self): # Unit 0: control, complete for t in range(8): - data.append({ - "unit": 0, - "period": t, - "outcome": 10.0 + t, - "treated": 0, - }) + data.append( + { + "unit": 0, + "period": t, + "outcome": 10.0 + t, + "treated": 0, + } + ) # Unit 1: treated from t=4, missing t=2 and t=6 for t in range(8): if t in [2, 6]: continue # Skip these periods treated = 1 if t >= 4 else 0 - data.append({ - "unit": 1, - "period": t, - "outcome": 10.0 + t + (2.0 if treated else 0), - "treated": treated, - }) + data.append( + { + "unit": 1, + "period": t, + "outcome": 10.0 + t + (2.0 if treated else 0), + "treated": treated, + } + ) # Unit 2: control, missing t=0 for t in range(8): if t == 0: continue - data.append({ - "unit": 2, - "period": t, - "outcome": 10.0 + t, - "treated": 0, - }) + data.append( + { + "unit": 2, + "period": t, + "outcome": 10.0 + t, + "treated": 0, + } + ) df = pd.DataFrame(data) @@ -2453,7 +2500,7 @@ def test_unbalanced_panel_multiple_missing_periods(self): lambda_unit_grid=[0.0], lambda_nn_grid=[0.0], n_bootstrap=5, - seed=42 + seed=42, ) # Should not raise error @@ -2475,11 +2522,11 @@ def test_mixed_grid_values_with_final_score_computation(self, simple_panel_data) use finite values only (0.0 = uniform weights per Eq. 3). """ trop_est = TROP( - 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 + 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 + seed=42, ) # This should complete without error @@ -2501,9 +2548,9 @@ def test_mixed_grid_values_with_final_score_computation(self, simple_panel_data) # but ATT should still be finite (falls back to defaults) pass else: - assert np.isfinite(results.loocv_score), ( - "LOOCV score should be finite when computed with converted inf values" - ) + assert np.isfinite( + results.loocv_score + ), "LOOCV score should be finite when computed with converted inf values" def test_violation_across_missing_gap_caught(self): """Test that 1→0 violations spanning missing periods are caught. @@ -2567,12 +2614,14 @@ def test_n_post_periods_counts_observed_treatment(self): if unit in [1, 2] and period == 5: continue # Skip - creates unbalanced panel treated = 1 if (unit in [1, 2] and period >= 3) else 0 - data.append({ - "unit": unit, - "period": period, - "outcome": 10.0 + period, - "treated": treated, - }) + data.append( + { + "unit": unit, + "period": period, + "outcome": 10.0 + period, + "treated": treated, + } + ) df = pd.DataFrame(data) trop_est = TROP( @@ -2591,9 +2640,9 @@ def test_n_post_periods_counts_observed_treatment(self): ) # Periods with D=1 observations: 3, 4 (not 5 - missing for treated units) - assert results.n_post_periods == 2, ( - f"Expected 2 post-periods with D=1, got {results.n_post_periods}" - ) + assert ( + results.n_post_periods == 2 + ), f"Expected 2 post-periods with D=1, got {results.n_post_periods}" class TestTROPNuclearNormSolver: @@ -2651,9 +2700,9 @@ def test_lowrank_objective_decreases(self): # Objective should be non-increasing (within numerical tolerance) for k in range(1, len(objectives)): - assert objectives[k] <= objectives[k - 1] + 1e-10, ( - f"Objective increased at step {k}: {objectives[k]} > {objectives[k-1]}" - ) + assert ( + objectives[k] <= objectives[k - 1] + 1e-10 + ), f"Objective increased at step {k}: {objectives[k]} > {objectives[k-1]}" def test_local_nonuniform_weights_objective(self): """Verify objective decreases with non-uniform weights (W_max < 1).""" @@ -2686,16 +2735,16 @@ def test_local_nonuniform_weights_objective(self): _, s_final, _ = np.linalg.svd(L_final, full_matrices=False) obj_final = f_final + lambda_nn * np.sum(s_final) - assert obj_final <= obj_init + 1e-10, ( - f"Objective did not decrease: {obj_final} > {obj_init}" - ) + assert ( + obj_final <= obj_init + 1e-10 + ), f"Objective did not decrease: {obj_final} > {obj_init}" # Soft-thresholding should reduce nuclear norm vs residual nuclear_norm_R = np.sum(np.linalg.svd(R, compute_uv=False)) nuclear_norm_L = np.sum(s_final) - assert nuclear_norm_L < nuclear_norm_R, ( - f"Nuclear norm not reduced: {nuclear_norm_L} >= {nuclear_norm_R}" - ) + assert ( + nuclear_norm_L < nuclear_norm_R + ), f"Nuclear norm not reduced: {nuclear_norm_L} >= {nuclear_norm_R}" def test_zero_weights_no_division_error(self): """Verify solver handles all-zero weights without ZeroDivisionError.""" @@ -2761,7 +2810,7 @@ def test_global_no_lowrank(self, simple_panel_data): method="global", lambda_time_grid=[0.0], lambda_unit_grid=[0.0], - lambda_nn_grid=[float('inf')], # Disable low-rank + lambda_nn_grid=[float("inf")], # Disable low-rank n_bootstrap=10, seed=42, ) @@ -2981,18 +3030,18 @@ def test_global_loocv_score_internal(self, simple_panel_data): ) # Setup data matrices - all_units = sorted(simple_panel_data['unit'].unique()) - all_periods = sorted(simple_panel_data['period'].unique()) + all_units = sorted(simple_panel_data["unit"].unique()) + all_periods = sorted(simple_panel_data["period"].unique()) n_units = len(all_units) n_periods = len(all_periods) Y = ( - simple_panel_data.pivot(index='period', columns='unit', values='outcome') + simple_panel_data.pivot(index="period", columns="unit", values="outcome") .reindex(index=all_periods, columns=all_units) .values ) D = ( - simple_panel_data.pivot(index='period', columns='unit', values='treated') + simple_panel_data.pivot(index="period", columns="unit", values="treated") .reindex(index=all_periods, columns=all_units) .fillna(0) .astype(int) @@ -3001,23 +3050,25 @@ def test_global_loocv_score_internal(self, simple_panel_data): control_mask = D == 0 control_obs = [ - (t, i) for t in range(n_periods) for i in range(n_units) + (t, i) + for t in range(n_periods) + for i in range(n_units) if control_mask[t, i] and not np.isnan(Y[t, i]) - ][:20] # Limit for speed + ][ + :20 + ] # Limit for speed treated_periods = 3 # From fixture: n_post = 3 # Score should be finite score = trop_est._loocv_score_global( - Y, D, control_obs, 0.0, 0.0, 0.0, - treated_periods, n_units, n_periods + Y, D, control_obs, 0.0, 0.0, 0.0, treated_periods, n_units, n_periods ) assert np.isfinite(score) or np.isinf(score), "Score should be finite or inf" # Score with larger lambda_nn should still work score2 = trop_est._loocv_score_global( - Y, D, control_obs, 1.0, 1.0, 0.1, - treated_periods, n_units, n_periods + Y, D, control_obs, 1.0, 1.0, 0.1, treated_periods, n_units, n_periods ) assert np.isfinite(score2) or np.isinf(score2), "Score should be finite or inf" @@ -3025,13 +3076,13 @@ def test_global_handles_nan_outcomes(self, simple_panel_data): """Global method handles NaN outcome values gracefully.""" # Introduce NaN in some control observations data = simple_panel_data.copy() - control_mask = data['treated'] == 0 + control_mask = data["treated"] == 0 control_indices = data[control_mask].index.tolist() # Set 5 random control observations to NaN np.random.seed(42) nan_indices = np.random.choice(control_indices, size=5, replace=False) - data.loc[nan_indices, 'outcome'] = np.nan + data.loc[nan_indices, "outcome"] = np.nan trop_est = TROP( method="global", @@ -3059,13 +3110,13 @@ def test_global_with_lowrank_handles_nan(self, simple_panel_data): """Global method with low-rank handles NaN values correctly.""" # Introduce NaN in some control observations data = simple_panel_data.copy() - control_mask = data['treated'] == 0 + control_mask = data["treated"] == 0 control_indices = data[control_mask].index.tolist() # Set 3 random control observations to NaN np.random.seed(123) nan_indices = np.random.choice(control_indices, size=3, replace=False) - data.loc[nan_indices, 'outcome'] = np.nan + data.loc[nan_indices, "outcome"] = np.nan trop_est = TROP( method="global", @@ -3098,7 +3149,7 @@ def test_global_nan_exclusion_behavior(self, simple_panel_data): data_full = simple_panel_data.copy() # Identify a specific control observation to "remove" - control_mask = data_full['treated'] == 0 + control_mask = data_full["treated"] == 0 control_indices = data_full[control_mask].index.tolist() # Pick a few specific observations to remove/set to NaN @@ -3107,7 +3158,7 @@ def test_global_nan_exclusion_behavior(self, simple_panel_data): # Create version with NaN data_nan = data_full.copy() - data_nan.loc[remove_indices, 'outcome'] = np.nan + data_nan.loc[remove_indices, "outcome"] = np.nan # Create version with rows removed data_dropped = data_full.drop(remove_indices) @@ -3162,17 +3213,17 @@ def test_global_unit_no_valid_pre_gets_zero_weight(self, simple_panel_data): data = simple_panel_data.copy() # Find a control unit (unit that never has treated=1) - unit_ever_treated = data.groupby('unit')['treated'].max() + unit_ever_treated = data.groupby("unit")["treated"].max() control_units = unit_ever_treated[unit_ever_treated == 0].index.tolist() target_unit = control_units[0] # Get pre-periods (periods where this control unit has treated=0) - unit_data = data[data['unit'] == target_unit] - pre_periods = sorted(unit_data[unit_data['treated'] == 0]['period'].unique())[:5] + unit_data = data[data["unit"] == target_unit] + pre_periods = sorted(unit_data[unit_data["treated"] == 0]["period"].unique())[:5] # Set all pre-period values for target_unit to NaN - mask = (data['unit'] == target_unit) & (data['period'].isin(pre_periods)) - data.loc[mask, 'outcome'] = np.nan + mask = (data["unit"] == target_unit) & (data["period"].isin(pre_periods)) + data.loc[mask, "outcome"] = np.nan trop_est = TROP( method="global", @@ -3192,7 +3243,9 @@ def test_global_unit_no_valid_pre_gets_zero_weight(self, simple_panel_data): time="period", ) - assert np.isfinite(results.att), "ATT should be finite even with unit having no pre-period data" + assert np.isfinite( + results.att + ), "ATT should be finite even with unit having no pre-period data" assert np.isfinite(results.se), "SE should be finite" def test_global_treated_pre_nan_handling(self, simple_panel_data): @@ -3207,10 +3260,10 @@ def test_global_treated_pre_nan_handling(self, simple_panel_data): data = simple_panel_data.copy() # Find treated units and pre-periods - treated_units = data[data['treated'] == 1]['unit'].unique() + treated_units = data[data["treated"] == 1]["unit"].unique() # Pre-periods are periods where treated=0 for treated units pre_periods = sorted( - data[(data['unit'].isin(treated_units)) & (data['treated'] == 0)]['period'].unique() + data[(data["unit"].isin(treated_units)) & (data["treated"] == 0)]["period"].unique() ) assert len(pre_periods) >= 2, "Need at least 2 pre-periods for this test" @@ -3219,11 +3272,11 @@ def test_global_treated_pre_nan_handling(self, simple_panel_data): # Set ALL treated units' outcomes at target_period to NaN # This makes average_treated[target_period] = NaN - mask = (data['unit'].isin(treated_units)) & (data['period'] == target_period) - data.loc[mask, 'outcome'] = np.nan + mask = (data["unit"].isin(treated_units)) & (data["period"] == target_period) + data.loc[mask, "outcome"] = np.nan # Verify we set NaN correctly - n_nan = data.loc[mask, 'outcome'].isna().sum() + n_nan = data.loc[mask, "outcome"].isna().sum() assert n_nan == len(treated_units), f"Should have {len(treated_units)} NaN, got {n_nan}" trop_est = TROP( @@ -3262,17 +3315,14 @@ def test_global_rejects_staggered_adoption(self): is_treated_unit = i < 5 # Units 0-4 are treated, 5-9 are control for t in range(10): treated = 1 if is_treated_unit and t >= first_treat else 0 - data.append({ - 'unit': i, - 'time': t, - 'outcome': np.random.randn(), - 'treated': treated - }) + data.append( + {"unit": i, "time": t, "outcome": np.random.randn(), "treated": treated} + ) df = pd.DataFrame(data) trop = TROP(method="global") with pytest.raises(ValueError, match="staggered adoption"): - trop.fit(df, 'outcome', 'treated', 'unit', 'time') + trop.fit(df, "outcome", "treated", "unit", "time") def test_global_method_alias(self, simple_panel_data): """method='global' runs and produces a valid positive ATT.""" @@ -3306,18 +3356,18 @@ def test_global_uses_control_only_weights(self, simple_panel_data): ) # Setup data matrices - all_units = sorted(simple_panel_data['unit'].unique()) - all_periods = sorted(simple_panel_data['period'].unique()) + all_units = sorted(simple_panel_data["unit"].unique()) + all_periods = sorted(simple_panel_data["period"].unique()) n_units = len(all_units) n_periods = len(all_periods) Y = ( - simple_panel_data.pivot(index='period', columns='unit', values='outcome') + simple_panel_data.pivot(index="period", columns="unit", values="outcome") .reindex(index=all_periods, columns=all_units) .values ) D = ( - simple_panel_data.pivot(index='period', columns='unit', values='treated') + simple_panel_data.pivot(index="period", columns="unit", values="treated") .reindex(index=all_periods, columns=all_units) .fillna(0) .astype(int) @@ -3331,13 +3381,11 @@ def test_global_uses_control_only_weights(self, simple_panel_data): ) # All treated cells should have zero weight - assert np.all(delta[D == 1] == 0.0), ( - "Treated observations should have zero weight after (1-W) masking" - ) + assert np.all( + delta[D == 1] == 0.0 + ), "Treated observations should have zero weight after (1-W) masking" # Some control cells should have non-zero weight - assert np.any(delta[D == 0] > 0.0), ( - "Some control observations should have positive weight" - ) + assert np.any(delta[D == 0] > 0.0), "Some control observations should have positive weight" def test_global_tau_is_posthoc_residual(self, simple_panel_data): """Verify ATT == mean(Y - mu - alpha - beta - L) over treated cells.""" @@ -3361,9 +3409,9 @@ def test_global_tau_is_posthoc_residual(self, simple_panel_data): tau_values = [v for v in results.treatment_effects.values() if np.isfinite(v)] assert len(tau_values) > 0, "Should have treatment effects" reconstructed_att = np.mean(tau_values) - assert np.isclose(results.att, reconstructed_att, atol=1e-10), ( - f"ATT ({results.att}) should equal mean of treatment effects ({reconstructed_att})" - ) + assert np.isclose( + results.att, reconstructed_att, atol=1e-10 + ), f"ATT ({results.att}) should equal mean of treatment effects ({reconstructed_att})" def test_global_heterogeneous_treatment_effects(self, simple_panel_data): """Treatment effects are heterogeneous (not all identical) with global method.""" @@ -3371,7 +3419,7 @@ def test_global_heterogeneous_treatment_effects(self, simple_panel_data): method="global", lambda_time_grid=[0.0], lambda_unit_grid=[0.0], - lambda_nn_grid=[float('inf')], + lambda_nn_grid=[float("inf")], n_bootstrap=10, seed=42, ) @@ -3385,24 +3433,24 @@ def test_global_heterogeneous_treatment_effects(self, simple_panel_data): te_values = list(results.treatment_effects.values()) # With post-hoc extraction, effects should vary across observations - assert len(set(te_values)) > 1, ( - "Treatment effects should be heterogeneous with post-hoc extraction" - ) + assert ( + len(set(te_values)) > 1 + ), "Treatment effects should be heterogeneous with post-hoc extraction" def test_global_treated_outcome_does_not_affect_fit(self, simple_panel_data): """Perturbing treated outcomes should not change (mu, alpha, beta, L).""" - all_units = sorted(simple_panel_data['unit'].unique()) - all_periods = sorted(simple_panel_data['period'].unique()) + all_units = sorted(simple_panel_data["unit"].unique()) + all_periods = sorted(simple_panel_data["period"].unique()) n_units = len(all_units) n_periods = len(all_periods) Y = ( - simple_panel_data.pivot(index='period', columns='unit', values='outcome') + simple_panel_data.pivot(index="period", columns="unit", values="outcome") .reindex(index=all_periods, columns=all_units) .values ) D = ( - simple_panel_data.pivot(index='period', columns='unit', values='treated') + simple_panel_data.pivot(index="period", columns="unit", values="treated") .reindex(index=all_periods, columns=all_units) .fillna(0) .astype(int) @@ -3423,9 +3471,7 @@ def test_global_treated_outcome_does_not_affect_fit(self, simple_panel_data): delta = trop_est._compute_global_weights( Y, D, 1.0, 1.0, treated_periods, n_units, n_periods ) - mu1, alpha1, beta1, L1 = trop_est._solve_global_with_lowrank( - Y, delta, 0.1, 100, 1e-6 - ) + mu1, alpha1, beta1, L1 = trop_est._solve_global_with_lowrank(Y, delta, 0.1, 100, 1e-6) # Perturb treated outcomes by large amount Y_perturbed = Y.copy() @@ -3450,8 +3496,7 @@ class TestTROPNValidTreated: """Tests for n_valid_treated consistency and NaN treated outcome handling.""" @staticmethod - def _make_panel(n_units=20, n_periods=8, n_treated=5, n_post=3, - effect=2.0, seed=42): + def _make_panel(n_units=20, n_periods=8, n_treated=5, n_post=3, effect=2.0, seed=42): """Helper: generate a clean panel DataFrame.""" rng = np.random.default_rng(seed) rows = [] @@ -3463,7 +3508,7 @@ def _make_panel(n_units=20, n_periods=8, n_treated=5, n_post=3, d = 1 if (is_treated and post) else 0 if d: y += effect - rows.append({'unit': i, 'time': t, 'outcome': y, 'treated': d}) + rows.append({"unit": i, "time": t, "outcome": y, "treated": d}) return pd.DataFrame(rows) def test_global_n_treated_obs_partial_nan(self): @@ -3471,11 +3516,11 @@ def test_global_n_treated_obs_partial_nan(self): df = self._make_panel() # Inject NaN into some treated outcomes - treated_mask = (df['treated'] == 1) + treated_mask = df["treated"] == 1 treated_idx = df[treated_mask].index.tolist() n_nan = 3 for idx in treated_idx[:n_nan]: - df.loc[idx, 'outcome'] = np.nan + df.loc[idx, "outcome"] = np.nan total_treated = int(treated_mask.sum()) @@ -3489,21 +3534,22 @@ def test_global_n_treated_obs_partial_nan(self): ) with warnings.catch_warnings(): warnings.simplefilter("ignore") - results = trop_est.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop_est.fit(df, "outcome", "treated", "unit", "time") - assert results.n_treated_obs == total_treated - n_nan, \ - f"Expected {total_treated - n_nan}, got {results.n_treated_obs}" + assert ( + results.n_treated_obs == total_treated - n_nan + ), f"Expected {total_treated - n_nan}, got {results.n_treated_obs}" assert np.isfinite(results.att) def test_local_n_treated_obs_partial_nan(self): """Local method: n_treated_obs reflects only finite outcomes.""" df = self._make_panel() - treated_mask = (df['treated'] == 1) + treated_mask = df["treated"] == 1 treated_idx = df[treated_mask].index.tolist() n_nan = 3 for idx in treated_idx[:n_nan]: - df.loc[idx, 'outcome'] = np.nan + df.loc[idx, "outcome"] = np.nan total_treated = int(treated_mask.sum()) @@ -3517,10 +3563,11 @@ def test_local_n_treated_obs_partial_nan(self): ) with warnings.catch_warnings(): warnings.simplefilter("ignore") - results = trop_est.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop_est.fit(df, "outcome", "treated", "unit", "time") - assert results.n_treated_obs == total_treated - n_nan, \ - f"Expected {total_treated - n_nan}, got {results.n_treated_obs}" + assert ( + results.n_treated_obs == total_treated - n_nan + ), f"Expected {total_treated - n_nan}, got {results.n_treated_obs}" assert np.isfinite(results.att) def test_local_nan_treated_not_poison_att(self): @@ -3528,9 +3575,9 @@ def test_local_nan_treated_not_poison_att(self): df = self._make_panel(effect=3.0) # Make ONE treated outcome NaN - treated_mask = (df['treated'] == 1) + treated_mask = df["treated"] == 1 first_treated_idx = df[treated_mask].index[0] - df.loc[first_treated_idx, 'outcome'] = np.nan + df.loc[first_treated_idx, "outcome"] = np.nan trop_est = TROP( method="local", @@ -3542,7 +3589,7 @@ def test_local_nan_treated_not_poison_att(self): ) with warnings.catch_warnings(): warnings.simplefilter("ignore") - results = trop_est.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop_est.fit(df, "outcome", "treated", "unit", "time") # ATT must be finite (not NaN from NaN poisoning) assert np.isfinite(results.att), f"ATT should be finite, got {results.att}" @@ -3554,7 +3601,7 @@ def test_global_all_treated_nan_warns(self): df = self._make_panel() # Set ALL treated outcomes to NaN - df.loc[df['treated'] == 1, 'outcome'] = np.nan + df.loc[df["treated"] == 1, "outcome"] = np.nan trop_est = TROP( method="global", @@ -3566,7 +3613,7 @@ def test_global_all_treated_nan_warns(self): ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - results = trop_est.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop_est.fit(df, "outcome", "treated", "unit", "time") # Should warn about all NaN treated nan_warnings = [x for x in w if "All treated outcomes are NaN" in str(x.message)] @@ -3578,7 +3625,7 @@ def test_local_all_treated_nan_warns(self): """Local method warns when all treated outcomes are NaN.""" df = self._make_panel() - df.loc[df['treated'] == 1, 'outcome'] = np.nan + df.loc[df["treated"] == 1, "outcome"] = np.nan trop_est = TROP( method="local", @@ -3590,7 +3637,7 @@ def test_local_all_treated_nan_warns(self): ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - results = trop_est.fit(df, 'outcome', 'treated', 'unit', 'time') + results = trop_est.fit(df, "outcome", "treated", "unit", "time") nan_warnings = [x for x in w if "All treated outcomes are NaN" in str(x.message)] assert len(nan_warnings) > 0, "Should warn about all-NaN treated outcomes" @@ -3619,16 +3666,24 @@ def test_global_bootstrap_zero_draws_returns_nan_se(self): # Disable Rust backend so Python fallback path is tested, # then patch _fit_global_with_fixed_lambda to always raise - trop_global_module = sys.modules['diff_diff.trop_global'] - with patch.object(trop_global_module, 'HAS_RUST_BACKEND', False), \ - patch.object(trop_global_module, '_rust_bootstrap_trop_variance_global', None), \ - patch.object(TROP, '_fit_global_with_fixed_lambda', - side_effect=ValueError("forced failure")): + trop_global_module = sys.modules["diff_diff.trop_global"] + with ( + patch.object(trop_global_module, "HAS_RUST_BACKEND", False), + patch.object(trop_global_module, "_rust_bootstrap_trop_variance_global", None), + patch.object( + TROP, "_fit_global_with_fixed_lambda", side_effect=ValueError("forced failure") + ), + ): with warnings.catch_warnings(): warnings.simplefilter("ignore") se, dist = trop_est._bootstrap_variance_global( - df, 'outcome', 'treated', 'unit', 'time', - (1.0, 1.0, 1e10), 3, + df, + "outcome", + "treated", + "unit", + "time", + (1.0, 1.0, 1e10), + 3, ) assert np.isnan(se), f"SE should be NaN when 0 draws succeed, got {se}" @@ -3650,12 +3705,15 @@ def test_local_bootstrap_zero_draws_returns_nan_se(self): ) # Patch _fit_with_fixed_lambda to always raise - with patch.object(TROP, '_fit_with_fixed_lambda', - side_effect=ValueError("forced failure")): + with patch.object(TROP, "_fit_with_fixed_lambda", side_effect=ValueError("forced failure")): with warnings.catch_warnings(): warnings.simplefilter("ignore") se, dist = trop_est._bootstrap_variance( - df, 'outcome', 'treated', 'unit', 'time', + df, + "outcome", + "treated", + "unit", + "time", (1.0, 1.0, 1e10), ) @@ -3678,10 +3736,14 @@ def _make_panel(): y = rng.normal(0, 1) if treated and t >= 4: y += 2.0 # treatment effect - rows.append({ - "unit": i, "time": t, "outcome": y, - "treated": 1 if treated and t >= 4 else 0, - }) + rows.append( + { + "unit": i, + "time": t, + "outcome": y, + "treated": 1 if treated and t >= 4 else 0, + } + ) return pd.DataFrame(rows) def test_global_absorbing_state_error_has_remediation_guidance(self): @@ -3735,7 +3797,7 @@ def test_method_dispatch_global_uses_fit_global(self): df = self._make_panel() trop_est = TROP(method="global", n_bootstrap=2, seed=42) - with patch.object(TROP, '_fit_global', wraps=trop_est._fit_global) as mock_fg: + with patch.object(TROP, "_fit_global", wraps=trop_est._fit_global) as mock_fg: with warnings.catch_warnings(): warnings.simplefilter("ignore") trop_est.fit(df, "outcome", "treated", "unit", "time") @@ -3746,13 +3808,121 @@ def test_method_dispatch_local_does_not_use_fit_global(self): from unittest.mock import patch df = self._make_panel() - trop_est = TROP(method="local", n_bootstrap=2, seed=42, - lambda_time_grid=[0.0], lambda_unit_grid=[0.0], - lambda_nn_grid=[np.inf]) + trop_est = TROP( + method="local", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) - with patch.object(TROP, '_fit_global') as mock_fg: + with patch.object(TROP, "_fit_global") as mock_fg: with warnings.catch_warnings(): warnings.simplefilter("ignore") trop_est.fit(df, "outcome", "treated", "unit", "time") mock_fg.assert_not_called() + +class TestSilentWarningAudit: + """Tests for UserWarning emissions added by the silent warning audit.""" + + @staticmethod + def _make_panel(n_units=20, n_periods=8, n_treated=5, n_post=3, seed=42): + rng = np.random.default_rng(seed) + rows = [] + for u in range(n_units): + for t in range(n_periods): + treated = 1 if (u < n_treated and t >= n_periods - n_post) else 0 + outcome = rng.standard_normal() + (2.0 if treated else 0.0) + rows.append({"unit": u, "time": t, "outcome": outcome, "treated": treated}) + return pd.DataFrame(rows) + + def test_item5_missing_treatment_fill_warning(self): + """Item 5: Warn when NaN treatment indicators filled with 0.""" + df = self._make_panel() + # Remove some observations to make panel unbalanced + df = df.drop(df[(df["unit"] == 0) & (df["time"].isin([1, 2]))].index).reset_index(drop=True) + trop_est = TROP( + method="global", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + trop_est.fit(df, "outcome", "treated", "unit", "time") + fill_warnings = [x for x in w if "missing treatment indicator" in str(x.message)] + assert len(fill_warnings) > 0, ( + f"Expected 'missing treatment indicator' warning. " + f"Got: {[str(x.message) for x in w]}" + ) + + def test_item5_balanced_panel_no_warning(self): + """Item 5 negative: Balanced panel should not warn about missing treatment.""" + df = self._make_panel() + trop_est = TROP( + method="global", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + trop_est.fit(df, "outcome", "treated", "unit", "time") + fill_warnings = [x for x in w if "missing treatment indicator" in str(x.message)] + assert len(fill_warnings) == 0 + + def test_item6_rust_loocv_fallback_warning(self): + """Item 6: Warn when Rust LOOCV falls back to Python.""" + from unittest.mock import patch + import diff_diff.trop_global as trop_global_mod + + df = self._make_panel() + trop_est = TROP( + method="global", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) + + with ( + patch.object(trop_global_mod, "HAS_RUST_BACKEND", True), + patch.object( + trop_global_mod, "_rust_loocv_grid_search_global", side_effect=RuntimeError("test") + ), + ): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + trop_est.fit(df, "outcome", "treated", "unit", "time") + rust_warnings = [x for x in w if "Rust backend failed" in str(x.message)] + assert len(rust_warnings) > 0, ( + f"Expected 'Rust backend failed' warning. " f"Got: {[str(x.message) for x in w]}" + ) + + def test_item1_lstsq_pinv_fallback_warning(self): + """Item 1: Warn when lstsq falls back to pseudo-inverse.""" + from unittest.mock import patch + + df = self._make_panel() + trop_est = TROP( + method="global", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) + + def failing_lstsq(*args, **kwargs): + raise np.linalg.LinAlgError("test failure") + + with patch("numpy.linalg.lstsq", side_effect=failing_lstsq): + with pytest.warns(UserWarning, match="pseudo-inverse"): + trop_est.fit(df, "outcome", "treated", "unit", "time") diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index ed311e001..c552b039a 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -663,12 +663,14 @@ def test_nan_propagation(self): for h, eff in results.event_study_effects.items(): if np.isnan(eff["effect"]): nan_horizons_found += 1 - assert_nan_inference({ - "se": eff["se"], - "t_stat": eff["t_stat"], - "p_value": eff["p_value"], - "conf_int": eff["conf_int"], - }) + assert_nan_inference( + { + "se": eff["se"], + "t_stat": eff["t_stat"], + "p_value": eff["p_value"], + "conf_int": eff["conf_int"], + } + ) assert nan_horizons_found > 0, "Should have at least one Prop 5 NaN horizon" # Normal results should have finite values @@ -1064,9 +1066,7 @@ def test_bootstrap_weights_mammen(self, ci_params): """Bootstrap with mammen weights should produce valid results.""" data = generate_test_data() n_boot = ci_params.bootstrap(50) - results = TwoStageDiD( - n_bootstrap=n_boot, bootstrap_weights="mammen", seed=42 - ).fit( + results = TwoStageDiD(n_bootstrap=n_boot, bootstrap_weights="mammen", seed=42).fit( data, outcome="outcome", unit="unit", time="time", first_treat="first_treat" ) @@ -1080,9 +1080,7 @@ def test_bootstrap_weights_webb(self, ci_params): """Bootstrap with webb weights should produce valid results.""" data = generate_test_data() n_boot = ci_params.bootstrap(50) - results = TwoStageDiD( - n_bootstrap=n_boot, bootstrap_weights="webb", seed=42 - ).fit( + results = TwoStageDiD(n_bootstrap=n_boot, bootstrap_weights="webb", seed=42).fit( data, outcome="outcome", unit="unit", time="time", first_treat="first_treat" ) @@ -1096,11 +1094,13 @@ def test_bootstrap_weights_event_study(self, ci_params): """Bootstrap with non-default weights should work for event study aggregation.""" data = generate_test_data() n_boot = ci_params.bootstrap(50) - results = TwoStageDiD( - n_bootstrap=n_boot, bootstrap_weights="mammen", seed=42 - ).fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", aggregate="event_study", + results = TwoStageDiD(n_bootstrap=n_boot, bootstrap_weights="mammen", seed=42).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", ) br = results.bootstrap_results @@ -1115,11 +1115,13 @@ def test_bootstrap_weights_group(self, ci_params): """Bootstrap with non-default weights should work for group aggregation.""" data = generate_test_data() n_boot = ci_params.bootstrap(50) - results = TwoStageDiD( - n_bootstrap=n_boot, bootstrap_weights="mammen", seed=42 - ).fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", aggregate="group", + results = TwoStageDiD(n_bootstrap=n_boot, bootstrap_weights="mammen", seed=42).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="group", ) br = results.bootstrap_results @@ -1225,9 +1227,72 @@ def test_sparse_fallback_path(self): finally: ts_mod._SPARSE_DENSE_THRESHOLD = orig - np.testing.assert_allclose( - result_dense.overall_att, result_sparse.overall_att, rtol=1e-10 - ) - np.testing.assert_allclose( - result_dense.overall_se, result_sparse.overall_se, rtol=1e-10 - ) + np.testing.assert_allclose(result_dense.overall_att, result_sparse.overall_att, rtol=1e-10) + np.testing.assert_allclose(result_dense.overall_se, result_sparse.overall_se, rtol=1e-10) + + +class TestSilentWarningAudit: + """Tests for UserWarning emissions added by the silent warning audit.""" + + def test_item2_nan_ytilde_masking_warning(self): + """Item 2: Warn when NaN y_tilde values are masked.""" + # never_treated_frac=0 forces some periods without untreated obs + data = generate_test_data(n_units=50, n_periods=10, never_treated_frac=0.0, seed=42) + ts = TwoStageDiD() + with pytest.warns(UserWarning, match="non-finite imputed outcomes"): + ts.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + + def test_item3_always_treated_survey_weight_note(self): + """Item 3: Enhanced always-treated warning mentions survey weights.""" + data = generate_test_data(n_units=50, n_periods=10, never_treated_frac=0.3, seed=42) + # Shift time so min_time > 0, then set some units always-treated + data["time"] = data["time"] + 1 # now min_time = 1 + min_time = data["time"].min() + # Pick treated units and make them always-treated (first_treat=1 <= min_time=1) + treated_units = data.loc[data["first_treat"] > 0, "unit"].unique()[:3] + data.loc[data["unit"].isin(treated_units), "first_treat"] = min_time + + from diff_diff.survey import SurveyDesign + + rng = np.random.default_rng(42) + unit_weights = {u: rng.uniform(0.5, 2.0) for u in data["unit"].unique()} + data["sw"] = data["unit"].map(unit_weights) + survey = SurveyDesign(weights="sw") + + ts = TwoStageDiD() + with pytest.warns(UserWarning, match="survey weights"): + ts.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + survey_design=survey, + ) + + def test_item3_always_treated_no_survey_note_without_weights(self): + """Item 3 negative: Without survey weights, no survey note.""" + data = generate_test_data(n_units=50, n_periods=10, never_treated_frac=0.3, seed=42) + data["time"] = data["time"] + 1 + min_time = data["time"].min() + treated_units = data.loc[data["first_treat"] > 0, "unit"].unique()[:3] + data.loc[data["unit"].isin(treated_units), "first_treat"] = min_time + + ts = TwoStageDiD() + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + ts.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + survey_notes = [x for x in w if "survey weights" in str(x.message)] + assert len(survey_notes) == 0, f"Unexpected survey note: {survey_notes}" From 6ca1f5fa9d9c8b249fc272b9eea5a99aaf9e12b8 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 13:55:43 -0400 Subject: [PATCH 2/6] Address AI review P1/P2/P3 findings P1-1: Materialize NaN entries for skipped (g,t) cells in vectorized and covariate regression paths (REGISTRY contract line 350). P1-2: Reject observed treatment=NaN before pivoting in TROP to distinguish bad input from structural panel gaps. P2-1: Consolidated skip warning now reports NaN cell reasons. P2-2: Extract _validate_and_pivot_treatment shared helper to trop_local.py, eliminating duplication between trop.py and trop_global.py. P2-3/P2-4: Add regression tests for treatment=NaN rejection and NaN-materialization of skipped cells. P3-1: Update return type annotations to Tuple[Dict, Dict, Dict]. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/staggered.py | 35 +++++++++++++++++++++++++++++++---- diff_diff/trop.py | 20 ++++---------------- diff_diff/trop_global.py | 17 +++-------------- diff_diff/trop_local.py | 39 +++++++++++++++++++++++++++++++++++++++ tests/test_staggered.py | 31 +++++++++++++++++++++++++++++++ tests/test_trop.py | 30 ++++++++++++++++++++++++++++++ 6 files changed, 138 insertions(+), 34 deletions(-) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index c9e8598aa..596fb40cd 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -848,7 +848,7 @@ def _compute_all_att_gt_vectorized( treatment_groups: List[Any], time_periods: List[Any], min_period: Any, - ) -> Tuple[Dict, Dict]: + ) -> Tuple[Dict, Dict, Dict]: """ Vectorized computation of all ATT(g,t) for the no-covariates regression case. @@ -1006,6 +1006,18 @@ def _compute_all_att_gt_vectorized( ses.append(se) task_keys.append((g, t)) + # Materialize NaN entries for skipped cells (REGISTRY contract: line 350) + for g, t in skipped_missing_period + skipped_empty_cell: + group_time_effects[(g, t)] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_treated": 0, + "n_control": 0, + } + # Batch inference for all (g,t) pairs at once if task_keys: df_survey_val = precomputed.get("df_survey") @@ -1040,7 +1052,7 @@ def _compute_all_att_gt_covariate_reg( treatment_groups: List[Any], time_periods: List[Any], min_period: Any, - ) -> Tuple[Dict, Dict]: + ) -> Tuple[Dict, Dict, Dict]: """ Optimized computation of all ATT(g,t) for the covariate regression case. @@ -1358,6 +1370,18 @@ def _compute_all_att_gt_covariate_reg( stacklevel=2, ) + # Materialize NaN entries for skipped cells (REGISTRY contract: line 350) + for g, t in skipped_missing_period + skipped_empty_cell: + group_time_effects[(g, t)] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_treated": 0, + "n_control": 0, + } + # Batch inference if task_keys: # Use survey df for replicate designs (propagated from precomputed) @@ -1794,7 +1818,9 @@ def fit( stacklevel=2, ) - # Consolidated (g,t) cell skip warning + # Consolidated (g,t) cell skip warning (vectorized/covariate paths) + # General/RC paths already materialize NaN cells which are reported by + # _aggregate_simple()'s own NaN-exclusion warning. _n_missing = len(_skip_info.get("missing_period", [])) _n_empty = len(_skip_info.get("empty_cell", [])) _n_total_skipped = _n_missing + _n_empty @@ -1807,7 +1833,8 @@ def fit( if _n_empty: _parts.append(f"{_n_empty} due to zero treated or control " f"observations") warnings.warn( - f"{_n_total_skipped} (group, time) cell(s) skipped: " f"{'; '.join(_parts)}.", + f"{_n_total_skipped} (group, time) cell(s) could not be " + f"estimated (NaN): {'; '.join(_parts)}.", UserWarning, stacklevel=2, ) diff --git a/diff_diff/trop.py b/diff_diff/trop.py index a5976a13c..cbe56f771 100644 --- a/diff_diff/trop.py +++ b/diff_diff/trop.py @@ -31,7 +31,7 @@ _rust_loocv_grid_search, ) from diff_diff.trop_global import TROPGlobalMixin -from diff_diff.trop_local import TROPLocalMixin +from diff_diff.trop_local import TROPLocalMixin, _validate_and_pivot_treatment from diff_diff.trop_results import ( _LAMBDA_INF, _PrecomputedStructures, @@ -518,22 +518,10 @@ def fit( .values ) - # For D matrix, track missing values BEFORE fillna to support unbalanced panels - # Issue 3 fix: Missing observations should not trigger spurious violations - D_raw = data.pivot(index=time, columns=unit, values=treatment).reindex( - index=all_periods, columns=all_units + # For D matrix, validate observed treatment and handle unbalanced panels + D, missing_mask = _validate_and_pivot_treatment( + data, time, unit, treatment, all_periods, all_units ) - missing_mask = pd.isna(D_raw).values # True where originally missing - n_missing_treatment = int(pd.isna(D_raw).sum().sum()) - if n_missing_treatment > 0: - warnings.warn( - f"{n_missing_treatment} missing treatment indicator(s) in the " - f"(time x unit) panel matrix filled with 0 (assumed " - f"untreated). This typically occurs in unbalanced panels.", - UserWarning, - stacklevel=2, - ) - D = D_raw.fillna(0).astype(int).values # Validate D is monotonic non-decreasing per unit (absorbing state) # D[t, i] must satisfy: once D=1, it must stay 1 for all subsequent periods diff --git a/diff_diff/trop_global.py b/diff_diff/trop_global.py index f2a12aa1e..3d17e4ac1 100644 --- a/diff_diff/trop_global.py +++ b/diff_diff/trop_global.py @@ -24,7 +24,7 @@ _rust_bootstrap_trop_variance_global, _rust_loocv_grid_search_global, ) -from diff_diff.trop_local import _soft_threshold_svd +from diff_diff.trop_local import _soft_threshold_svd, _validate_and_pivot_treatment from diff_diff.trop_results import TROPResults from diff_diff.utils import safe_inference @@ -563,20 +563,9 @@ def _fit_global( .values ) - D_raw = data.pivot(index=time, columns=unit, values=treatment).reindex( - index=all_periods, columns=all_units + D, missing_mask = _validate_and_pivot_treatment( + data, time, unit, treatment, all_periods, all_units ) - missing_mask = pd.isna(D_raw).values - n_missing_treatment = int(pd.isna(D_raw).sum().sum()) - if n_missing_treatment > 0: - warnings.warn( - f"{n_missing_treatment} missing treatment indicator(s) in the " - f"(time x unit) panel matrix filled with 0 (assumed " - f"untreated). This typically occurs in unbalanced panels.", - UserWarning, - stacklevel=2, - ) - D = D_raw.fillna(0).astype(int).values # Validate absorbing state violating_units = [] diff --git a/diff_diff/trop_local.py b/diff_diff/trop_local.py index 4ca780c8c..fcfe59b4f 100644 --- a/diff_diff/trop_local.py +++ b/diff_diff/trop_local.py @@ -27,6 +27,45 @@ from diff_diff.trop_results import _PrecomputedStructures +def _validate_and_pivot_treatment(data, time, unit, treatment, all_periods, all_units): + """Validate treatment column and create D matrix with missing mask. + + Rejects observed rows with missing treatment values (data quality error), + then pivots to (time x unit) matrix. Structural gaps from unbalanced panels + are filled with 0 (assumed untreated) and flagged with a warning. + + Returns + ------- + D : ndarray + Treatment matrix (n_periods x n_units), int. + missing_mask : ndarray + Boolean mask of structurally absent cells (n_periods x n_units). + """ + n_nan_observed = int(data[treatment].isna().sum()) + if n_nan_observed > 0: + raise ValueError( + f"{n_nan_observed} observation(s) have missing treatment values. " + f"TROP requires non-missing treatment indicators for all observed " + f"rows. Remove or impute missing values before fitting." + ) + + D_raw = data.pivot(index=time, columns=unit, values=treatment).reindex( + index=all_periods, columns=all_units + ) + missing_mask = pd.isna(D_raw).values + n_missing_structural = int(missing_mask.sum()) + if n_missing_structural > 0: + warnings.warn( + f"{n_missing_structural} missing treatment indicator(s) in the " + f"(time x unit) panel matrix filled with 0 (assumed " + f"untreated). This typically occurs in unbalanced panels.", + UserWarning, + stacklevel=3, + ) + D = D_raw.fillna(0).astype(int).values + return D, missing_mask + + # Module-level convergence tolerance for SVD singular value truncation. # Singular values below this threshold after soft-thresholding are treated # as zero to improve numerical stability. diff --git a/tests/test_staggered.py b/tests/test_staggered.py index e64573c31..998265be4 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -4079,3 +4079,34 @@ def test_item4_no_skip_warning_normal_data(self): ) skip_warnings = [x for x in w if "cell(s) skipped" in str(x.message)] assert len(skip_warnings) == 0, f"Unexpected skip warning: {skip_warnings}" + + def test_skipped_cells_materialized_as_nan(self): + """P1-1: Skipped (g,t) cells appear in results with NaN effect.""" + # Sparse panel: only times 3, 5, 7. Cohort at t=5 needs base=4 + # (universal), which is missing → NaN materialization. + rng = np.random.default_rng(42) + rows = [] + for u in range(30): + for t in [3, 5, 7]: + ft = 0 if u < 10 else 5 + outcome = rng.standard_normal() + (2.0 if (ft > 0 and t >= ft) else 0.0) + rows.append({"unit": u, "time": t, "outcome": outcome, "first_treat": ft}) + data = pd.DataFrame(rows) + + # estimation_method="reg" uses the vectorized path with NaN materialization + cs = CallawaySantAnna(base_period="universal", estimation_method="reg") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + results = cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + nan_cells = {k: v for k, v in results.group_time_effects.items() if np.isnan(v["effect"])} + assert len(nan_cells) > 0, "Expected NaN-materialized cells" + for k, v in nan_cells.items(): + assert np.isnan(v["se"]), f"Cell {k} se should be NaN" + assert np.isnan(v["t_stat"]), f"Cell {k} t_stat should be NaN" + assert np.isnan(v["p_value"]), f"Cell {k} p_value should be NaN" diff --git a/tests/test_trop.py b/tests/test_trop.py index d39cfd6be..b498d4d74 100644 --- a/tests/test_trop.py +++ b/tests/test_trop.py @@ -3926,3 +3926,33 @@ def failing_lstsq(*args, **kwargs): with patch("numpy.linalg.lstsq", side_effect=failing_lstsq): with pytest.warns(UserWarning, match="pseudo-inverse"): trop_est.fit(df, "outcome", "treated", "unit", "time") + + def test_observed_treatment_nan_raises_global(self): + """P1-2: Observed treatment=NaN raises ValueError (global method).""" + df = self._make_panel() + df.loc[df.index[5], "treated"] = np.nan + trop_est = TROP( + method="global", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) + with pytest.raises(ValueError, match="missing treatment values"): + trop_est.fit(df, "outcome", "treated", "unit", "time") + + def test_observed_treatment_nan_raises_local(self): + """P1-2: Observed treatment=NaN raises ValueError (local method).""" + df = self._make_panel() + df.loc[df.index[5], "treated"] = np.nan + trop_est = TROP( + method="local", + n_bootstrap=2, + seed=42, + lambda_time_grid=[0.0], + lambda_unit_grid=[0.0], + lambda_nn_grid=[np.inf], + ) + with pytest.raises(ValueError, match="missing treatment values"): + trop_est.fit(df, "outcome", "treated", "unit", "time") From bf9874d7c784479d4e7f8d77ed1d3595e4274cd4 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 15:12:00 -0400 Subject: [PATCH 3/6] Fix CI failures and address CI review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Revert NaN materialization for skipped (g,t) cells (broke test_non_consecutive_time_periods and test_some_units_treated_first_period) - Fix inf/int64 TypeError on pandas >=2.0 in test_item8 (cast to float) - Extract _mask_nan_ytilde helper in TwoStageDiD for consistent y_tilde warning across static, event_study, and group Stage 2 paths - Remove redundant _aggregate_simple NaN warning (upstream warnings suffice) - Fix test string mismatch ("cell(s) skipped" → "could not be estimated") - Add event_study/group aggregation warning tests for TwoStageDiD - Add TODO.md entry for NaN materialization as future tech debt Co-Authored-By: Claude Opus 4.6 (1M context) --- TODO.md | 1 + diff_diff/staggered.py | 26 +--- diff_diff/staggered_aggregation.py | 15 +-- diff_diff/two_stage.py | 201 +++++++++++++++++++---------- tests/test_staggered.py | 45 ++----- tests/test_two_stage.py | 28 ++++ 6 files changed, 174 insertions(+), 142 deletions(-) diff --git a/TODO.md b/TODO.md index 0b60cb7c2..573fe4212 100644 --- a/TODO.md +++ b/TODO.md @@ -52,6 +52,7 @@ Deferred items from PR reviews that were not addressed before merge. | Issue | Location | PR | Priority | |-------|----------|----|----------| +| CallawaySantAnna: materialize NaN entries for non-estimable (g,t) cells per REGISTRY.md line 350; currently omitted from group_time_effects dict | `staggered.py` | #256 | Medium | | ImputationDiD dense `(A0'A0).toarray()` scales O((U+T+K)^2), OOM risk on large panels | `imputation.py` | #141 | Medium (deferred — only triggers when sparse solver fails) | | Multi-absorb weighted demeaning needs iterative alternating projections for N > 1 absorbed FE with survey weights; unweighted multi-absorb also uses single-pass (pre-existing, exact only for balanced panels) | `estimators.py` | #218 | Medium | | Replicate-weight survey df — **Resolved**. `df_survey = rank(replicate_weights) - 1` matching R's `survey::degf()`. For IF paths, `n_valid - 1` when dropped replicates reduce effective count. | `survey.py` | #238 | Resolved | diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 596fb40cd..c639b6657 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -1006,18 +1006,6 @@ def _compute_all_att_gt_vectorized( ses.append(se) task_keys.append((g, t)) - # Materialize NaN entries for skipped cells (REGISTRY contract: line 350) - for g, t in skipped_missing_period + skipped_empty_cell: - group_time_effects[(g, t)] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_treated": 0, - "n_control": 0, - } - # Batch inference for all (g,t) pairs at once if task_keys: df_survey_val = precomputed.get("df_survey") @@ -1370,18 +1358,6 @@ def _compute_all_att_gt_covariate_reg( stacklevel=2, ) - # Materialize NaN entries for skipped cells (REGISTRY contract: line 350) - for g, t in skipped_missing_period + skipped_empty_cell: - group_time_effects[(g, t)] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_treated": 0, - "n_control": 0, - } - # Batch inference if task_keys: # Use survey df for replicate designs (propagated from precomputed) @@ -1834,7 +1810,7 @@ def fit( _parts.append(f"{_n_empty} due to zero treated or control " f"observations") warnings.warn( f"{_n_total_skipped} (group, time) cell(s) could not be " - f"estimated (NaN): {'; '.join(_parts)}.", + f"estimated: {'; '.join(_parts)}.", UserWarning, stacklevel=2, ) diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index 7f122a255..fd364defb 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -105,18 +105,11 @@ def _aggregate_simple( weights = np.array(weights_list, dtype=float) groups_for_gt = np.array(groups_for_gt) - # Exclude NaN effects from aggregation (R's aggte() convention) + # Exclude NaN effects from aggregation (R's aggte() convention). + # No warning here — upstream code (consolidated skip warning in fit(), + # covariate regression NaN warning) already informed the user. finite_mask = np.isfinite(effects) - n_nan = int(np.sum(~finite_mask)) - if n_nan > 0: - import warnings - - warnings.warn( - f"{n_nan} group-time effect(s) are NaN and excluded from overall ATT " - "aggregation. Inspect group_time_effects for details.", - UserWarning, - stacklevel=2, - ) + if not np.all(finite_mask): effects = effects[finite_mask] weights = weights[finite_mask] gt_pairs = [gt for gt, m in zip(gt_pairs, finite_mask) if m] diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 5b0a892bc..c07674f2d 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -247,9 +247,7 @@ def fit( _resolve_survey_for_fit(survey_design, data, "analytical") ) - _uses_replicate_ts = ( - resolved_survey is not None and resolved_survey.uses_replicate_variance - ) + _uses_replicate_ts = resolved_survey is not None and resolved_survey.uses_replicate_variance if _uses_replicate_ts and self.n_bootstrap > 0: raise ValueError( "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " @@ -527,25 +525,46 @@ def fit( if aggregate in ("event_study", "all"): event_study_effects = self._stage2_event_study( - df=df, unit=unit, time=time, first_treat=first_treat, - covariates=covariates, omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, unit_fe=unit_fe, time_fe=time_fe, - grand_mean=grand_mean, delta_hat=delta_hat, - cluster_var=cluster_var, treatment_groups=treatment_groups, - ref_period=ref_period, balance_e=balance_e, - kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, - survey_weight_type=survey_weight_type, survey_df=_survey_df, + df=df, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + ref_period=ref_period, + balance_e=balance_e, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + survey_weight_type=survey_weight_type, + survey_df=_survey_df, ) if aggregate in ("group", "all"): group_effects = self._stage2_group( - df=df, unit=unit, time=time, first_treat=first_treat, - covariates=covariates, omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, unit_fe=unit_fe, time_fe=time_fe, - grand_mean=grand_mean, delta_hat=delta_hat, - cluster_var=cluster_var, treatment_groups=treatment_groups, - kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, - survey_weight_type=survey_weight_type, survey_df=_survey_df, + df=df, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + survey_weight_type=survey_weight_type, + survey_df=_survey_df, ) # Replicate variance override: derive keys from actual outputs, then refit @@ -556,12 +575,12 @@ def fit( # Derive keys from actual outputs (excludes filtered/Prop5 horizons) _sorted_es_periods_ts = sorted( - e for e in (event_study_effects or {}).keys() + e + for e in (event_study_effects or {}).keys() if np.isfinite(event_study_effects[e]["effect"]) ) _sorted_groups_ts = sorted( - g for g in (group_effects or {}).keys() - if np.isfinite(group_effects[g]["effect"]) + g for g in (group_effects or {}).keys() if np.isfinite(group_effects[g]["effect"]) ) _n_es_ts = len(_sorted_es_periods_ts) _n_grp_ts = len(_sorted_groups_ts) @@ -573,49 +592,92 @@ def fit( def _refit_ts(w_r): ufe_r, tfe_r, gm_r, delta_r, kcm_r = self._fit_untreated_model( - df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + df, + outcome, + unit, + time, + covariates, + omega_0_mask, + weights=w_r, ) y_tilde_r = self._residualize( - df, outcome, unit, time, covariates, - ufe_r, tfe_r, gm_r, delta_r, + df, + outcome, + unit, + time, + covariates, + ufe_r, + tfe_r, + gm_r, + delta_r, ) df_tmp = df.copy() df_tmp["_y_tilde"] = y_tilde_r results = [] att_r, _ = self._stage2_static( - df=df_tmp, unit=unit, time=time, first_treat=first_treat, - covariates=covariates, omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, - grand_mean=gm_r, delta_hat=delta_r, cluster_var=cluster_var, - kept_cov_mask=kcm_r, survey_weights=w_r, + df=df_tmp, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=ufe_r, + time_fe=tfe_r, + grand_mean=gm_r, + delta_hat=delta_r, + cluster_var=cluster_var, + kept_cov_mask=kcm_r, + survey_weights=w_r, survey_weight_type="pweight", ) results.append(att_r) if _sorted_es_periods_ts: es_r = self._stage2_event_study( - df=df_tmp, unit=unit, time=time, first_treat=first_treat, - covariates=covariates, omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, - grand_mean=gm_r, delta_hat=delta_r, - cluster_var=cluster_var, treatment_groups=treatment_groups, - ref_period=ref_period, balance_e=balance_e, - kept_cov_mask=kcm_r, survey_weights=w_r, - survey_weight_type="pweight", survey_df=None, + df=df_tmp, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=ufe_r, + time_fe=tfe_r, + grand_mean=gm_r, + delta_hat=delta_r, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + ref_period=ref_period, + balance_e=balance_e, + kept_cov_mask=kcm_r, + survey_weights=w_r, + survey_weight_type="pweight", + survey_df=None, ) for e in _sorted_es_periods_ts: results.append(es_r[e]["effect"] if e in es_r else np.nan) if _sorted_groups_ts: grp_r = self._stage2_group( - df=df_tmp, unit=unit, time=time, first_treat=first_treat, - covariates=covariates, omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, - grand_mean=gm_r, delta_hat=delta_r, - cluster_var=cluster_var, treatment_groups=treatment_groups, - kept_cov_mask=kcm_r, survey_weights=w_r, - survey_weight_type="pweight", survey_df=None, + df=df_tmp, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=ufe_r, + time_fe=tfe_r, + grand_mean=gm_r, + delta_hat=delta_r, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + kept_cov_mask=kcm_r, + survey_weights=w_r, + survey_weight_type="pweight", + survey_df=None, ) for g in _sorted_groups_ts: results.append(grp_r[g]["effect"] if g in grp_r else np.nan) @@ -652,9 +714,9 @@ def _refit_ts(w_r): # Override group SEs (only for identified effects) for j, g in enumerate(_sorted_groups_ts): if group_effects is not None and g in group_effects: - se_g = float(np.sqrt(max( - _vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0 - ))) + se_g = float( + np.sqrt(max(_vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0)) + ) eff_g = group_effects[g]["effect"] t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) group_effects[g]["se"] = se_g @@ -1024,6 +1086,26 @@ def _residualize( # Stage 2 specifications # ========================================================================= + @staticmethod + def _mask_nan_ytilde(y_tilde): + """Mask non-finite y_tilde values and warn if any found. + + Returns the boolean mask of non-finite values. Modifies y_tilde in-place + (sets NaN values to 0.0). + """ + nan_mask = ~np.isfinite(y_tilde) + if nan_mask.any(): + n_nan = int(nan_mask.sum()) + warnings.warn( + f"{n_nan} observation(s) have non-finite imputed outcomes " + f"(y_tilde) from unidentified fixed effects. These " + f"observations are excluded from ATT estimation.", + UserWarning, + stacklevel=3, + ) + y_tilde[nan_mask] = 0.0 + return nan_mask + def _stage2_static( self, df: pd.DataFrame, @@ -1048,21 +1130,7 @@ def _stage2_static( Returns (att, se). """ y_tilde = df["_y_tilde"].values.copy() - - # Handle NaN y_tilde (from unidentified FEs — e.g., rank condition violations) - # Set to 0 so solve_ols doesn't reject; these obs have X_2=0 (untreated) - # or contribute NaN treatment effects (excluded from point estimate). - nan_mask = ~np.isfinite(y_tilde) - if nan_mask.any(): - n_nan = int(nan_mask.sum()) - warnings.warn( - f"{n_nan} observation(s) have non-finite imputed outcomes " - f"(y_tilde) from unidentified fixed effects. These " - f"observations are excluded from ATT estimation.", - UserWarning, - stacklevel=2, - ) - y_tilde[nan_mask] = 0.0 + nan_mask = self._mask_nan_ytilde(y_tilde) D = omega_1_mask.values.astype(float) # Zero out treatment indicator for NaN y_tilde obs (don't count in ATT) @@ -1131,10 +1199,7 @@ def _stage2_event_study( ) -> Dict[int, Dict[str, Any]]: """Event study Stage 2: OLS of y_tilde on relative-time dummies.""" y_tilde = df["_y_tilde"].values.copy() - # Handle NaN y_tilde (unidentified FEs) - nan_mask = ~np.isfinite(y_tilde) - if nan_mask.any(): - y_tilde[nan_mask] = 0.0 + nan_mask = self._mask_nan_ytilde(y_tilde) rel_times = df["_rel_time"].values n = len(df) @@ -1350,9 +1415,7 @@ def _stage2_group( ) -> Dict[Any, Dict[str, Any]]: """Group (cohort) Stage 2: OLS of y_tilde on cohort dummies.""" y_tilde = df["_y_tilde"].values.copy() - nan_mask = ~np.isfinite(y_tilde) - if nan_mask.any(): - y_tilde[nan_mask] = 0.0 + nan_mask = self._mask_nan_ytilde(y_tilde) n = len(df) # Build Stage 2 design: one column per cohort (no intercept) diff --git a/tests/test_staggered.py b/tests/test_staggered.py index 998265be4..73c16a364 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -3963,6 +3963,8 @@ def test_item8_inf_to_zero_warning_in_fit(self): data = generate_staggered_data(seed=42) # Set some units to inf (never-treated encoding) + # Cast to float first for pandas >=2.0 compatibility + data["first_treat"] = data["first_treat"].astype(float) never_units = data.loc[data["first_treat"] == 0, "unit"].unique()[:5] data.loc[data["unit"].isin(never_units), "first_treat"] = np.inf @@ -3981,7 +3983,8 @@ def test_item8_inf_to_zero_warning_in_diagnose_propensity(self): import warnings data = generate_staggered_data_with_covariates(seed=42) - # Set some units to inf + # Cast to float first for pandas >=2.0 compatibility + data["first_treat"] = data["first_treat"].astype(float) never_units = data.loc[data["first_treat"] == 0, "unit"].unique()[:5] data.loc[data["unit"].isin(never_units), "first_treat"] = np.inf @@ -4054,13 +4057,12 @@ def test_item4_consolidated_skip_warning(self): except ValueError: pass # May fail if all cells skipped - skip_warnings = [x for x in w if "cell(s) skipped" in str(x.message)] + skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] # Should have at least one skip warning if cells were skipped # (The sparse panel structure should cause missing period skips) if skip_warnings: - assert "missing base/post period" in str( - skip_warnings[0].message - ) or "zero treated or control" in str(skip_warnings[0].message) + msg = str(skip_warnings[0].message) + assert "missing base/post period" in msg or "zero treated" in msg def test_item4_no_skip_warning_normal_data(self): """Item 4 negative: No skip warning on well-formed balanced data.""" @@ -4077,36 +4079,5 @@ def test_item4_no_skip_warning_normal_data(self): time="time", first_treat="first_treat", ) - skip_warnings = [x for x in w if "cell(s) skipped" in str(x.message)] + skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] assert len(skip_warnings) == 0, f"Unexpected skip warning: {skip_warnings}" - - def test_skipped_cells_materialized_as_nan(self): - """P1-1: Skipped (g,t) cells appear in results with NaN effect.""" - # Sparse panel: only times 3, 5, 7. Cohort at t=5 needs base=4 - # (universal), which is missing → NaN materialization. - rng = np.random.default_rng(42) - rows = [] - for u in range(30): - for t in [3, 5, 7]: - ft = 0 if u < 10 else 5 - outcome = rng.standard_normal() + (2.0 if (ft > 0 and t >= ft) else 0.0) - rows.append({"unit": u, "time": t, "outcome": outcome, "first_treat": ft}) - data = pd.DataFrame(rows) - - # estimation_method="reg" uses the vectorized path with NaN materialization - cs = CallawaySantAnna(base_period="universal", estimation_method="reg") - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - results = cs.fit( - data, - outcome="outcome", - unit="unit", - time="time", - first_treat="first_treat", - ) - nan_cells = {k: v for k, v in results.group_time_effects.items() if np.isnan(v["effect"])} - assert len(nan_cells) > 0, "Expected NaN-materialized cells" - for k, v in nan_cells.items(): - assert np.isnan(v["se"]), f"Cell {k} se should be NaN" - assert np.isnan(v["t_stat"]), f"Cell {k} t_stat should be NaN" - assert np.isnan(v["p_value"]), f"Cell {k} p_value should be NaN" diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index c552b039a..bae7ff328 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -1296,3 +1296,31 @@ def test_item3_always_treated_no_survey_note_without_weights(self): ) survey_notes = [x for x in w if "survey weights" in str(x.message)] assert len(survey_notes) == 0, f"Unexpected survey note: {survey_notes}" + + def test_item2_nan_ytilde_event_study(self): + """Item 2: y_tilde warning fires for aggregate='event_study'.""" + data = generate_test_data(n_units=50, n_periods=10, never_treated_frac=0.0, seed=42) + ts = TwoStageDiD() + with pytest.warns(UserWarning, match="non-finite imputed outcomes"): + ts.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + def test_item2_nan_ytilde_group(self): + """Item 2: y_tilde warning fires for aggregate='group'.""" + data = generate_test_data(n_units=50, n_periods=10, never_treated_frac=0.0, seed=42) + ts = TwoStageDiD() + with pytest.warns(UserWarning, match="non-finite imputed outcomes"): + ts.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="group", + ) From 0df5d33992578213d6348dff316cc2291c7db1e3 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 15:29:11 -0400 Subject: [PATCH 4/6] Address AI review P1/P2/P3 findings Update REGISTRY.md to document actual behavior: non-estimable (g,t) cells are omitted from group_time_effects (not NaN-materialized). NaN materialization tracked as future work in TODO.md. Strengthen skip-warning test: unconditional assertion, deterministic test data with both succeeding and failing cohorts. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/methodology/REGISTRY.md | 2 +- tests/test_staggered.py | 49 ++++++++++++++++++------------------ 2 files changed, 26 insertions(+), 25 deletions(-) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index bce12bd76..f1ea62cf4 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -347,7 +347,7 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1: *Edge cases:* - Groups with single observation: included but may have high variance -- Missing group-time cells: ATT(g,t) set to NaN +- Missing group-time cells: omitted from `group_time_effects` (non-estimable cells are not included); a consolidated warning is emitted for the `reg` estimation path listing skip reasons - **Note:** When `balance_e` is specified, cohorts with NaN effects at the anchor horizon are excluded from the balanced panel - Anticipation: `anticipation` parameter shifts reference period - Group aggregation includes periods t >= g - anticipation (not just t >= g) diff --git a/tests/test_staggered.py b/tests/test_staggered.py index 73c16a364..bf090f3e6 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -4021,17 +4021,23 @@ def test_item4_consolidated_skip_warning(self): """Item 4: Consolidated warning when (g,t) cells are skipped.""" import warnings - # Create data with a cohort whose base period is outside the panel range. - # Cohort with first_treat=1 needs base_period=0 (universal) which exists, - # but cohort with first_treat=2 needs base_period=1 — which exists. - # To trigger missing period skips, create a very sparse panel. + # Two cohorts: g=4 succeeds (base=3 exists), g=6 fails (base=5 + # exists but post periods 7 need base=5 which exists — so use + # g=8 whose base=7 exists). Actually simplest: periods [1,2,4,5] + # with cohort g=4 (base=3 missing → some skips) and g=2 (base=1 + # exists → succeeds). rng = np.random.default_rng(42) - n_units = 30 + n_units = 40 rows = [] for u in range(n_units): - # Only include time periods 3, 5, 7 (skip 0, 1, 2, 4, 6) - for t in [3, 5, 7]: - ft = 0 if u < 10 else 5 # Cohort at t=5, or never-treated + for t in [1, 2, 4, 5]: + # u < 10: never-treated; u < 25: cohort g=2; rest: cohort g=4 + if u < 10: + ft = 0 + elif u < 25: + ft = 2 # base=1 exists → succeeds + else: + ft = 4 # base=3 missing → skipped outcome = rng.standard_normal() + (2.0 if (ft > 0 and t >= ft) else 0.0) rows.append( { @@ -4043,26 +4049,21 @@ def test_item4_consolidated_skip_warning(self): ) data = pd.DataFrame(rows) - cs = CallawaySantAnna(base_period="universal") + cs = CallawaySantAnna(base_period="universal", estimation_method="reg") with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - try: - cs.fit( - data, - outcome="outcome", - unit="unit", - time="time", - first_treat="first_treat", - ) - except ValueError: - pass # May fail if all cells skipped + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] - # Should have at least one skip warning if cells were skipped - # (The sparse panel structure should cause missing period skips) - if skip_warnings: - msg = str(skip_warnings[0].message) - assert "missing base/post period" in msg or "zero treated" in msg + assert len(skip_warnings) > 0, "Expected consolidated skip warning" + msg = str(skip_warnings[0].message) + assert "missing base/post period" in msg def test_item4_no_skip_warning_normal_data(self): """Item 4 negative: No skip warning on well-formed balanced data.""" From c339fd619eed21d9544bbc5f42be40676403324b Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 16:02:35 -0400 Subject: [PATCH 5/6] Address AI review P1/P2/P3 findings Extend consolidated skip warning to all CS estimation paths (general DR/IPW and repeated cross-section), not just vectorized/covariate reg. Update REGISTRY.md with labeled **Note:** documenting omission behavior per reviewer rubric requirements. Fix misleading comment in _aggregate_simple about upstream warnings. Add DR-path and panel=False skip warning tests. Filter expected skip warnings in power tests that use not_yet_treated. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/staggered.py | 15 +++++++--- diff_diff/staggered_aggregation.py | 4 +-- docs/methodology/REGISTRY.md | 3 +- tests/test_power.py | 6 ++++ tests/test_staggered.py | 47 ++++++++++++++++++++++++++++++ 5 files changed, 68 insertions(+), 7 deletions(-) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index c639b6657..c00111a6e 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -1623,6 +1623,7 @@ def fit( has_survey = resolved_survey is not None _skip_info = {"missing_period": [], "empty_cell": []} + _n_skipped_other = 0 if not self.panel: # --- Repeated cross-section path --- @@ -1678,6 +1679,8 @@ def fit( if inf_info is not None: influence_func_info[(g, t)] = inf_info + else: + _n_skipped_other += 1 elif covariates is None and self.estimation_method == "reg": # Fast vectorized path for the common no-covariates regression case @@ -1767,6 +1770,8 @@ def fit( if inf_info is not None: influence_func_info[(g, t)] = inf_info + else: + _n_skipped_other += 1 if not group_time_effects: raise ValueError( @@ -1794,12 +1799,10 @@ def fit( stacklevel=2, ) - # Consolidated (g,t) cell skip warning (vectorized/covariate paths) - # General/RC paths already materialize NaN cells which are reported by - # _aggregate_simple()'s own NaN-exclusion warning. + # Consolidated (g,t) cell skip warning (all paths) _n_missing = len(_skip_info.get("missing_period", [])) _n_empty = len(_skip_info.get("empty_cell", [])) - _n_total_skipped = _n_missing + _n_empty + _n_total_skipped = _n_missing + _n_empty + _n_skipped_other if _n_total_skipped > 0: _parts = [] if _n_missing: @@ -1808,6 +1811,10 @@ def fit( ) if _n_empty: _parts.append(f"{_n_empty} due to zero treated or control " f"observations") + if _n_skipped_other: + _parts.append( + f"{_n_skipped_other} due to insufficient data or " f"non-estimable cells" + ) warnings.warn( f"{_n_total_skipped} (group, time) cell(s) could not be " f"estimated: {'; '.join(_parts)}.", diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index fd364defb..d90139d4b 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -106,8 +106,8 @@ def _aggregate_simple( groups_for_gt = np.array(groups_for_gt) # Exclude NaN effects from aggregation (R's aggte() convention). - # No warning here — upstream code (consolidated skip warning in fit(), - # covariate regression NaN warning) already informed the user. + # No warning here — fit() emits a consolidated skip warning covering + # all estimation paths (vectorized, covariate, general, RC). finite_mask = np.isfinite(effects) if not np.all(finite_mask): effects = effects[finite_mask] diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index f1ea62cf4..7668d09ca 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -347,7 +347,8 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1: *Edge cases:* - Groups with single observation: included but may have high variance -- Missing group-time cells: omitted from `group_time_effects` (non-estimable cells are not included); a consolidated warning is emitted for the `reg` estimation path listing skip reasons +- Missing group-time cells: omitted from `group_time_effects` with a consolidated warning listing skip reasons and counts + - **Note:** Non-estimable cells (missing base/post period, zero treated/control, insufficient data) are omitted rather than stored as NaN. A consolidated UserWarning is emitted from `fit()` across all estimation paths. R's `did` package also omits these cells from `aggte()` results. - **Note:** When `balance_e` is specified, cohorts with NaN effects at the anchor horizon are excluded from the balanced panel - Anticipation: `anticipation` parameter shifts reference period - Group aggregation includes periods t >= g - anticipation (not just t >= g) diff --git a/tests/test_power.py b/tests/test_power.py index 47e996a21..116a97a50 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -1513,6 +1513,8 @@ def _custom_staggered(**kwargs): with warnings.catch_warnings(): warnings.simplefilter("error", UserWarning) + # Skip warning is expected for not_yet_treated (some cells non-estimable) + warnings.filterwarnings("ignore", message=".*could not be estimated.*") simulate_power( CallawaySantAnna(control_group="not_yet_treated"), data_generator=_custom_staggered, @@ -1533,6 +1535,8 @@ def test_staggered_dgp_no_warn_with_dgp_kwargs_override(self): """data_generator_kwargs with cohort_periods suppresses warning.""" with warnings.catch_warnings(): warnings.simplefilter("error", UserWarning) + # Skip warning is expected for not_yet_treated (some cells non-estimable) + warnings.filterwarnings("ignore", message=".*could not be estimated.*") result = simulate_power( CallawaySantAnna(control_group="not_yet_treated"), n_periods=6, @@ -2014,6 +2018,8 @@ def test_staggered_multi_cohort_no_warn(self): """CS with cohort_periods=[2, 4] does NOT warn.""" with warnings.catch_warnings(): warnings.simplefilter("error") + # Skip warning is expected for not_yet_treated (some cells non-estimable) + warnings.filterwarnings("ignore", message=".*could not be estimated.*") simulate_power( CallawaySantAnna(control_group="not_yet_treated"), n_units=60, diff --git a/tests/test_staggered.py b/tests/test_staggered.py index bf090f3e6..ef47ff0d2 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -4082,3 +4082,50 @@ def test_item4_no_skip_warning_normal_data(self): ) skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] assert len(skip_warnings) == 0, f"Unexpected skip warning: {skip_warnings}" + + def test_skip_warning_dr_path(self): + """Skip warning fires for default DR path (general path).""" + data = generate_staggered_data( + n_units=50, + n_periods=6, + n_cohorts=3, + never_treated_frac=0.0, + seed=42, + ) + cs = CallawaySantAnna(control_group="not_yet_treated") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] + assert len(skip_warnings) > 0, "Expected skip warning for DR path" + assert "insufficient data" in str(skip_warnings[0].message) + + def test_skip_warning_panel_false(self): + """Skip warning fires for panel=False (RC path).""" + data = generate_staggered_data( + n_units=80, + n_periods=6, + n_cohorts=3, + never_treated_frac=0.0, + seed=42, + ) + # panel=False needs unique unit IDs (repeated cross-section) + data["unit"] = np.arange(len(data)) + cs = CallawaySantAnna(panel=False, control_group="not_yet_treated") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] + assert len(skip_warnings) > 0, "Expected skip warning for RC path" From 7b290d29ea6d91770a780cb06e0d285872ddd84e Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 16:15:49 -0400 Subject: [PATCH 6/6] Fix survey zero-mass skip warning gap and TODO contradiction Make survey zero-effective-mass branches consistent with skip contract: - _compute_att_gt_fast: return None (not np.nan) for zero-mass cells - vectorized survey path: track zero-mass skips in skipped_empty_cell - _compute_att_gt_rc: return None (not np.nan) for zero-mass cells All three now feed into the consolidated skip warning. Update TODO.md to describe actual remaining work (downstream consumer updates for optional NaN materialization) instead of contradicting the REGISTRY.md omission contract. Add survey-weighted zero-mass skip warning test. Co-Authored-By: Claude Opus 4.6 (1M context) --- TODO.md | 2 +- diff_diff/staggered.py | 9 +++++---- tests/test_staggered.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/TODO.md b/TODO.md index 573fe4212..d23caa201 100644 --- a/TODO.md +++ b/TODO.md @@ -52,7 +52,7 @@ Deferred items from PR reviews that were not addressed before merge. | Issue | Location | PR | Priority | |-------|----------|----|----------| -| CallawaySantAnna: materialize NaN entries for non-estimable (g,t) cells per REGISTRY.md line 350; currently omitted from group_time_effects dict | `staggered.py` | #256 | Medium | +| CallawaySantAnna: consider materializing NaN entries for non-estimable (g,t) cells in group_time_effects dict (currently omitted with consolidated warning); would require updating downstream consumers (event study, balance_e, aggregation) | `staggered.py` | #256 | Low | | ImputationDiD dense `(A0'A0).toarray()` scales O((U+T+K)^2), OOM risk on large panels | `imputation.py` | #141 | Medium (deferred — only triggers when sparse solver fails) | | Multi-absorb weighted demeaning needs iterative alternating projections for N > 1 absorbed FE with survey weights; unweighted multi-absorb also uses single-pass (pre-existing, exact only for balanced panels) | `estimators.py` | #218 | Medium | | Replicate-weight survey df — **Resolved**. `df_survey = rank(replicate_weights) - 1` matching R's `survey::degf()`. For IF paths, `n_valid - 1` when dropped replicates reduce effective count. | `survey.py` | #238 | Resolved | diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index c00111a6e..f4e63cae1 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -731,9 +731,9 @@ def _compute_att_gt_fast( # Guard against zero effective mass after subpopulation filtering if sw_treated is not None and np.sum(sw_treated) <= 0: - return np.nan, np.nan, 0, 0, None, None + return None, 0.0, 0, 0, None, None if sw_control is not None and np.sum(sw_control) <= 0: - return np.nan, np.nan, 0, 0, None, None + return None, 0.0, 0, 0, None, None # Get covariates if specified (from the base period) X_treated = None @@ -947,6 +947,7 @@ def _compute_all_att_gt_vectorized( sw_c = survey_w[control_valid] # Guard against zero effective mass if np.sum(sw_t) <= 0 or np.sum(sw_c) <= 0: + skipped_empty_cell.append((g, t)) continue sw_t_norm = sw_t / np.sum(sw_t) sw_c_norm = sw_c / np.sum(sw_c) @@ -2892,9 +2893,9 @@ def _compute_att_gt_rc( # Guard against zero effective mass if sw_gt is not None: if np.sum(sw_gt) <= 0 or np.sum(sw_gs) <= 0: - return np.nan, np.nan, 0, 0, None, None + return None, 0.0, 0, 0, None, None if np.sum(sw_ct) <= 0 or np.sum(sw_cs) <= 0: - return np.nan, np.nan, 0, 0, None, None + return None, 0.0, 0, 0, None, None # Get covariates if specified obs_covariates = precomputed.get("obs_covariates") diff --git a/tests/test_staggered.py b/tests/test_staggered.py index ef47ff0d2..4eb0e5124 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -4129,3 +4129,36 @@ def test_skip_warning_panel_false(self): ) skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] assert len(skip_warnings) > 0, "Expected skip warning for RC path" + + def test_skip_warning_survey_zero_mass(self): + """Skip warning fires when survey weights produce zero effective mass.""" + from diff_diff.survey import SurveyDesign + + data = generate_staggered_data( + n_units=60, + n_periods=6, + n_cohorts=2, + never_treated_frac=0.3, + seed=42, + ) + # Set survey weights to 0 for ALL units in one cohort to force + # zero effective mass in that cohort's cells + data["sw"] = 1.0 + first_cohort = sorted(data.loc[data["first_treat"] > 0, "first_treat"].unique())[0] + cohort_units = data.loc[data["first_treat"] == first_cohort, "unit"].unique() + data.loc[data["unit"].isin(cohort_units), "sw"] = 0.0 + + survey = SurveyDesign(weights="sw") + cs = CallawaySantAnna(estimation_method="reg") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + cs.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + survey_design=survey, + ) + skip_warnings = [x for x in w if "could not be estimated" in str(x.message)] + assert len(skip_warnings) > 0, "Expected skip warning for zero-mass survey cells"