add to all the other notebooks + fix failing tests

Author: drbenvincent

causalpy/experiments/base.py

@@ -20,19 +20,27 @@
 
 import arviz as az
 import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
 from sklearn.base import RegressorMixin
 
 from causalpy.pymc_models import PyMCModel
 from causalpy.reporting import (
     EffectSummary,
     _compute_statistics,
+    _compute_statistics_did_ols,
+    _compute_statistics_ols,
     _detect_experiment_type,
     _effect_summary_did,
+    _effect_summary_rd,
     _extract_counterfactual,
     _extract_window,
     _generate_prose,
+    _generate_prose_did_ols,
+    _generate_prose_ols,
     _generate_table,
+    _generate_table_did_ols,
+    _generate_table_ols,
 )
 from causalpy.skl_models import create_causalpy_compatible_class
 
@@ -129,34 +137,34 @@ def effect_summary(
         treated_unit: Optional[str] = None,
     ) -> EffectSummary:
         """
-        Generate a decision-ready summary of causal effects from posterior draws.
+        Generate a decision-ready summary of causal effects.
 
-        Supports Interrupted Time Series (ITS), Synthetic Control, and
-        Difference-in-Differences (DiD) experiments. Automatically detects experiment
-        type and generates appropriate summary.
+        Supports Interrupted Time Series (ITS), Synthetic Control, Difference-in-Differences (DiD),
+        and Regression Discontinuity (RD) experiments. Works with both PyMC (Bayesian) and OLS models.
+        Automatically detects experiment type and model type, generating appropriate summary.
 
         Parameters
         ----------
         window : str, tuple, or slice, default="post"
-            Time window for analysis (ITS/SC only, ignored for DiD):
+            Time window for analysis (ITS/SC only, ignored for DiD/RD):
             - "post": All post-treatment time points (default)
             - (start, end): Tuple of start and end times (handles both datetime and integer indices)
             - slice: Python slice object for integer indices
         direction : {"increase", "decrease", "two-sided"}, default="increase"
-            Direction for tail probability calculation:
+            Direction for tail probability calculation (PyMC only, ignored for OLS):
             - "increase": P(effect > 0)
             - "decrease": P(effect < 0)
             - "two-sided": Two-sided p-value, report 1-p as "probability of effect"
         alpha : float, default=0.05
-            Significance level for HDI intervals (1-alpha confidence level)
+            Significance level for HDI/CI intervals (1-alpha confidence level)
         cumulative : bool, default=True
-            Whether to include cumulative effect statistics (ITS/SC only, ignored for DiD)
+            Whether to include cumulative effect statistics (ITS/SC only, ignored for DiD/RD)
         relative : bool, default=True
             Whether to include relative effect statistics (% change vs counterfactual)
-            (ITS/SC only, ignored for DiD)
+            (ITS/SC only, ignored for DiD/RD)
         min_effect : float, optional
-            Region of Practical Equivalence (ROPE) threshold. If provided, reports
-            P(|effect| > min_effect) for two-sided or P(effect > min_effect) for one-sided.
+            Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).
+            If provided, reports P(|effect| > min_effect) for two-sided or P(effect > min_effect) for one-sided.
         treated_unit : str, optional
             For multi-unit experiments (Synthetic Control), specify which treated unit
             to analyze. If None and multiple units exist, uses first unit.
@@ -168,38 +176,62 @@ def effect_summary(
 
         Examples
         --------
-        Interrupted Time Series:
+        Interrupted Time Series (PyMC):
 
             import causalpy as cp
-            result = cp.InterruptedTimeSeries(...)
+            result = cp.InterruptedTimeSeries(..., model=cp.pymc_models.LinearRegression(...))
             stats = result.effect_summary()
             print(stats.table)
             print(stats.text)
 
+        Interrupted Time Series (OLS):
+
+            from sklearn.linear_model import LinearRegression
+            result = cp.InterruptedTimeSeries(..., model=LinearRegression())
+            stats = result.effect_summary()
+            print(stats.table)
+
         Difference-in-Differences:
 
             result = cp.DifferenceInDifferences(...)
             stats = result.effect_summary()
             print(stats.table)
-        """
-        # Validate model type
-        if not isinstance(self.model, PyMCModel):
-            raise ValueError(
-                "effect_summary currently only supports PyMC models. "
-                "OLS model support is planned for future release."
-            )
 
