diff-diff 1.2.0__tar.gz → 1.2.1__tar.gz

{diff_diff-1.2.0 → diff_diff-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 1.2.0
+Version: 1.2.1
 Summary: A library for Difference-in-Differences causal inference analysis
 Author: diff-diff contributors
 License-Expression: MIT

{diff_diff-1.2.0 → diff_diff-1.2.1}/diff_diff/__init__.py

@@ -98,7 +98,7 @@ from diff_diff.visualization import (
     plot_sensitivity,
 )
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"
 __all__ = [
     # Estimators
     "DifferenceInDifferences",

{diff_diff-1.2.0 → diff_diff-1.2.1}/diff_diff/staggered.py

@@ -1603,30 +1603,25 @@ class CallawaySantAnna:
                         weights * bootstrap_atts_gt[b, gt_indices]
                     )
 
-        # Compute bootstrap statistics
-        # ATT(g,t) statistics
+        # Compute bootstrap statistics for ATT(g,t)
         gt_ses = {}
         gt_cis = {}
         gt_p_values = {}
 
         for j, gt in enumerate(gt_pairs):
-            original_effect = original_atts[j]
-            boot_dist = bootstrap_atts_gt[:, j]
-
-            se = float(np.std(boot_dist, ddof=1))
-            ci = self._compute_percentile_ci(boot_dist, self.alpha)
-            p_value = self._compute_bootstrap_pvalue(original_effect, boot_dist)
-
+            se, ci, p_value = self._compute_effect_bootstrap_stats(
+                original_atts[j], bootstrap_atts_gt[:, j]
+            )
             gt_ses[gt] = se
             gt_cis[gt] = ci
             gt_p_values[gt] = p_value
 
-        # Overall ATT statistics
-        overall_se = float(np.std(bootstrap_overall, ddof=1))
-        overall_ci = self._compute_percentile_ci(bootstrap_overall, self.alpha)
-        overall_p_value = self._compute_bootstrap_pvalue(original_overall, bootstrap_overall)
+        # Compute bootstrap statistics for overall ATT
+        overall_se, overall_ci, overall_p_value = self._compute_effect_bootstrap_stats(
+            original_overall, bootstrap_overall
+        )
 
-        # Event study statistics
+        # Compute bootstrap statistics for event study effects
         event_study_ses = None
         event_study_cis = None
         event_study_p_values = None
@@ -1637,16 +1632,14 @@ class CallawaySantAnna:
             event_study_p_values = {}
 
             for e in rel_periods:
-                original_effect = event_study_info[e]['effect']
-                boot_dist = bootstrap_event_study[e]
-
-                event_study_ses[e] = float(np.std(boot_dist, ddof=1))
-                event_study_cis[e] = self._compute_percentile_ci(boot_dist, self.alpha)
-                event_study_p_values[e] = self._compute_bootstrap_pvalue(
-                    original_effect, boot_dist
+                se, ci, p_value = self._compute_effect_bootstrap_stats(
+                    event_study_info[e]['effect'], bootstrap_event_study[e]
                 )
+                event_study_ses[e] = se
+                event_study_cis[e] = ci
+                event_study_p_values[e] = p_value
 
-        # Group effect statistics
+        # Compute bootstrap statistics for group effects
         group_effect_ses = None
         group_effect_cis = None
         group_effect_p_values = None
@@ -1657,14 +1650,12 @@ class CallawaySantAnna:
             group_effect_p_values = {}
 
             for g in groups:
-                original_effect = group_agg_info[g]['effect']
-                boot_dist = bootstrap_group[g]
-
-                group_effect_ses[g] = float(np.std(boot_dist, ddof=1))
-                group_effect_cis[g] = self._compute_percentile_ci(boot_dist, self.alpha)
-                group_effect_p_values[g] = self._compute_bootstrap_pvalue(
-                    original_effect, boot_dist
+                se, ci, p_value = self._compute_effect_bootstrap_stats(
+                    group_agg_info[g]['effect'], bootstrap_group[g]
                 )
+                group_effect_ses[g] = se
+                group_effect_cis[g] = ci
+                group_effect_p_values[g] = p_value
 
         return CSBootstrapResults(
             n_bootstrap=self.n_bootstrap,
@@ -1817,6 +1808,35 @@ class CallawaySantAnna:
 
         return float(p_value)
 
+    def _compute_effect_bootstrap_stats(
+        self,
+        original_effect: float,
+        boot_dist: np.ndarray,
+    ) -> Tuple[float, Tuple[float, float], float]:
+        """
+        Compute bootstrap statistics for a single effect.
+
+        Parameters
+        ----------
+        original_effect : float
+            Original point estimate.
+        boot_dist : np.ndarray
+            Bootstrap distribution of the effect.
+
+        Returns
+        -------
+        se : float
+            Bootstrap standard error.
+        ci : Tuple[float, float]
+            Percentile confidence interval.
+        p_value : float
+            Bootstrap p-value.
+        """
+        se = float(np.std(boot_dist, ddof=1))
+        ci = self._compute_percentile_ci(boot_dist, self.alpha)
+        p_value = self._compute_bootstrap_pvalue(original_effect, boot_dist)
+        return se, ci, p_value
+
     def get_params(self) -> Dict[str, Any]:
         """Get estimator parameters (sklearn-compatible)."""
         return {

{diff_diff-1.2.0 → diff_diff-1.2.1}/diff_diff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 1.2.0
+Version: 1.2.1
 Summary: A library for Difference-in-Differences causal inference analysis
 Author: diff-diff contributors
 License-Expression: MIT

{diff_diff-1.2.0 → diff_diff-1.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "diff-diff"
-version = "1.2.0"
+version = "1.2.1"
 description = "A library for Difference-in-Differences causal inference analysis"
 readme = "README.md"
 license = "MIT"

{diff_diff-1.2.0 → diff_diff-1.2.1}/tests/test_estimators.py

@@ -2090,6 +2090,159 @@ class TestSyntheticDiD:
         assert results1.att == results2.att
         assert results1.se == results2.se
 
+    def test_insufficient_pre_periods_warning(self):
+        """Test that SDID warns with very few pre-treatment periods."""
+        np.random.seed(42)
+
+        # Create data with only 2 pre-treatment periods
+        n_control = 8
+        n_periods = 4  # 2 pre, 2 post
+        post_periods = [2, 3]
+
+        data = []
+        # Treated unit
+        for t in range(n_periods):
+            y = 10.0 + t * 0.5 + np.random.normal(0, 0.3)
+            if t in post_periods:
+                y += 3.0
+            data.append({
+                "unit": 0,
+                "period": t,
+                "outcome": y,
+                "treated": 1,
+            })
+
+        # Control units
+        for unit in range(1, n_control + 1):
+            for t in range(n_periods):
+                y = 8.0 + unit * 0.2 + t * 0.4 + np.random.normal(0, 0.3)
+                data.append({
+                    "unit": unit,
+                    "period": t,
+                    "outcome": y,
+                    "treated": 0,
+                })
+
+        df = pd.DataFrame(data)
+
+        sdid = SyntheticDiD(n_bootstrap=30, seed=42)
+
+        # Should work but may warn about few pre-periods
+        # (Depending on implementation - some may warn, some may not)
+        results = sdid.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            unit="unit",
+            time="period",
+            post_periods=post_periods
+        )
+
+        # Results should still be valid
+        assert np.isfinite(results.att)
+        assert results.se > 0
+
+    def test_single_pre_period_edge_case(self):
+        """Test SDID with single pre-treatment period (extreme edge case)."""
+        np.random.seed(42)
+
+        n_control = 5
+        n_periods = 3  # 1 pre, 2 post
+        post_periods = [1, 2]
+
+        data = []
+        # Treated unit
+        for t in range(n_periods):
+            y = 10.0 + np.random.normal(0, 0.2)
+            if t in post_periods:
+                y += 2.0
+            data.append({
+                "unit": 0,
+                "period": t,
+                "outcome": y,
+                "treated": 1,
+            })
+
+        # Control units
+        for unit in range(1, n_control + 1):
+            for t in range(n_periods):
+                y = 9.0 + np.random.normal(0, 0.2)
+                data.append({
+                    "unit": unit,
+                    "period": t,
+                    "outcome": y,
+                    "treated": 0,
+                })
+
+        df = pd.DataFrame(data)
+
+        sdid = SyntheticDiD(n_bootstrap=30, seed=42)
+
+        # With single pre-period, time weights will be trivially [1.0]
+        results = sdid.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            unit="unit",
+            time="period",
+            post_periods=post_periods
+        )
+
+        # Should still produce results
+        assert np.isfinite(results.att)
+        # Time weights should have single entry
+        assert len(results.time_weights) == 1
+
+    def test_more_pre_periods_than_control_units(self):
+        """Test SDID when n_pre_periods > n_control_units (underdetermined)."""
+        np.random.seed(42)
+
+        n_control = 3  # Few control units
+        n_periods = 10  # Many periods
+        post_periods = [8, 9]  # 8 pre-treatment periods
+
+        data = []
+        # Treated unit
+        for t in range(n_periods):
+            y = 10.0 + t * 0.2 + np.random.normal(0, 0.3)
+            if t in post_periods:
+                y += 2.5
+            data.append({
+                "unit": 0,
+                "period": t,
+                "outcome": y,
+                "treated": 1,
+            })
+
+        # Control units
+        for unit in range(1, n_control + 1):
+            for t in range(n_periods):
+                y = 8.0 + t * 0.15 + np.random.normal(0, 0.3)
+                data.append({
+                    "unit": unit,
+                    "period": t,
+                    "outcome": y,
+                    "treated": 0,
+                })
+
+        df = pd.DataFrame(data)
+
+        # Use regularization to help with underdetermined system
+        sdid = SyntheticDiD(lambda_reg=1.0, n_bootstrap=30, seed=42)
+
+        results = sdid.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            unit="unit",
+            time="period",
+            post_periods=post_periods
+        )
+
+        # Should produce valid results with regularization
+        assert np.isfinite(results.att)
+        assert results.se > 0
+
 
 class TestSyntheticWeightsUtils:
     """Tests for synthetic weight utility functions."""
@@ -2179,3 +2332,368 @@ class TestSyntheticWeightsUtils:
         # Control: 12 - 10 = 2
         # SDID: 5 - 2 = 3
         assert abs(tau - 3.0) < 1e-6
+
+
+# =============================================================================
+# Edge Case Tests
+# =============================================================================
+
+
+class TestUnbalancedPanels:
+    """Tests for handling unbalanced panels with missing periods."""
+
+    def test_did_with_missing_periods(self):
+        """Test DifferenceInDifferences handles missing periods gracefully."""
+        # Create unbalanced panel - some units missing some periods
+        np.random.seed(42)
+        data = []
+
+        for unit in range(20):
+            is_treated = unit < 10
+            # Some units missing period 0, some missing period 1
+            periods = [0, 1]
+            if unit % 5 == 0:
+                periods = [1]  # Missing pre-period
+            elif unit % 7 == 0:
+                periods = [0]  # Missing post-period
+
+            for period in periods:
+                y = 10.0 + unit * 0.1
+                if period == 1:
+                    y += 5.0
+                if is_treated and period == 1:
+                    y += 3.0
+                y += np.random.normal(0, 1)
+
+                data.append({
+                    "unit": unit,
+                    "period": period,
+                    "treated": int(is_treated),
+                    "post": period,
+                    "outcome": y,
+                })
+
+        df = pd.DataFrame(data)
+
+        did = DifferenceInDifferences()
+        results = did.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            time="post"
+        )
+
+        # Should still produce valid results
+        assert np.isfinite(results.att)
+        assert results.se > 0
+        assert results.n_obs == len(df)
+
+    def test_twfe_with_unbalanced_panel(self):
+        """Test TwoWayFixedEffects handles unbalanced panels."""
+        from diff_diff import TwoWayFixedEffects
+
+        np.random.seed(42)
+        data = []
+
+        for unit in range(15):
+            is_treated = unit < 8
+            unit_effect = np.random.normal(0, 2)
+
+            # Create unbalanced panel - varying number of periods per unit
+            if unit < 5:
+                periods = [0, 1, 2, 3]  # Full panel
+            elif unit < 10:
+                periods = [0, 1, 3]  # Missing period 2
+            else:
+                periods = [1, 2, 3]  # Missing period 0
+
+            for period in periods:
+                time_effect = period * 0.5
+                post = 1 if period >= 2 else 0
+
+                y = 10.0 + unit_effect + time_effect
+                if is_treated and post:
+                    y += 3.0
+                y += np.random.normal(0, 0.5)
+
+                data.append({
+                    "unit": unit,
+                    "period": period,
+                    "treated": int(is_treated),
+                    "post": post,
+                    "outcome": y,
+                })
+
+        df = pd.DataFrame(data)
+
+        twfe = TwoWayFixedEffects()
+        results = twfe.fit(
+            df,
+            outcome="outcome",
+            treatment="post",
+            unit="unit",
+            time="period"
+        )
+
+        # Should produce valid results
+        assert np.isfinite(results.att)
+        assert results.se > 0
+
+    def test_multiperiod_with_sparse_data(self):
+        """Test MultiPeriodDiD with sparse data across periods."""
+        np.random.seed(42)
+        data = []
+
+        n_units = 30
+        for unit in range(n_units):
+            is_treated = unit < n_units // 2
+
+            # Each unit observed in random subset of periods
+            available_periods = np.random.choice([0, 1, 2, 3, 4], size=3, replace=False)
+            available_periods = sorted(available_periods)
+
+            for period in available_periods:
+                y = 10.0 + np.random.normal(0, 1)
+                if period >= 2:
+                    y += 2.0  # Time effect
+                if is_treated and period >= 2:
+                    y += 3.0  # Treatment effect
+
+                data.append({
+                    "unit": unit,
+                    "period": period,
+                    "treated": int(is_treated),
+                    "outcome": y,
+                })
+
+        df = pd.DataFrame(data)
+
+        mp_did = MultiPeriodDiD()
+        results = mp_did.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            time="period",
+            reference_period=1
+        )
+
+        # Should produce valid results
+        assert np.isfinite(results.avg_att)
+        assert len(results.period_effects) > 0
+
+
+class TestSingleTreatedUnit:
+    """Tests for scenarios with only one treated unit."""
+
+    def test_did_single_treated_unit(self):
+        """Test DifferenceInDifferences with single treated unit."""
+        np.random.seed(42)
+        data = []
+
+        # 1 treated unit, 10 control units
+        for unit in range(11):
+            is_treated = unit == 0
+
+            for period in [0, 1]:
+                y = 10.0 + np.random.normal(0, 0.5)
+                if period == 1:
+                    y += 2.0
+                if is_treated and period == 1:
+                    y += 5.0  # Large effect for single unit
+
+                data.append({
+                    "unit": unit,
+                    "period": period,
+                    "treated": int(is_treated),
+                    "post": period,
+                    "outcome": y,
+                })
+
+        df = pd.DataFrame(data)
+
+        did = DifferenceInDifferences()
+        results = did.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            time="post"
+        )
+
+        # Should produce valid results
+        assert np.isfinite(results.att)
+        assert results.se > 0
+        assert results.n_treated == 2  # 1 unit × 2 periods
+
+    def test_sdid_single_treated_unit(self):
+        """Test SyntheticDiD with single treated unit (primary use case)."""
+        np.random.seed(42)
+
+        n_control = 10
+        n_periods = 6
+        post_periods = [4, 5]
+
+        data = []
+
+        # Generate treated unit
+        treated_base = 15.0
+        treated_trend = 0.5
+        for t in range(n_periods):
+            y = treated_base + treated_trend * t + np.random.normal(0, 0.3)
+            if t in post_periods:
+                y += 3.0  # Treatment effect
+            data.append({
+                "unit": 0,
+                "period": t,
+                "outcome": y,
+                "treated": 1,
+            })
+
+        # Generate control units
+        for unit in range(1, n_control + 1):
+            unit_base = 10.0 + np.random.normal(0, 2)
+            unit_trend = 0.4 + np.random.normal(0, 0.1)
+            for t in range(n_periods):
+                y = unit_base + unit_trend * t + np.random.normal(0, 0.3)
+                data.append({
+                    "unit": unit,
+                    "period": t,
+                    "outcome": y,
+                    "treated": 0,
+                })
+
+        df = pd.DataFrame(data)
+
+        sdid = SyntheticDiD(n_bootstrap=50, seed=42)
+        results = sdid.fit(
+            df,
+            outcome="outcome",
+            treatment="treated",
+            unit="unit",
+            time="period",
+            post_periods=post_periods
+        )
+
+        # SDID is designed for single/few treated units
+        assert np.isfinite(results.att)
+        assert results.se > 0
+        # Effect should be roughly correct
+        assert abs(results.att - 3.0) < 2.0
+
+
+class TestCollinearityDetection:
+    """Tests for handling perfect or near collinearity."""
+
+    def test_did_with_redundant_covariate_raises_error(self):
+        """Test DiD raises clear error for perfectly collinear covariates."""
+        np.random.seed(42)
+        data = pd.DataFrame({
+            "outcome": np.random.normal(10, 1, 100),
+            "treated": np.repeat([0, 1], 50),
+            "post": np.tile([0, 1], 50),
+            "x1": np.random.normal(0, 1, 100),
+        })
+        # Add perfectly collinear covariate
+        data["x2"] = data["x1"] * 2 + 3
+
+        did = DifferenceInDifferences()
+
+        # Should raise a clear error about collinearity
+        with pytest.raises(ValueError, match="rank-deficient"):
+            did.fit(
+                data,
+                outcome="outcome",
+                treatment="treated",
+                time="post",
+                covariates=["x1", "x2"]
+            )
+
+    def test_did_with_constant_covariate_raises_error(self):
+        """Test DiD raises clear error for constant covariates."""
+        np.random.seed(42)
+        data = pd.DataFrame({
+            "outcome": np.random.normal(10, 1, 100),
+            "treated": np.repeat([0, 1], 50),
+            "post": np.tile([0, 1], 50),
+            "constant_x": np.ones(100),  # Constant covariate
+        })
+
+        did = DifferenceInDifferences()
+
+        # Constant covariate is collinear with intercept
+        # Should raise clear error
+        with pytest.raises(ValueError, match="rank-deficient"):
+            did.fit(
+                data,
+                outcome="outcome",
+                treatment="treated",
+                time="post",
+                covariates=["constant_x"]
+            )
+
+    def test_did_with_near_collinear_covariates(self):
+        """Test DiD handles near-collinear covariates (not perfectly collinear)."""
+        np.random.seed(42)
+        data = pd.DataFrame({
+            "outcome": np.random.normal(10, 1, 100),
+            "treated": np.repeat([0, 1], 50),
+            "post": np.tile([0, 1], 50),
+            "x1": np.random.normal(0, 1, 100),
+        })
+        # Add near-collinear covariate (small noise breaks perfect collinearity)
+        data["x2"] = data["x1"] * 2 + 3 + np.random.normal(0, 0.1, 100)
+
+        did = DifferenceInDifferences()
+
+        # Near-collinear should work (not perfectly rank-deficient)
+        results = did.fit(
+            data,
+            outcome="outcome",
+            treatment="treated",
+            time="post",
+            covariates=["x1", "x2"]
+        )
+
+        assert np.isfinite(results.att)
+
+    def test_twfe_with_absorbed_covariate(self):
+        """Test TWFE handles covariate absorbed by fixed effects."""
+        from diff_diff import TwoWayFixedEffects
+
+        np.random.seed(42)
+        n_units = 20
+        n_periods = 4
+
+        data = []
+        for unit in range(n_units):
+            # Unit-specific covariate (absorbed by unit FE)
+            unit_x = np.random.normal(0, 1)
+
+            for period in range(n_periods):
+                y = 10.0 + unit * 0.5 + period * 0.3 + np.random.normal(0, 0.5)
+                post = 1 if period >= 2 else 0
+                if unit < n_units // 2 and post:
+                    y += 2.0
+
+                data.append({
+                    "unit": unit,
+                    "period": period,
+                    "outcome": y,
+                    "treated": int(unit < n_units // 2),
+                    "post": post,
+                    "unit_covariate": unit_x,  # Same for all periods within unit
+                })
+
+        df = pd.DataFrame(data)
+
+        twfe = TwoWayFixedEffects()
+        # unit_covariate is absorbed by unit fixed effects
+        results = twfe.fit(
+            df,
+            outcome="outcome",
+            treatment="post",
+            unit="unit",
+            time="period"
+        )
+
+        assert np.isfinite(results.att)
+        assert results.se > 0

{diff_diff-1.2.0 → diff_diff-1.2.1}/tests/test_staggered.py

@@ -1125,3 +1125,234 @@ class TestCallawaySantAnnaBootstrap:
                 time='time',
                 first_treat='first_treat'
             )
