Merge pull request #151 from pymc-labs/remove-unecessary-kwargs

remove `treated` and `untreated` kwargs for DifferenceInDifferences

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -254,8 +254,6 @@ def __init__(
         formula: str,
         time_variable_name: str,
         group_variable_name: str,
-        treated: str,
-        untreated: str,
         model=None,
         **kwargs,
     ):
@@ -265,10 +263,6 @@ def __init__(
         self.formula = formula
         self.time_variable_name = time_variable_name
         self.group_variable_name = group_variable_name
-        self.treated = treated  # level of the group_variable_name that was treated
-        self.untreated = (
-            untreated  # level of the group_variable_name that was untreated
-        )
         y, X = dmatrices(formula, self.data)
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
@@ -277,11 +271,9 @@ def __init__(
         self.outcome_variable_name = y.design_info.column_names[0]
 
         # Input validation ----------------------------------------------------
-        # Check that `treated` appears in the module formula
         assert (
             "post_treatment" in formula
         ), "A predictor called `post_treatment` should be in the dataframe"
-        # Check that we have `treated` in the incoming dataframe
         assert (
             "post_treatment" in self.data.columns
         ), "Require a boolean column labelling observations which are `treated`"
@@ -293,26 +285,22 @@ def __init__(
         Require a `unit` column to label unique units.
         This is used for plotting purposes
         """
-        # Check that `group_variable_name` has TWO levels, representing the
-        # treated/untreated. But it does not matter what the actual names of
-        # the levels are.
-        assert (
-            len(pd.Categorical(self.data[self.group_variable_name]).categories) == 2
+        # Check that `group_variable_name` is dummy coded. It should be 0 or 1
+        assert not set(self.data[self.group_variable_name]).difference(
+            set([0, 1])
         ), f"""
-            There must be 2 levels of the grouping variable {self.group_variable_name}
-            .I.e. the treated and untreated.
+            The grouping variable {self.group_variable_name} should be dummy coded.
+            Consisting of 0's and 1's only.
         """
 
-        # DEVIATION FROM SKL EXPERIMENT CODE =============================
         COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.X.shape[0])}
         self.model.fit(X=self.X, y=self.y, coords=COORDS)
-        # ================================================================
 
         # predicted outcome for control group
         self.x_pred_control = (
             self.data
             # just the untreated group
-            .query(f"{self.group_variable_name} == @self.untreated")
+            .query(f"{self.group_variable_name} == 0")
             # drop the outcome variable
             .drop(self.outcome_variable_name, axis=1)
             # We may have multiple units per time point, we only want one time point
@@ -328,7 +316,7 @@ def __init__(
         self.x_pred_treatment = (
             self.data
             # just the treated group
-            .query(f"{self.group_variable_name} == @self.treated")
+            .query(f"{self.group_variable_name} == 1")
             # drop the outcome variable
             .drop(self.outcome_variable_name, axis=1)
             # We may have multiple units per time point, we only want one time point
@@ -345,7 +333,7 @@ def __init__(
         self.x_pred_counterfactual = (
             self.data
             # just the treated group
-            .query(f"{self.group_variable_name} == @self.treated")
+            .query(f"{self.group_variable_name} == 1")
             # just the treatment period(s)
             .query("post_treatment == True")
             # drop the outcome variable

causalpy/tests/test_integration_pymc_examples.py

@@ -14,8 +14,6 @@ def test_did():
         formula="y ~ 1 + group*post_treatment",
         time_variable_name="t",
         group_variable_name="group",
-        treated=1,
-        untreated=0,
         model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     assert isinstance(df, pd.DataFrame)
@@ -59,8 +57,6 @@ def test_did_banks_simple():
         formula="bib ~ 1 + district * post_treatment",
         time_variable_name="year",
         group_variable_name="district",
-        treated=1,
-        untreated=0,
         model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     assert isinstance(df, pd.DataFrame)
@@ -100,8 +96,6 @@ def test_did_banks_multi():
         formula="bib ~ 1 + year + district + post_treatment + district:post_treatment",
         time_variable_name="year",
         group_variable_name="district",
-        treated=1,
-        untreated=0,
         model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     assert isinstance(df, pd.DataFrame)

docs/notebooks/did_pymc.ipynb

@@ -239,8 +239,6 @@
     "    formula=\"y ~ 1 + group*post_treatment\",\n",
     "    time_variable_name=\"t\",\n",
     "    group_variable_name=\"group\",\n",
-    "    treated=1,\n",
-    "    untreated=0,\n",
     "    model=cp.pymc_models.LinearRegression(sample_kwargs={\"random_seed\": seed}),\n",
     ")"
    ]

docs/notebooks/did_pymc_banks.ipynb

@@ -482,8 +482,6 @@
     "    formula=\"bib ~ 1 + district * post_treatment\",\n",
     "    time_variable_name=\"year\",\n",
     "    group_variable_name=\"district\",\n",
-    "    treated=1,\n",
-    "    untreated=0,\n",
     "    model=cp.pymc_models.LinearRegression(\n",
     "        sample_kwargs={\"target_accept\": 0.95, \"random_seed\": seed}\n",
     "    ),\n",
@@ -647,8 +645,6 @@
     "    formula=\"bib ~ 1 + year + district + post_treatment + district:post_treatment\",\n",
     "    time_variable_name=\"year\",\n",
     "    group_variable_name=\"district\",\n",
-    "    treated=1,\n",
-    "    untreated=0,\n",
     "    model=cp.pymc_models.LinearRegression(sample_kwargs={\"random_seed\": seed}),\n",
     ")"
    ]