Merge pull request #226 from pymc-labs/rd-bandwidth

Regression Discontinuity: add ability to specify bandwidth

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -1,4 +1,5 @@
-from typing import Union
+import warnings
+from typing import Optional, Union
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -541,7 +542,7 @@ def summary(self):
 
 class RegressionDiscontinuity(ExperimentalDesign):
     """
-    A class to analyse regression discontinuity experiments.
+    A class to analyse sharp regression discontinuity experiments.
 
     :param data:
         A pandas dataframe
@@ -556,6 +557,9 @@ class RegressionDiscontinuity(ExperimentalDesign):
     :param epsilon:
         A small scalar value which determines how far above and below the treatment
         threshold to evaluate the causal impact.
+    :param bandwidth:
+        Data outside of the bandwidth (relative to the discontinuity) is not used to fit
+        the model.
     """
 
     def __init__(
@@ -566,6 +570,7 @@ def __init__(
         model=None,
         running_variable_name: str = "x",
         epsilon: float = 0.001,
+        bandwidth: Optional[float] = None,
         **kwargs,
     ):
         super().__init__(model=model, **kwargs)
@@ -575,9 +580,22 @@ def __init__(
         self.running_variable_name = running_variable_name
         self.treatment_threshold = treatment_threshold
         self.epsilon = epsilon
+        self.bandwidth = bandwidth
         self._input_validation()
 
-        y, X = dmatrices(formula, self.data)
+        if self.bandwidth is not None:
+            fmin = self.treatment_threshold - self.bandwidth
+            fmax = self.treatment_threshold + self.bandwidth
+            filtered_data = self.data.query(f"{fmin} <= x <= {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,
+                )
+            y, X = dmatrices(formula, filtered_data)
+        else:
+            y, X = dmatrices(formula, self.data)
+
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
         self.labels = X.design_info.column_names
@@ -594,11 +612,14 @@ def __init__(
         self.score = self.model.score(X=self.X, y=self.y)
 
         # get the model predictions of the observed data
-        xi = np.linspace(
-            np.min(self.data[self.running_variable_name]),
-            np.max(self.data[self.running_variable_name]),
-            200,
-        )
+        if self.bandwidth is not None:
+            xi = np.linspace(fmin, fmax, 200)
+        else:
+            xi = np.linspace(
+                np.min(self.data[self.running_variable_name]),
+                np.max(self.data[self.running_variable_name]),
+                200,
+            )
         self.x_pred = pd.DataFrame(
             {self.running_variable_name: xi, "treated": self._is_treated(xi)}
         )

causalpy/skl_experiments.py

@@ -1,3 +1,6 @@
+import warnings
+from typing import Optional
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -346,7 +349,7 @@ def plot(self):
 
 class RegressionDiscontinuity(ExperimentalDesign):
     """
-    A class to analyse regression discontinuity experiments.
+    A class to analyse sharp regression discontinuity experiments.
 
     :param data:
         A pandas dataframe
@@ -361,6 +364,9 @@ class RegressionDiscontinuity(ExperimentalDesign):
     :param epsilon:
         A small scalar value which determines how far above and below the treatment
         threshold to evaluate the causal impact.
+    :param bandwidth:
+        Data outside of the bandwidth (relative to the discontinuity) is not used to fit
+        the model.
     """
 
     def __init__(
@@ -371,15 +377,30 @@ def __init__(
         model=None,
         running_variable_name="x",
         epsilon: float = 0.001,
+        bandwidth: Optional[float] = None,
         **kwargs,
     ):
         super().__init__(model=model, **kwargs)
         self.data = data
         self.formula = formula
         self.running_variable_name = running_variable_name
         self.treatment_threshold = treatment_threshold
+        self.bandwidth = bandwidth
         self.epsilon = epsilon
-        y, X = dmatrices(formula, self.data)
+
+        if self.bandwidth is not None:
+            fmin = self.treatment_threshold - self.bandwidth
+            fmax = self.treatment_threshold + self.bandwidth
+            filtered_data = self.data.query(f"{fmin} <= x <= {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,
+                )
+            y, X = dmatrices(formula, filtered_data)
+        else:
+            y, X = dmatrices(formula, self.data)
+
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
         self.labels = X.design_info.column_names
@@ -396,11 +417,14 @@ def __init__(
         self.score = self.model.score(X=self.X, y=self.y)
 
         # get the model predictions of the observed data
-        xi = np.linspace(
-            np.min(self.data[self.running_variable_name]),
-            np.max(self.data[self.running_variable_name]),
-            1000,
-        )
+        if self.bandwidth is not None:
+            xi = np.linspace(fmin, fmax, 200)
+        else:
+            xi = np.linspace(
+                np.min(self.data[self.running_variable_name]),
+                np.max(self.data[self.running_variable_name]),
+                200,
+            )
         self.x_pred = pd.DataFrame(
             {self.running_variable_name: xi, "treated": self._is_treated(xi)}
         )

causalpy/tests/test_integration_pymc_examples.py

@@ -120,6 +120,23 @@ def test_rd():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
+@pytest.mark.integration
+def test_rd_bandwidth():
+    df = cp.load_data("rd")
+    result = cp.pymc_experiments.RegressionDiscontinuity(
+        df,
+        formula="y ~ 1 + x + treated + x:treated",
+        model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        treatment_threshold=0.5,
+        epsilon=0.001,
+        bandwidth=0.3,
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.RegressionDiscontinuity)
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
+
+
 @pytest.mark.integration
 def test_rd_drinking():
     df = (

causalpy/tests/test_integration_skl_examples.py

@@ -88,6 +88,21 @@ def test_rd_linear_main_effects():
     assert isinstance(result, cp.skl_experiments.RegressionDiscontinuity)
 
 
+@pytest.mark.integration
+def test_rd_linear_main_effects_bandwidth():
+    data = cp.load_data("rd")
+    result = cp.skl_experiments.RegressionDiscontinuity(
+        data,
+        formula="y ~ 1 + x + treated",
+        model=LinearRegression(),
+        treatment_threshold=0.5,
+        epsilon=0.001,
+        bandwidth=0.3,
+    )
+    assert isinstance(data, pd.DataFrame)
+    assert isinstance(result, cp.skl_experiments.RegressionDiscontinuity)
+
+
 @pytest.mark.integration
 def test_rd_linear_with_interaction():
     data = cp.load_data("rd")