Add Donut Regression Discontinuity functionality (#610)

* Add donut RDD support to RegressionDiscontinuity

Introduces a donut_hole parameter to RegressionDiscontinuity, allowing exclusion of observations within a specified distance from the treatment threshold for robustness against manipulation or heaping. Updates plotting, summary, and input validation to support donut RDD, adds comprehensive tests for donut_hole behavior and validation, expands glossary with donut RDD concepts, and provides a new notebook demonstrating donut RDD usage. References on donut RDD, manipulation, and heaping are added to the bibliography.

* update cell tags

* Increase sample size and strengthen manipulation in donut RDD notebook

Raised the number of generated observations from 500 to 1000 and increased manipulation parameters to better demonstrate the donut RDD approach. Updated output and data table examples to reflect the new data generation settings.

* edit cell tag

* add glossary terms

* Expand explanation of heaping and manipulation in Donut RDD

* Add regression discontinuity edge case tests

Added tests for warnings when bandwidth or donut_hole filters leave few datapoints, for unrecognized model types, and for donut hole boundary lines in OLS and Bayesian plots. Also updated interrogate badge coverage from 96.3% to 96.4%.

* Fix plot, legend, warning, and validation issues in donut RDD

- Conditional two-layer scatter: only show excluded/fit data distinction
  when data is actually excluded; default case shows single "data" layer
- Fix Bayesian plot legend to include all labeled artists (donut
  boundaries, threshold, scatter labels) instead of only posterior mean
- Fix malformed warning when filter_desc is empty on small datasets
- Use ValueError instead of DataException for donut_hole param validation
- Update tests to expect ValueError for donut_hole validation

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: drbenvincent

causalpy/experiments/regression_discontinuity.py

@@ -60,6 +60,12 @@ class RegressionDiscontinuity(BaseExperiment):
     :param bandwidth:
         Data outside of the bandwidth (relative to the discontinuity) is not used to fit
         the model.
+    :param donut_hole:
+        Observations within this distance from the treatment threshold are excluded from
+        model fitting. Used as a robustness check when observations closest to the
+        threshold may be problematic (e.g., due to manipulation or heaping). Defaults
+        to 0.0 (no exclusion). Must be non-negative and less than bandwidth if bandwidth
+        is finite.
 
     Example
     --------
@@ -94,6 +100,7 @@ def __init__(
         running_variable_name: str = "x",
         epsilon: float = 0.001,
         bandwidth: float = np.inf,
+        donut_hole: float = 0.0,
         **kwargs: dict,
     ) -> None:
         super().__init__(model=model)
@@ -104,28 +111,43 @@ def __init__(
         self.treatment_threshold = treatment_threshold
         self.epsilon = epsilon
         self.bandwidth = bandwidth
+        self.donut_hole = donut_hole
         self.input_validation()
         self._build_design_matrices()
         self._prepare_data()
         self.algorithm()
 
     def _build_design_matrices(self) -> None:
-        """Build design matrices from formula and data, applying bandwidth filtering."""
+        """Build design matrices from formula and data, applying bandwidth and donut hole filtering."""
+        x_vals = self.data[self.running_variable_name]
+        c = self.treatment_threshold
+        mask = pd.Series(True, index=self.data.index)
+
         if self.bandwidth is not np.inf:
-            fmin = self.treatment_threshold - self.bandwidth
-            fmax = self.treatment_threshold + self.bandwidth
-            filtered_data = self.data.query(
-                f"{fmin} <= {self.running_variable_name} <= {fmax}"
-            )
-            if len(filtered_data) <= 10:
-                warnings.warn(
-                    f"Choice of bandwidth parameter has lead to only {len(filtered_data)} remaining datapoints. Consider increasing the bandwidth parameter.",  # noqa: E501
-                    UserWarning,
-                    stacklevel=2,
+            mask &= np.abs(x_vals - c) <= self.bandwidth
+
+        if self.donut_hole > 0:
+            mask &= np.abs(x_vals - c) >= self.donut_hole
+
+        self.fit_data = self.data.loc[mask]
+
+        if len(self.fit_data) <= 10:
+            filter_desc = []
+            if self.bandwidth is not np.inf:
+                filter_desc.append(f"bandwidth={self.bandwidth}")
+            if self.donut_hole > 0:
+                filter_desc.append(f"donut_hole={self.donut_hole}")
+            if filter_desc:
+                msg = (
+                    f"Choice of {' and '.join(filter_desc)} parameters has led to only "
+                    f"{len(self.fit_data)} remaining datapoints. "
+                    f"Consider adjusting these parameters."
                 )
-            y, X = dmatrices(self.formula, filtered_data)
-        else:
-            y, X = dmatrices(self.formula, self.data)
+            else:
+                msg = f"Only {len(self.fit_data)} datapoints in the dataset."
+            warnings.warn(msg, UserWarning, stacklevel=2)
+
+        y, X = dmatrices(self.formula, self.fit_data)
 
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
@@ -227,6 +249,16 @@ def input_validation(self) -> None:
                 """The treated variable should be dummy coded. Consisting of 0's and 1's only."""  # noqa: E501
             )
 
+        # Validate donut_hole parameter
+        if self.donut_hole < 0:
+            raise ValueError("donut_hole must be non-negative.")
+
+        if self.bandwidth is not np.inf and self.donut_hole >= self.bandwidth:
+            raise ValueError(
+                f"donut_hole ({self.donut_hole}) must be less than bandwidth "
+                f"({self.bandwidth}) when bandwidth is finite."
+            )
+
         # Convert integer treated variable to boolean if needed
         if self.data["treated"].dtype in ["int64", "int32"]:
             # Make a copy to avoid SettingWithCopyWarning
@@ -249,10 +281,13 @@ def summary(self, round_to: int | None = None) -> None:
         :param round_to:
             Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers.
         """
-        print("Difference in Differences experiment")
+        print("Regression Discontinuity experiment")
         print(f"Formula: {self.formula}")
         print(f"Running variable: {self.running_variable_name}")
         print(f"Threshold on running variable: {self.treatment_threshold}")
+        print(f"Bandwidth: {self.bandwidth}")
+        print(f"Donut hole: {self.donut_hole}")
+        print(f"Observations used for fit: {len(self.fit_data)}")
         print("\nResults:")
         print(
             f"Discontinuity at threshold = {convert_to_string(self.discontinuity_at_threshold)}"
@@ -265,28 +300,39 @@ def _bayesian_plot(
     ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for regression discontinuity designs."""
         fig, ax = plt.subplots()
-        # Plot raw data
+
+        # Plot data: use two layers only when there are excluded observations
+        has_exclusion = len(self.fit_data) < len(self.data)
+        if has_exclusion:
+            sns.scatterplot(
+                self.data,
+                x=self.running_variable_name,
+                y=self.outcome_variable_name,
+                color="lightgray",
+                ax=ax,
+                label="excluded data",
+            )
         sns.scatterplot(
-            self.data,
+            self.fit_data,
             x=self.running_variable_name,
             y=self.outcome_variable_name,
-            c="k",
+            color="k",
             ax=ax,
+            label="fit data" if has_exclusion else "data",
         )
 
         # Plot model fit to data
-        h_line, h_patch = plot_xY(
+        plot_xY(
             self.x_pred[self.running_variable_name],
             self.pred["posterior_predictive"].mu.isel(treated_units=0),
             ax=ax,
             plot_hdi_kwargs={"color": "C1"},
+            label="Posterior mean",
         )
-        handles = [(h_line, h_patch)]
-        labels = ["Posterior mean"]
 
         # create strings to compose title
         title_info = f"{round_num(self.score['unit_0_r2'], round_to)} (std = {round_num(self.score['unit_0_r2_std'], round_to)})"
-        r2 = f"Bayesian $R^2$ on all data = {title_info}"
+        r2 = f"Bayesian $R^2$ on fit data = {title_info}"
         percentiles = self.discontinuity_at_threshold.quantile([0.03, 1 - 0.03]).values
         ci = (
             r"$CI_{94\%}$"
@@ -296,34 +342,61 @@ def _bayesian_plot(
             Discontinuity at threshold = {round_num(self.discontinuity_at_threshold.mean(), round_to)},
             """
         ax.set(title=r2 + "\n" + discon + ci)
-        # Intervention line
+
+        # Treatment threshold line
         ax.axvline(
             x=self.treatment_threshold,
             ls="-",
             lw=3,
             color="r",
             label="treatment threshold",
         )
-        ax.legend(
-            handles=(h_tuple for h_tuple in handles),
-            labels=labels,
-            fontsize=LEGEND_FONT_SIZE,
-        )
+
+        # Add donut hole boundary lines if donut_hole > 0
+        if self.donut_hole > 0:
+            ax.axvline(
+                x=self.treatment_threshold - self.donut_hole,
+                ls="--",
+                lw=2,
+                color="orange",
+                label="donut boundary",
+            )
+            ax.axvline(
+                x=self.treatment_threshold + self.donut_hole,
+                ls="--",
+                lw=2,
+                color="orange",
+            )
+
+        ax.legend(fontsize=LEGEND_FONT_SIZE)
         return (fig, ax)
 
     def _ols_plot(
         self, round_to: int | None = None, **kwargs: dict
     ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for regression discontinuity designs."""
         fig, ax = plt.subplots()
-        # Plot raw data
+
+        # Plot data: use two layers only when there are excluded observations
+        has_exclusion = len(self.fit_data) < len(self.data)
+        if has_exclusion:
+            sns.scatterplot(
+                self.data,
+                x=self.running_variable_name,
+                y=self.outcome_variable_name,
+                color="lightgray",
+                ax=ax,
+                label="excluded data",
+            )
         sns.scatterplot(
-            self.data,
+            self.fit_data,
             x=self.running_variable_name,
             y=self.outcome_variable_name,
-            c="k",  # hue="treated",
+            color="k",
             ax=ax,
+            label="fit data" if has_exclusion else "data",
         )
+
         # Plot model fit to data
         ax.plot(
             self.x_pred[self.running_variable_name],
@@ -332,18 +405,37 @@ def _ols_plot(
             markersize=10,
             label="model fit",
         )
+
         # create strings to compose title
-        r2 = f"$R^2$ on all data = {round_num(float(self.score), round_to)}"
+        r2 = f"$R^2$ on fit data = {round_num(float(self.score), round_to)}"
         discon = f"Discontinuity at threshold = {round_num(self.discontinuity_at_threshold, round_to)}"
         ax.set(title=r2 + "\n" + discon)
-        # Intervention line
+
+        # Treatment threshold line
         ax.axvline(
             x=self.treatment_threshold,
             ls="-",
             lw=3,
             color="r",
             label="treatment threshold",
         )
+
+        # Add donut hole boundary lines if donut_hole > 0
+        if self.donut_hole > 0:
+            ax.axvline(
+                x=self.treatment_threshold - self.donut_hole,
+                ls="--",
+                lw=2,
+                color="orange",
+                label="donut boundary",
+            )
+            ax.axvline(
+                x=self.treatment_threshold + self.donut_hole,
+                ls="--",
+                lw=2,
+                color="orange",
+            )
+
         ax.legend(fontsize=LEGEND_FONT_SIZE)
         return (fig, ax)