Add result_extractor param, fix stale power.rst docs, and add missing API symbols

- Add `result_extractor` parameter to simulate_power, simulate_mde, and
  simulate_sample_size for unregistered estimators with non-standard schemas
- Fix power.rst: correct PowerAnalysis method names, example code, and add
  SimulationMDEResults/SimulationSampleSizeResults/simulate_mde/simulate_sample_size
- Add 4 missing symbols to docs/api/index.rst autosummary
- Add api/power.rst to doc snippet smoke tests
- Add tests for custom result_extractor and MDE forwarding

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

Author: igerber

diff_diff/power.py

@@ -1214,6 +1214,7 @@ def simulate_power(
     data_generator: Optional[Callable] = None,
     data_generator_kwargs: Optional[Dict[str, Any]] = None,
     estimator_kwargs: Optional[Dict[str, Any]] = None,
+    result_extractor: Optional[Callable] = None,
     progress: bool = True,
 ) -> SimulationPowerResults:
     """
@@ -1257,6 +1258,11 @@ def simulate_power(
         Additional keyword arguments for data generator.
     estimator_kwargs : dict, optional
         Additional keyword arguments for estimator.fit().
+    result_extractor : callable, optional
+        Custom function to extract results from the estimator output.
+        Takes the estimator result object and returns a tuple of
+        ``(att, se, p_value, conf_int)``. Useful for unregistered
+        estimators with non-standard result schemas.
     progress : bool, default=True
         Whether to print progress updates.
 
@@ -1439,6 +1445,8 @@ def simulate_power(
                 # --- Extract results ---
                 if profile is not None:
                     att, se, p_val, ci = profile.result_extractor(result)
+                elif result_extractor is not None:
+                    att, se, p_val, ci = result_extractor(result)
                 else:
                     att = result.att if hasattr(result, "att") else result.avg_att
                     se = result.se if hasattr(result, "se") else result.avg_se
@@ -1717,6 +1725,7 @@ def simulate_mde(
     data_generator: Optional[Callable] = None,
     data_generator_kwargs: Optional[Dict[str, Any]] = None,
     estimator_kwargs: Optional[Dict[str, Any]] = None,
+    result_extractor: Optional[Callable] = None,
     progress: bool = True,
 ) -> SimulationMDEResults:
     """