+
+
+# =============================================================================
+# Edge Case Tests: Single Cohort
+# =============================================================================
+
+
+class TestCallawaySantAnnaSingleCohort:
+    """Tests for CallawaySantAnna with a single treatment cohort."""
+
+    def test_single_cohort_basic(self):
+        """Test CS estimator with single treatment cohort."""
+        np.random.seed(42)
+
+        n_units = 60
+        n_periods = 8
+        treatment_period = 4
+
+        # Generate data with single cohort
+        data = []
+        for unit in range(n_units):
+            # 40% never-treated, 60% treated at period 4
+            if unit < int(n_units * 0.4):
+                first_treat = 0  # Never treated
+            else:
+                first_treat = treatment_period  # Single cohort
+
+            unit_fe = np.random.normal(0, 2)
+
+            for t in range(n_periods):
+                time_fe = t * 0.3
+                y = 10.0 + unit_fe + time_fe
+
+                # Treatment effect for treated units after treatment
+                if first_treat > 0 and t >= first_treat:
+                    y += 2.5
+
+                y += np.random.normal(0, 0.5)
+
+                data.append({
+                    'unit': unit,
+                    'time': t,
+                    'outcome': y,
+                    'first_treat': first_treat,
+                })
+
+        df = pd.DataFrame(data)
+
+        cs = CallawaySantAnna()
+        results = cs.fit(
+            df,
+            outcome='outcome',
+            unit='unit',
+            time='time',
+            first_treat='first_treat'
+        )
+
+        # Should produce valid results
+        assert results.overall_att is not None
+        assert np.isfinite(results.overall_att)
+        assert results.overall_se > 0
+
+        # Should have effects for single group only
+        groups = set(g for g, t in results.group_time_effects.keys())
+        assert len(groups) == 1
+        assert treatment_period in groups
+
+        # ATT should be roughly correct
+        assert abs(results.overall_att - 2.5) < 1.5
+
+    def test_single_cohort_event_study(self):
+        """Test event study aggregation with single cohort."""
+        np.random.seed(42)
+
+        n_units = 80
+        n_periods = 12
+        treatment_period = 6  # Start later to have both pre and post periods
+
+        data = []
+        for unit in range(n_units):
+            if unit < int(n_units * 0.3):
+                first_treat = 0
+            else:
+                first_treat = treatment_period
+
+            unit_fe = np.random.normal(0, 1)
+
+            for t in range(n_periods):
+                y = 10.0 + unit_fe + t * 0.2
+
+                if first_treat > 0 and t >= first_treat:
+                    # Dynamic effect: grows over time
+                    periods_since = t - first_treat
+                    y += 2.0 + 0.3 * periods_since
+
+                y += np.random.normal(0, 0.4)
+
+                data.append({
+                    'unit': unit,
+                    'time': t,
+                    'outcome': y,
+                    'first_treat': first_treat,
+                })
+
+        df = pd.DataFrame(data)
+
+        cs = CallawaySantAnna()
+        results = cs.fit(
+            df,
+            outcome='outcome',
+            unit='unit',
+            time='time',
+            first_treat='first_treat',
+            aggregate='event_study'
+        )
+
+        assert results.event_study_effects is not None
+        assert len(results.event_study_effects) > 0
+
+        # Event study should have multiple relative periods
+        rel_periods = sorted(results.event_study_effects.keys())
+        assert len(rel_periods) >= 2, f"Expected multiple periods, got {rel_periods}"
+
+        # With single cohort, all effects are for the same group
+        # Post-treatment effects (e >= 0) should show positive effect
+        post_periods = [e for e in rel_periods if e >= 0]
+        if post_periods:
+            # At least some post-periods should show positive effect
+            post_effects = [results.event_study_effects[e]['effect'] for e in post_periods]
+            assert any(e > 0.5 for e in post_effects), f"Expected positive post-period effects, got {post_effects}"
+
+    def test_single_cohort_with_bootstrap(self):
+        """Test bootstrap inference with single cohort."""
+        np.random.seed(42)
+
+        n_units = 50
+        n_periods = 6
+        treatment_period = 3
+
+        data = []
+        for unit in range(n_units):
+            if unit < int(n_units * 0.4):
+                first_treat = 0
+            else:
+                first_treat = treatment_period
+
+            for t in range(n_periods):
+                y = 10.0 + np.random.normal(0, 1)
+                if first_treat > 0 and t >= first_treat:
+                    y += 3.0
+
+                data.append({
+                    'unit': unit,
+                    'time': t,
+                    'outcome': y,
+                    'first_treat': first_treat,
+                })
+
+        df = pd.DataFrame(data)
+
+        cs = CallawaySantAnna(n_bootstrap=99, seed=42)
+        results = cs.fit(
+            df,
+            outcome='outcome',
+            unit='unit',
+            time='time',
+            first_treat='first_treat'
+        )
+
+        assert results.bootstrap_results is not None
+        assert results.bootstrap_results.overall_att_se > 0
+        assert results.bootstrap_results.overall_att_ci[0] < results.bootstrap_results.overall_att_ci[1]
+
+    def test_single_cohort_not_yet_treated_control(self):
+        """Test single cohort with not_yet_treated control group.
+
+        With a single cohort, not_yet_treated should behave same as
+        never_treated after the treatment period.
+        """
+        np.random.seed(42)
+
+        n_units = 60
+        n_periods = 8
+        treatment_period = 4
+
+        data = []
+        for unit in range(n_units):
+            if unit < int(n_units * 0.4):
+                first_treat = 0
+            else:
+                first_treat = treatment_period
+
+            for t in range(n_periods):
+                y = 10.0 + np.random.normal(0, 0.5)
+                if first_treat > 0 and t >= first_treat:
+                    y += 2.0
+
+                data.append({
+                    'unit': unit,
+                    'time': t,
+                    'outcome': y,
+                    'first_treat': first_treat,
+                })
+
+        df = pd.DataFrame(data)
+
+        cs_never = CallawaySantAnna(control_group='never_treated')
+        results_never = cs_never.fit(
+            df,
+            outcome='outcome',
+            unit='unit',
+            time='time',
+            first_treat='first_treat'
+        )
+
+        cs_not_yet = CallawaySantAnna(control_group='not_yet_treated')
+        results_not_yet = cs_not_yet.fit(
+            df,
+            outcome='outcome',
+            unit='unit',
+            time='time',
+            first_treat='first_treat'
+        )
+
+        # Both should produce valid results
+        assert np.isfinite(results_never.overall_att)
+        assert np.isfinite(results_not_yet.overall_att)
+
+        # Results may differ slightly due to different comparison groups
+        # but should be in similar range
+        assert abs(results_never.overall_att - results_not_yet.overall_att) < 1.0

