fix: address CI review round 2 - block heterogeneous TE, fix hi-bracket, docstrings

- P0: reject heterogeneous_te_by_strata=True in survey power validation
  (DGP population_att diverges from input treatment_effect, making
  bias/coverage/RMSE metrics misleading)
- P1: fix simulate_sample_size auto-bracketing hi to respect abs_min
  (hi = max(2*lo, abs_min, 100) prevents probing below survey floor)
- P2: add survey_config parameter to simulate_power/mde/sample_size
  docstrings with constraints and supported estimators
- P2: add regression tests for heterogeneous_te rejection and
  large-floor (n_strata=10, psu_per_stratum=10, min_viable_n=200)
  auto-bracketing branch

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

Author: igerber

diff_diff/power.py

@@ -1812,6 +1812,13 @@ def simulate_power(
         estimators with non-standard result schemas.
     progress : bool, default=True
         Whether to print progress updates.
+    survey_config : SurveyPowerConfig, optional
+        When provided, generates survey-structured data via
+        ``generate_survey_did_data`` and injects ``SurveyDesign`` into
+        estimator ``fit()``. Mutually exclusive with ``data_generator``.
+        Supported estimators: DiD, TWFE, MultiPeriod, CS, SA, Imputation,
+        TwoStage, Stacked, Efficient. Unsupported: TROP, SyntheticDiD,
+        TripleDifference. ``heterogeneous_te_by_strata`` must be False.
 
     Returns
     -------
@@ -1906,6 +1913,13 @@ def simulate_power(
                 f"No survey power profile for {estimator_name}. "
                 f"Supported: {sorted(_SURVEY_FIT_BUILDERS.keys())}."
             )
+        if survey_config.heterogeneous_te_by_strata:
+            raise ValueError(
+                "heterogeneous_te_by_strata=True is not supported with "
+                "simulation power analysis. The DGP's population ATT diverges "
+                "from the input treatment_effect under heterogeneous effects, "
+                "which would make bias/coverage/RMSE metrics misleading."
+            )
 
     data_gen_kwargs = data_generator_kwargs or {}
     est_kwargs = estimator_kwargs or {}
@@ -2482,6 +2496,9 @@ def simulate_mde(
         Forwarded to ``simulate_power()``.
     progress : bool, default=True
         Whether to print progress updates.
+    survey_config : SurveyPowerConfig, optional
+        Survey-aware simulation config. Forwarded to ``simulate_power()``.
+        See :func:`simulate_power` for details and constraints.
 
     Returns
     -------
@@ -2693,6 +2710,11 @@ def simulate_sample_size(
         Forwarded to ``simulate_power()``.
     progress : bool, default=True
         Whether to print progress updates.
+    survey_config : SurveyPowerConfig, optional
+        Survey-aware simulation config. Forwarded to ``simulate_power()``.
+        When set, the bisection floor is raised to
+        ``survey_config.min_viable_n`` to ensure viable survey structure.
+        See :func:`simulate_power` for details and constraints.
 
     Returns
     -------
@@ -2847,7 +2869,7 @@ def _power_at_n(n: int) -> float:
                 )
             # Fall through to bisection with lo..hi bracket
         else:
-            hi = max(100, 2 * min_n)
+            hi = max(2 * lo, abs_min, 100)
             for _ in range(10):
                 if _power_at_n(hi) >= power:
                     break

tests/test_power.py

@@ -2377,6 +2377,36 @@ def test_survey_sample_size_large_effect_floor(self):
         )
         assert result.required_n >= cfg.min_viable_n  # must be >= 80
 
+    def test_survey_sample_size_large_floor_auto_bracket(self):
+        """Auto-bracketing hi respects min_viable_n > 100."""
+        cfg = SurveyPowerConfig(n_strata=10, psu_per_stratum=10)
+        # min_viable_n = 10 * 10 * 2 = 200, which exceeds default hi=100
+        result = simulate_sample_size(
+            CallawaySantAnna(),
+            treatment_effect=0.5,
+            sigma=3.0,
+            n_simulations=10,
+            max_steps=3,
+            seed=42,
+            survey_config=cfg,
+            n_periods=4,
+            treatment_period=2,
+            progress=False,
+        )
+        assert result.required_n >= cfg.min_viable_n  # must be >= 200
+
+    def test_survey_rejects_heterogeneous_te(self):
+        """heterogeneous_te_by_strata=True rejected with simulation power."""
+        cfg = SurveyPowerConfig(heterogeneous_te_by_strata=True)
+        with pytest.raises(ValueError, match="heterogeneous_te_by_strata"):
+            simulate_power(
+                CallawaySantAnna(),
+                n_simulations=1,
+                seed=42,
+                survey_config=cfg,
+                **_SIM_KW,
+            )
+
     # -- Closed-form deff tests --
 
     def test_closed_form_deff_default(self):