Add the epsilon kwarg into the sci-kit learn version of the RegressionDiscontinuity class

Author: drbenvincent

causalpy/skl_experiments.py

@@ -346,13 +346,21 @@ 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 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.
     """
 
     def __init__(
@@ -362,13 +370,15 @@ def __init__(
         treatment_threshold,
         model=None,
         running_variable_name="x",
+        epsilon: float = 0.001,
         **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.epsilon = epsilon
         y, X = dmatrices(formula, self.data)
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
@@ -404,7 +414,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]),
             }