diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index f1857b62..5809def0 100644 --- a/causalpy/skl_experiments.py +++ b/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,6 +370,7 @@ def __init__( treatment_threshold, model=None, running_variable_name="x", + epsilon: float = 0.001, **kwargs, ): super().__init__(model=model, **kwargs) @@ -369,6 +378,7 @@ def __init__( 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]), }