{diff_diff-1.2.0 → diff_diff-1.2.1}/tests/test_wild_bootstrap.py

@@ -621,3 +621,184 @@ class TestWildBootstrapResults:
 
         captured = capsys.readouterr()
         assert "Wild Cluster Bootstrap Results" in captured.out
+
+
+# =============================================================================
+# Edge Case Tests: Few Clusters (< 5)
+# =============================================================================
+
+
+class TestFewClustersEdgeCases:
+    """Tests for wild bootstrap behavior with very few clusters."""
+
+    def test_three_clusters_still_works(self):
+        """Test wild bootstrap works with 3 clusters (minimum viable)."""
+        np.random.seed(42)
+
+        n_clusters = 3
+        obs_per_cluster = 40
+
+        data = []
+        for cluster in range(n_clusters):
+            is_treated = cluster < 2  # 2 treated, 1 control cluster
+            cluster_effect = np.random.normal(0, 2)
+
+            for obs in range(obs_per_cluster):
+                for period in [0, 1]:
+                    y = 10.0 + cluster_effect
+                    if period == 1:
+                        y += 5.0
+                    if is_treated and period == 1:
+                        y += 3.0
+                    y += np.random.normal(0, 1)
+
+                    data.append({
+                        "cluster": cluster,
+                        "unit": cluster * obs_per_cluster + obs,
+                        "period": period,
+                        "treated": int(is_treated),
+                        "post": period,
+                        "outcome": y,
+                    })
+
+        df = pd.DataFrame(data)
+
+        did = DifferenceInDifferences(
+            cluster="cluster",
+            inference="wild_bootstrap",
+            n_bootstrap=99,
+            bootstrap_weights="webb",  # Webb recommended for few clusters
+            seed=42
+        )
+
+        # Should warn about few clusters but still produce valid results
+        with pytest.warns(UserWarning, match="Only 3 clusters"):
+            results = did.fit(
+                df,
+                outcome="outcome",
+                treatment="treated",
+                time="post"
+            )
+
+        assert results.se > 0
+        assert results.inference_method == "wild_bootstrap"
+        assert results.n_clusters == 3
+
+    def test_two_clusters_minimum(self):
+        """Test wild bootstrap works with exactly 2 clusters (absolute minimum)."""
+        np.random.seed(42)
+
+        n_clusters = 2
+        obs_per_cluster = 50
+
+        data = []
+        for cluster in range(n_clusters):
+            is_treated = cluster == 0
+            cluster_effect = np.random.normal(0, 2)
+
+            for obs in range(obs_per_cluster):
+                for period in [0, 1]:
+                    y = 10.0 + cluster_effect
+                    if period == 1:
+                        y += 5.0
+                    if is_treated and period == 1:
+                        y += 3.0
+                    y += np.random.normal(0, 1)
+
+                    data.append({
+                        "cluster": cluster,
+                        "unit": cluster * obs_per_cluster + obs,
+                        "period": period,
+                        "treated": int(is_treated),
+                        "post": period,
+                        "outcome": y,
+                    })
+
+        df = pd.DataFrame(data)
+
+        did = DifferenceInDifferences(
+            cluster="cluster",
+            inference="wild_bootstrap",
+            n_bootstrap=99,
+            bootstrap_weights="webb",
+            seed=42
+        )
+
+        # Should warn about few clusters
+        with pytest.warns(UserWarning, match="Only 2 clusters"):
+            results = did.fit(
+                df,
+                outcome="outcome",
+                treatment="treated",
+                time="post"
+            )
+
+        # Results should still be valid (though may have high variance)
+        assert results.se > 0
+        assert np.isfinite(results.att)
+        assert results.n_clusters == 2
+
+    def test_few_clusters_webb_vs_rademacher(self, few_cluster_data):
+        """Test that Webb weights produce different (often more conservative) SEs than Rademacher with few clusters."""
+        did_webb = DifferenceInDifferences(
+            cluster="cluster",
+            inference="wild_bootstrap",
+            n_bootstrap=199,
+            bootstrap_weights="webb",
+            seed=42
+        )
+
+        did_rademacher = DifferenceInDifferences(
+            cluster="cluster",
+            inference="wild_bootstrap",
+            n_bootstrap=199,
+            bootstrap_weights="rademacher",
+            seed=42
+        )
+
+        with pytest.warns(UserWarning):
+            results_webb = did_webb.fit(
+                few_cluster_data,
+                outcome="outcome",
+                treatment="treated",
+                time="post"
+            )
+
+        with pytest.warns(UserWarning):
+            results_rademacher = did_rademacher.fit(
+                few_cluster_data,
+                outcome="outcome",
+                treatment="treated",
+                time="post"
+            )
+
+        # Both should produce valid results
+        assert results_webb.se > 0
+        assert results_rademacher.se > 0
+        # ATT should be identical (same point estimate)
+        assert results_webb.att == results_rademacher.att
+        # SEs will differ due to different weight distributions
+        # (This is expected, not necessarily one > other)
+
+    def test_few_clusters_confidence_intervals_valid(self, few_cluster_data):
+        """Test that CIs are valid even with few clusters."""
+        did = DifferenceInDifferences(
+            cluster="cluster",
+            inference="wild_bootstrap",
+            n_bootstrap=199,
+            bootstrap_weights="webb",
+            seed=42
+        )
+
+        with pytest.warns(UserWarning):
+            results = did.fit(
+                few_cluster_data,
+                outcome="outcome",
+                treatment="treated",
+                time="post"
+            )
+
+        lower, upper = results.conf_int
+        assert lower < upper
+        # CI should contain the point estimate
+        assert lower < results.att < upper