+        Regression Discontinuity:
+
+            result = cp.RegressionDiscontinuity(...)
+            stats = result.effect_summary()
+            print(stats.table)
+        """
         # Detect experiment type
         experiment_type = _detect_experiment_type(self)
 
-        if experiment_type == "did":
-            # Difference-in-Differences: scalar effect, no time dimension
-            return _effect_summary_did(
+        # Check if PyMC or OLS model
+        is_pymc = isinstance(self.model, PyMCModel)
+
+        if experiment_type == "rd":
+            # Regression Discontinuity: scalar effect, no time dimension
+            return _effect_summary_rd(
                 self,
                 direction=direction,
                 alpha=alpha,
                 min_effect=min_effect,
             )
+        elif experiment_type == "did":
+            # Difference-in-Differences: scalar effect, no time dimension
+            if is_pymc:
+                return _effect_summary_did(
+                    self,
+                    direction=direction,
+                    alpha=alpha,
+                    min_effect=min_effect,
+                )
+            else:
+                # OLS DiD
+                stats = _compute_statistics_did_ols(self, alpha=alpha)
+                table = _generate_table_did_ols(stats)
+                text = _generate_prose_did_ols(stats, alpha=alpha)
+                return EffectSummary(table=table, text=text)
         else:
             # ITS or Synthetic Control: time-series effects
             # Extract windowed impact data
@@ -212,29 +244,63 @@ def effect_summary(
                 self, window_coords, treated_unit=treated_unit
             )
 
-            # Compute statistics
-            hdi_prob = 1 - alpha
-            stats = _compute_statistics(
-                windowed_impact,
-                counterfactual,
-                hdi_prob=hdi_prob,
-                direction=direction,
-                cumulative=cumulative,
-                relative=relative,
-                min_effect=min_effect,
-            )
-
-            # Generate table
-            table = _generate_table(stats, cumulative=cumulative, relative=relative)
-
-            # Generate prose
-            text = _generate_prose(
-                stats,
-                window_coords,
-                alpha=alpha,
-                direction=direction,
-                cumulative=cumulative,
-                relative=relative,
-            )
+            if is_pymc:
+                # PyMC model: use posterior draws
+                hdi_prob = 1 - alpha
+                stats = _compute_statistics(
+                    windowed_impact,
+                    counterfactual,
+                    hdi_prob=hdi_prob,
+                    direction=direction,
+                    cumulative=cumulative,
+                    relative=relative,
+                    min_effect=min_effect,
+                )
+
+                # Generate table
+                table = _generate_table(stats, cumulative=cumulative, relative=relative)
+
+                # Generate prose
+                text = _generate_prose(
+                    stats,
+                    window_coords,
+                    alpha=alpha,
+                    direction=direction,
+                    cumulative=cumulative,
+                    relative=relative,
+                )
+            else:
+                # OLS model: use point estimates and CIs
+                # Convert to numpy arrays if needed
+                if hasattr(windowed_impact, "values"):
+                    impact_array = windowed_impact.values
+                else:
+                    impact_array = np.asarray(windowed_impact)
+                if hasattr(counterfactual, "values"):
+                    counterfactual_array = counterfactual.values
+                else:
+                    counterfactual_array = np.asarray(counterfactual)
+
+                stats = _compute_statistics_ols(
+                    impact_array,
+                    counterfactual_array,
+                    alpha=alpha,
+                    cumulative=cumulative,
+                    relative=relative,
+                )
+
+                # Generate table
+                table = _generate_table_ols(
+                    stats, cumulative=cumulative, relative=relative
+                )
+
+                # Generate prose
+                text = _generate_prose_ols(
+                    stats,
+                    window_coords,
+                    alpha=alpha,
+                    cumulative=cumulative,
+                    relative=relative,
+                )
 
             return EffectSummary(table=table, text=text)