Fix MultiPeriodDiD bootstrap+survey fallback and relabel n_treated/n_control as observations from PR #218 review (round 14)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/estimators.py

@@ -838,13 +838,16 @@ def fit(  # type: ignore[override]
         ValueError
             If required parameters are missing or data validation fails.
         """
-        # Warn if wild bootstrap is requested but not supported
+        # Fall back to analytical inference if wild bootstrap requested
+        # (must happen before _resolve_survey_for_fit which rejects bootstrap+survey)
+        effective_inference = self.inference
         if self.inference == "wild_bootstrap":
             warnings.warn(
                 "Wild bootstrap inference is not yet supported for MultiPeriodDiD. "
                 "Using analytical inference instead.",
                 UserWarning,
             )
+            effective_inference = "analytical"
 
         # Validate basic inputs
         if outcome is None or treatment is None or time is None:
@@ -992,7 +995,7 @@ def fit(  # type: ignore[override]
         from diff_diff.survey import _resolve_effective_cluster, _resolve_survey_for_fit
 
         resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
-            _resolve_survey_for_fit(survey_design, data, self.inference)
+            _resolve_survey_for_fit(survey_design, data, effective_inference)
         )
 
         # Handle absorbed fixed effects (within-transformation)

diff_diff/results.py

@@ -34,9 +34,9 @@ class DiDResults:
     n_obs : int
         Number of observations used in estimation.
     n_treated : int
-        Number of treated units.
+        Number of treated observations.
     n_control : int
-        Number of control units.
+        Number of control observations.
     """
 
     att: float
@@ -93,8 +93,8 @@ def summary(self, alpha: Optional[float] = None) -> str:
             "=" * 70,
             "",
             f"{'Observations:':<25} {self.n_obs:>10}",
-            f"{'Treated units:':<25} {self.n_treated:>10}",
-            f"{'Control units:':<25} {self.n_control:>10}",
+            f"{'Treated obs:':<25} {self.n_treated:>10}",
+            f"{'Control obs:':<25} {self.n_control:>10}",
         ]
 
         if self.r_squared is not None:
@@ -645,9 +645,9 @@ class SyntheticDiDResults:
     n_obs : int
         Number of observations used in estimation.
     n_treated : int
-        Number of treated units.
+        Number of treated observations.
     n_control : int
-        Number of control units.
+        Number of control observations.
     unit_weights : dict
         Dictionary mapping control unit IDs to their synthetic weights.
     time_weights : dict
@@ -714,8 +714,8 @@ def summary(self, alpha: Optional[float] = None) -> str:
             "=" * 75,
             "",
             f"{'Observations:':<25} {self.n_obs:>10}",
-            f"{'Treated units:':<25} {self.n_treated:>10}",
-            f"{'Control units:':<25} {self.n_control:>10}",
+            f"{'Treated obs:':<25} {self.n_treated:>10}",
+            f"{'Control obs:':<25} {self.n_control:>10}",
             f"{'Pre-treatment periods:':<25} {len(self.pre_periods):>10}",
             f"{'Post-treatment periods:':<25} {len(self.post_periods):>10}",
         ]

tests/test_survey.py

@@ -2793,3 +2793,34 @@ def test_zero_score_dispersion_stratified_psu(self):
         vcov = compute_survey_vcov(X, residuals, resolved=resolved)
         # Zero residuals → zero scores → zero V_h per stratum → zero vcov
         np.testing.assert_array_equal(vcov, np.zeros((2, 2)))
+
+
+class TestRound14Fixes:
+    """Tests for PR #218 review round 14 fixes."""
+
+    def test_multiperiod_bootstrap_survey_fallback(self):
+        """MultiPeriodDiD with wild_bootstrap + survey_design falls back gracefully."""
+        np.random.seed(42)
+        n = 40
+        df = pd.DataFrame(
+            {
+                "outcome": np.random.randn(n),
+                "treated": np.array([1] * 20 + [0] * 20),
+                "time": np.tile([0, 1, 2, 3], 10),
+                "w": np.ones(n),
+            }
+        )
+        sd = SurveyDesign(weights="w", weight_type="pweight")
+        mpd = MultiPeriodDiD(inference="wild_bootstrap")
+        # Should warn about fallback and produce valid analytical results
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            result = mpd.fit(
+                df,
+                outcome="outcome",
+                treatment="treated",
+                time="time",
+                post_periods=[2, 3],
+                survey_design=sd,
+            )
+        assert np.isfinite(result.avg_att)