Merge pull request #56 from pymc-labs/bayesian-r2

#43 use Bayesian R2 measures for pymc models + rerun notebook

causalpy/pymc_experiments.py

@@ -95,7 +95,9 @@ def plot(self):
             self.datapost.index, self.post_pred["posterior_predictive"].y_hat, ax=ax[0]
         )
         ax[0].plot(self.datapost.index, self.post_y, "k.")
-        ax[0].set(title=f"$R^2$ on pre-intervention data = {self.score:.3f}")
+        ax[0].set(
+            title=f"Pre-intervention Bayesian $R^2$: {self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
+        )
 
         plot_xY(self.datapre.index, self.pre_impact, ax=ax[1])
         plot_xY(self.datapost.index, self.post_impact, ax=ax[1])
@@ -393,7 +395,7 @@ def plot(self):
             ax=ax,
         )
         # create strings to compose title
-        r2 = f"$R^2$ on all data = {self.score:.3f}"
+        r2 = f"Bayesian $R^2$ on all data = {self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
         percentiles = self.discontinuity_at_threshold.quantile([0.03, 1 - 0.03]).values
         ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
         discon = f"Discontinuity at threshold = {self.discontinuity_at_threshold.mean():.2f}, "

causalpy/pymc_models.py

@@ -1,7 +1,7 @@
 import arviz as az
 import numpy as np
 import pymc as pm
-from sklearn.metrics import r2_score
+from arviz import r2_score
 
 
 class ModelBuilder(pm.Model):
@@ -39,12 +39,21 @@ def predict(self, X):
         return post_pred
 
     def score(self, X, y):
-        """Score the predictions :math:`R^2` given inputs ``X`` and outputs ``y``."""
+        """Score the Bayesian :math:`R^2` given inputs ``X`` and outputs ``y``.
+
+        .. caution::
+
+           The Bayesian :math:`R^2` is not the same as the traditional coefficient of determination, https://en.wikipedia.org/wiki/Coefficient_of_determination.
+
+        """
         yhat = self.predict(X)
-        yhat = az.extract(yhat, group="posterior_predictive", var_names="y_hat").mean(
-            dim="sample"
-        )
-        return r2_score(y, yhat)
+        yhat = az.extract(
+            yhat, group="posterior_predictive", var_names="y_hat"
+        ).T.values
+        # Note: First argument must be a 1D array
+        return r2_score(y.flatten(), yhat)
+
+    # .stack(sample=("chain", "draw")
 
 
 class WeightedSumFitter(ModelBuilder):