@@ -1759,6 +1768,9 @@ def simulate_mde(
         Additional keyword arguments for data generator.
     estimator_kwargs : dict, optional
         Additional keyword arguments for estimator.fit().
+    result_extractor : callable, optional
+        Custom function to extract results from the estimator output.
+        Forwarded to ``simulate_power()``.
     progress : bool, default=True
         Whether to print progress updates.
 
@@ -1789,6 +1801,7 @@ def simulate_mde(
         data_generator=data_generator,
         data_generator_kwargs=data_generator_kwargs,
         estimator_kwargs=estimator_kwargs,
+        result_extractor=result_extractor,
         progress=False,
     )
 
@@ -1911,6 +1924,7 @@ def simulate_sample_size(
     data_generator: Optional[Callable] = None,
     data_generator_kwargs: Optional[Dict[str, Any]] = None,
     estimator_kwargs: Optional[Dict[str, Any]] = None,
+    result_extractor: Optional[Callable] = None,
     progress: bool = True,
 ) -> SimulationSampleSizeResults:
     """
@@ -1951,6 +1965,9 @@ def simulate_sample_size(
         Additional keyword arguments for data generator.
     estimator_kwargs : dict, optional
         Additional keyword arguments for estimator.fit().
+    result_extractor : callable, optional
+        Custom function to extract results from the estimator output.
+        Forwarded to ``simulate_power()``.
     progress : bool, default=True
         Whether to print progress updates.
 
@@ -1988,6 +2005,7 @@ def simulate_sample_size(
         data_generator=data_generator,
         data_generator_kwargs=data_generator_kwargs,
         estimator_kwargs=estimator_kwargs,
+        result_extractor=result_extractor,
         progress=False,
     )
 

docs/api/index.rst

@@ -148,10 +148,14 @@ Power analysis for study design:
    diff_diff.PowerAnalysis
    diff_diff.PowerResults
    diff_diff.SimulationPowerResults
+   diff_diff.SimulationMDEResults
+   diff_diff.SimulationSampleSizeResults
    diff_diff.compute_power
    diff_diff.compute_mde
    diff_diff.compute_sample_size
    diff_diff.simulate_power
+   diff_diff.simulate_mde
+   diff_diff.simulate_sample_size
 
 Pre-Trends Power Analysis
 -------------------------

docs/api/power.rst

@@ -30,9 +30,11 @@ Main class for analytical power calculations.
 
    .. autosummary::
 
-      ~PowerAnalysis.compute_power
-      ~PowerAnalysis.compute_mde
-      ~PowerAnalysis.compute_sample_size
+      ~PowerAnalysis.power
+      ~PowerAnalysis.mde
+      ~PowerAnalysis.sample_size
+      ~PowerAnalysis.power_curve
+      ~PowerAnalysis.sample_size_curve
 
 Example
 ~~~~~~~
@@ -41,29 +43,19 @@ Example
 
    from diff_diff import PowerAnalysis
 
-   # Create power analysis object
-   pa = PowerAnalysis(
-       effect_size=0.5,
-       n_treated=100,
-       n_control=100,
-       n_pre=4,
-       n_post=4,
-       sigma=1.0,
-       rho=0.5,  # Within-unit correlation
-       alpha=0.05
-   )
+   pa = PowerAnalysis(alpha=0.05, power=0.80)
 
    # Compute power
-   power = pa.compute_power()
-   print(f"Power: {power:.2%}")
+   result = pa.power(effect_size=0.5, n_treated=100, n_control=100, sigma=1.0)
+   print(f"Power: {result.power:.2%}")
 
    # Compute MDE at 80% power
-   mde = pa.compute_mde(power=0.80)
-   print(f"MDE: {mde:.3f}")
+   result = pa.mde(n_treated=100, n_control=100, sigma=1.0)
+   print(f"MDE: {result.mde:.3f}")
 
    # Required sample size
-   n = pa.compute_sample_size(power=0.80)
-   print(f"Required N per group: {n}")
+   result = pa.sample_size(effect_size=0.5, sigma=1.0)
+   print(f"Required N: {result.required_n}")
 
 PowerResults
 ------------
@@ -85,6 +77,26 @@ Results from simulation-based power analysis.
    :undoc-members:
    :show-inheritance:
 
+SimulationMDEResults
+--------------------
+
+Results from simulation-based MDE search.
+
+.. autoclass:: diff_diff.SimulationMDEResults
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+SimulationSampleSizeResults
+---------------------------
+
+Results from simulation-based sample size search.
+
+.. autoclass:: diff_diff.SimulationSampleSizeResults
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Convenience Functions
 ---------------------
 
@@ -116,6 +128,20 @@ Simulation-based power for any DiD estimator.
 
 .. autofunction:: diff_diff.simulate_power
 
+simulate_mde
+~~~~~~~~~~~~~
+
+Simulation-based MDE for any DiD estimator.
+
+.. autofunction:: diff_diff.simulate_mde
+
+simulate_sample_size
+~~~~~~~~~~~~~~~~~~~~
+
+Simulation-based sample size for any DiD estimator.
+
+.. autofunction:: diff_diff.simulate_sample_size
+
 Complete Example
 ----------------
 
@@ -125,8 +151,8 @@ Complete Example
        PowerAnalysis,
        compute_mde,
        simulate_power,
+       simulate_mde,
        DifferenceInDifferences,
-       plot_power_curve,
    )
 
    # Quick MDE calculation
@@ -145,20 +171,23 @@ Complete Example
    # Simulation-based power for DiD estimator
    sim_results = simulate_power(
        estimator=DifferenceInDifferences(),
-       effect_size=0.5,
-       n_treated=100,
-       n_control=100,
-       n_periods=8,
-       treatment_start=4,
+       treatment_effect=5.0,
+       n_units=100,
+       n_periods=4,
+       treatment_period=2,
        sigma=1.0,
-       n_simulations=1000
+       n_simulations=20,
    )
    print(f"Simulated power: {sim_results.power:.2%}")
 
-   # Power curve
-   pa = PowerAnalysis(n_treated=100, n_control=100, n_pre=4, n_post=4, sigma=1.0)
-   ax = plot_power_curve(pa, effect_range=(0, 1), n_points=50)
-   ax.figure.savefig('power_curve.png')
+   # Simulation-based MDE
+   mde_results = simulate_mde(
+       estimator=DifferenceInDifferences(),
+       n_units=100,
+       n_simulations=10,
+       max_steps=5,
+   )
+   print(f"Simulated MDE: {mde_results.mde:.3f}")
 
 See Also
 --------