Merge pull request #190 from pymc-labs/small-did-notebook-fix

fix small error in DiD notebook
pymc-labs · May 26, 2023 · 39714f2 · 39714f2
2 parents ed61617 + d934d9e
commit 39714f2
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 112 deletions.
diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py
@@ -134,7 +134,6 @@ def _input_validation(self, data, treatment_time):
             )
 
     def plot(self):
-
         """Plot the results"""
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
@@ -488,7 +487,13 @@ def _plot_causal_impact_arrow(self, ax):
             self.y_pred_counterfactual["posterior_predictive"].mu.mean().data
         )
         # Calculate the x position to plot at
-        diff = np.ptp(self.x_pred_treatment[self.time_variable_name].values)
+        # Note that we force to be float to avoid a type error using np.ptp with boolean
+        # values
+        diff = np.ptp(
+            np.array(self.x_pred_treatment[self.time_variable_name].values).astype(
+                float
+            )
+        )
         x = np.max(self.x_pred_treatment[self.time_variable_name].values) + 0.1 * diff
         # Plot the arrow
         ax.annotate(

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