#76 stop evaluating for multiple units per time point

pymc-labs · Dec 26, 2022 · e3bc8cd · e3bc8cd
1 parent 0ed7240
commit e3bc8cd
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 79 deletions.
diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py
@@ -298,9 +298,13 @@ def __init__(
         self.x_pred_control = (
             self.data
             # just the untreated group
-            .query(f"{self.group_variable_name} == @self.untreated")  # 🔥
+            .query(f"{self.group_variable_name} == @self.untreated")
             # 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
+            .groupby(self.time_variable_name)
+            .first()
+            .reset_index()
         )
         assert not self.x_pred_control.empty
         (new_x,) = build_design_matrices([self._x_design_info], self.x_pred_control)
@@ -310,9 +314,13 @@ 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} == @self.treated")
             # 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
+            .groupby(self.time_variable_name)
+            .first()
+            .reset_index()
         )
         assert not self.x_pred_treatment.empty
         (new_x,) = build_design_matrices([self._x_design_info], self.x_pred_treatment)
@@ -322,14 +330,17 @@ 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} == @self.treated")
             # just the treatment period(s)
-            # TODO: the line below might need some work to be more robust
             .query("post_treatment == True")
             # drop the outcome variable
             .drop(self.outcome_variable_name, axis=1)
             # DO AN INTERVENTION. Set the post_treatment variable to False
             .assign(post_treatment=False)
+            # We may have multiple units per time point, we only want one time point
+            .groupby(self.time_variable_name)
+            .first()
+            .reset_index()
         )
         assert not self.x_pred_counterfactual.empty
         (new_x,) = build_design_matrices(

diff --git a/docs/notebooks/did_pymc.ipynb b/docs/notebooks/did_pymc.ipynb
diff --git a/docs/notebooks/did_pymc_banks.ipynb b/docs/notebooks/did_pymc_banks.ipynb