Merge branch 'main' into feature_instrumental_variables

Author: NathanielF

causalpy/pymc_experiments.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Union
+from typing import Optional, Union
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -543,20 +543,24 @@ def summary(self):
 
 class RegressionDiscontinuity(ExperimentalDesign):
     """
-    A class to analyse regression discontinuity experiments.
-
-    :param data: A pandas dataframe
-    :param formula: A statistical model formula
-    :param treatment_threshold: A scalar threshold value at which the treatment
-                                is applied
-    :param model: A PyMC model
-    :param running_variable_name: The name of the predictor variable that the treatment
-                                  threshold is based upon
-
-    .. note::
-
-        There is no pre/post intervention data distinction for the regression
-        discontinuity design, we fit all the data available.
+    A class to analyse sharp regression discontinuity experiments.
+
+    :param data:
+        A pandas dataframe
+    :param formula:
+        A statistical model formula
+    :param treatment_threshold:
+        A scalar threshold value at which the treatment is applied
+    :param model:
+        A PyMC model
+    :param running_variable_name:
+        The name of the predictor variable that the treatment threshold is based upon
+    :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,8 @@ def __init__(
         treatment_threshold: float,
         model=None,
         running_variable_name: str = "x",
+        epsilon: float = 0.001,
+        bandwidth: Optional[float] = None,
         **kwargs,
     ):
         super().__init__(model=model, **kwargs)
@@ -574,9 +580,23 @@ def __init__(
         self.formula = formula
         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
@@ -593,11 +613,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)}
         )
@@ -611,7 +634,10 @@ def __init__(
         self.x_discon = pd.DataFrame(
             {
                 self.running_variable_name: np.array(
-                    [self.treatment_threshold - 0.001, self.treatment_threshold + 0.001]
+                    [
+                        self.treatment_threshold - self.epsilon,
+                        self.treatment_threshold + self.epsilon,
+                    ]
                 ),
                 "treated": np.array([0, 1]),
             }

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,13 +349,24 @@ def plot(self):
 
 class RegressionDiscontinuity(ExperimentalDesign):
     """
-    Analyse data from regression discontinuity experiments.
-
-    .. note::
-
-        There is no pre/post intervention data distinction for the regression
-        discontinuity design, we fit all the data available.
-
+    A class to analyse sharp regression discontinuity experiments.
+
+    :param data:
+        A pandas dataframe
+    :param formula:
+        A statistical model formula
+    :param treatment_threshold:
+        A scalar threshold value at which the treatment is applied
+    :param model:
+        A sci-kit learn model object
+    :param running_variable_name:
+        The name of the predictor variable that the treatment threshold is based upon
+    :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__(
@@ -362,14 +376,31 @@ def __init__(
         treatment_threshold,
         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
-        y, X = dmatrices(formula, self.data)
+        self.bandwidth = bandwidth
+        self.epsilon = epsilon
+
+        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
@@ -386,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)}
         )
@@ -404,7 +438,10 @@ def __init__(
         self.x_discon = pd.DataFrame(
             {
                 self.running_variable_name: np.array(
-                    [self.treatment_threshold - 0.001, self.treatment_threshold + 0.001]
+                    [
+                        self.treatment_threshold - self.epsilon,
+                        self.treatment_threshold + self.epsilon,
+                    ]
                 ),
                 "treated": np.array([0, 1]),
             }

causalpy/tests/test_integration_pymc_examples.py

@@ -112,6 +112,24 @@ def test_rd():
         formula="y ~ 1 + bs(x, df=6) + treated",
         model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         treatment_threshold=0.5,
+        epsilon=0.001,
+    )
+    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_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)

causalpy/tests/test_integration_skl_examples.py

@@ -36,6 +36,7 @@ def test_rd_drinking():
         running_variable_name="age",
         model=LinearRegression(),
         treatment_threshold=21,
+        epsilon=0.001,
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.skl_experiments.RegressionDiscontinuity)
@@ -81,6 +82,22 @@ def test_rd_linear_main_effects():
         formula="y ~ 1 + x + treated",
         model=LinearRegression(),
         treatment_threshold=0.5,
+        epsilon=0.001,
+    )
+    assert isinstance(data, pd.DataFrame)
+    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)
@@ -94,6 +111,7 @@ def test_rd_linear_with_interaction():
         formula="y ~ 1 + x + treated + x:treated",
         model=LinearRegression(),
         treatment_threshold=0.5,
+        epsilon=0.001,
     )
     assert isinstance(data, pd.DataFrame)
     assert isinstance(result, cp.skl_experiments.RegressionDiscontinuity)
@@ -108,6 +126,7 @@ def test_rd_linear_with_gaussian_process():
         formula="y ~ 1 + x + treated",
         model=GaussianProcessRegressor(kernel=kernel),
         treatment_threshold=0.5,
+        epsilon=0.001,
     )
     assert isinstance(data, pd.DataFrame)
     assert isinstance(result, cp.skl_experiments.RegressionDiscontinuity)