Address PR #208 review: add TWFE to registry, fix unregistered-estimator fallback

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

Author: igerber

diff_diff/power.py

@@ -136,6 +136,15 @@ def _basic_fit_kwargs(
     return dict(outcome="outcome", treatment="treated", time="post")
 
 
+def _twfe_fit_kwargs(
+    data: pd.DataFrame,
+    n_units: int,
+    n_periods: int,
+    treatment_period: int,
+) -> Dict[str, Any]:
+    return dict(outcome="outcome", treatment="treated", time="post", unit="unit")
+
+
 def _multiperiod_fit_kwargs(
     data: pd.DataFrame,
     n_units: int,
@@ -255,6 +264,13 @@ def _get_registry() -> Dict[str, _EstimatorProfile]:
             result_extractor=_extract_simple,
             min_n=20,
         ),
+        "TwoWayFixedEffects": _EstimatorProfile(
+            default_dgp=generate_did_data,
+            dgp_kwargs_builder=_basic_dgp_kwargs,
+            fit_kwargs_builder=_twfe_fit_kwargs,
+            result_extractor=_extract_simple,
+            min_n=20,
+        ),
         "MultiPeriodDiD": _EstimatorProfile(
             default_dgp=generate_did_data,
             dgp_kwargs_builder=_basic_dgp_kwargs,
@@ -1303,7 +1319,9 @@ def simulate_power(
     if profile is None and data_generator is None:
         raise ValueError(
             f"Estimator '{estimator_name}' not in registry. "
-            f"Provide a custom data_generator and estimator_kwargs."
+            f"Provide a custom data_generator and estimator_kwargs "
+            f"(the full dict of keyword arguments for estimator.fit(), "
+            f"e.g. dict(outcome='y', treatment='treat', time='period'))."
         )
 
     # When a custom data_generator is provided, bypass registry DGP
@@ -1414,8 +1432,7 @@ def simulate_power(
                         )
                         fit_kwargs.update(est_kwargs)
                     else:
-                        fit_kwargs = dict(outcome="outcome", treatment="treated", time="post")
-                        fit_kwargs.update(est_kwargs)
+                        fit_kwargs = dict(est_kwargs)
 
                 result = estimator.fit(data, **fit_kwargs)
 

tests/test_power.py

@@ -44,6 +44,7 @@
     _staggered_fit_kwargs,
     _trop_fit_kwargs,
 )
+from diff_diff.prep import generate_did_data
 
 
 class TestPowerAnalysis:
@@ -800,17 +801,10 @@ def test_continuous_did_not_in_registry(self):
                 progress=False,
             )
 
-    def test_twfe_not_in_registry(self):
-        """TwoWayFixedEffects is not in registry and raises without custom data_generator."""
+    def test_twfe_in_registry(self):
+        """TwoWayFixedEffects is in the registry."""
         registry = _get_registry()
-        assert "TwoWayFixedEffects" not in registry
-
-        with pytest.raises(ValueError, match="not in registry"):
-            simulate_power(
-                TwoWayFixedEffects(),
-                n_simulations=5,
-                progress=False,
-            )
+        assert "TwoWayFixedEffects" in registry
 
     def test_unknown_estimator_raises_without_data_generator(self):
         """Unknown estimator without data_generator raises ValueError."""
@@ -1066,6 +1060,83 @@ def test_sdid_sample_size(self):
         assert isinstance(result, SimulationSampleSizeResults)
         assert result.required_n > 0
 
+    @pytest.mark.slow
+    def test_twfe(self):
+        result = simulate_power(
+            TwoWayFixedEffects(),
+            n_simulations=5,
+            seed=42,
+            progress=False,
+        )
+        self._assert_valid_result(result, "TwoWayFixedEffects")
+
+    @pytest.mark.slow
+    def test_twfe_mde(self):
+        result = simulate_mde(
+            TwoWayFixedEffects(),
+            n_simulations=5,
+            effect_range=(0.5, 5.0),
+            seed=42,
+            progress=False,
+        )
+        assert isinstance(result, SimulationMDEResults)
+        assert result.mde > 0
+
+    @pytest.mark.slow
+    def test_twfe_sample_size(self):
+        result = simulate_sample_size(
+            TwoWayFixedEffects(),
+            n_simulations=5,
+            n_range=(20, 100),
+            seed=42,
+            progress=False,
+        )
+        assert isinstance(result, SimulationSampleSizeResults)
+        assert result.required_n > 0
+
+    @pytest.mark.slow
+    def test_custom_fallback_unregistered_estimator(self):
+        """Unregistered estimator works with custom data_generator and estimator_kwargs."""
+
+        class _UnregisteredEstimator:
+            """Unregistered wrapper for testing custom fallback."""
+
+            def __init__(self):
+                self._inner = DifferenceInDifferences()
+
+            def fit(self, data, **kwargs):
+                return self._inner.fit(data, **kwargs)
+
+        result = simulate_power(
+            _UnregisteredEstimator(),
+            data_generator=generate_did_data,
+            estimator_kwargs=dict(outcome="outcome", treatment="treated", time="post"),
+            n_simulations=5,
+            seed=42,
+            progress=False,
+        )
+        assert 0 <= result.power <= 1
+        assert result.n_simulations > 0
+
+    def test_custom_fallback_missing_kwargs_raises(self):
+        """Unregistered estimator with no estimator_kwargs fails on fit."""
+
+        class _UnregisteredEstimator:
+            def __init__(self):
+                self._inner = DifferenceInDifferences()
+
+            def fit(self, data, **kwargs):
+                return self._inner.fit(data, **kwargs)
+
+        with pytest.raises((ValueError, TypeError, RuntimeError)):
+            simulate_power(
+                _UnregisteredEstimator(),
+                data_generator=generate_did_data,
+                n_simulations=5,
+                seed=42,
+                progress=False,
+            )
+
 
 # ---------------------------------------------------------------------------
 # simulate_mde tests