towards a more unified scoring (r2) approach

Author: drbenvincent

causalpy/experiments/interrupted_time_series.py

@@ -239,8 +239,8 @@ def _bayesian_plot(
 
         ax[0].set(
             title=f"""
-            Pre-intervention Bayesian $R^2$: {round_num(self.score.r2, round_to)}
-            (std = {round_num(self.score.r2_std, round_to)})
+            Pre-intervention Bayesian $R^2$: {round_num(self.score["unit_r2"], round_to)}
+            (std = {round_num(self.score["unit_r2_std"], round_to)})
             """
         )
 

causalpy/experiments/regression_discontinuity.py

@@ -256,7 +256,7 @@ def _bayesian_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]
         labels = ["Posterior mean"]
 
         # create strings to compose title
-        title_info = f"{round_num(self.score.r2, round_to)} (std = {round_num(self.score.r2_std, round_to)})"
+        title_info = f"{round_num(self.score['unit_r2'], round_to)} (std = {round_num(self.score['unit_r2_std'], round_to)})"
         r2 = f"Bayesian $R^2$ on all data = {title_info}"
         percentiles = self.discontinuity_at_threshold.quantile([0.03, 1 - 0.03]).values
         ci = (

causalpy/experiments/regression_kink.py

@@ -227,7 +227,7 @@ def _bayesian_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]
         labels = ["Posterior mean"]
 
         # create strings to compose title
-        title_info = f"{round_num(self.score.r2, round_to)} (std = {round_num(self.score.r2_std, round_to)})"
+        title_info = f"{round_num(self.score['unit_r2'], round_to)} (std = {round_num(self.score['unit_r2_std'], round_to)})"
         r2 = f"Bayesian $R^2$ on all data = {title_info}"
         percentiles = self.gradient_change.quantile([0.03, 1 - 0.03]).values
         ci = (

causalpy/pymc_models.py

@@ -68,8 +68,8 @@ class PyMCModel(pm.Model):
     >>> model.fit(X, y)
     Inference data...
     >>> model.score(X, y)  # doctest: +ELLIPSIS
-    r2        ...
-    r2_std    ...
+    unit_r2        ...
+    unit_r2_std    ...
     dtype: float64
     >>> X_new = rng.normal(loc=0, scale=1, size=(20, 2))
     >>> model.predict(X_new)
@@ -203,27 +203,32 @@ def score(self, X: xr.DataArray, y: xr.DataArray) -> pd.Series:
         mu = self.predict(X)
         mu_data = az.extract(mu, group="posterior_predictive", var_names="mu")
 
-        # Handle both single and multiple treated units
+        # Always use the multiple treated unit convention for consistency
+        scores = {}
+
         if "treated_units" in mu_data.dims:
-            # Multiple treated units - we need to score each unit separately
+            # Multiple treated units - score each unit separately
             treated_units = mu_data.coords["treated_units"].values
-            scores = {}
-
             for unit in treated_units:
                 unit_mu = mu_data.sel(treated_units=unit).T  # (sample, obs_ind)
                 unit_y = y.sel(treated_units=unit).data
                 unit_score = r2_score(unit_y, unit_mu.data)
-
-                # Flatten the r2_score results into the expected format
                 scores[f"{unit}_r2"] = unit_score["r2"]
                 scores[f"{unit}_r2_std"] = unit_score["r2_std"]
-
-            return pd.Series(scores)
         else:
-            # Single treated unit - transpose to match expected format
+            # Single treated unit - determine unit name and use same format
+            if hasattr(y, "coords") and "treated_units" in y.coords:
+                unit_name = y.coords["treated_units"].values[0]
+            else:
+                unit_name = "unit"  # Fallback for backwards compatibility
+
             mu_data = mu_data.T
             y_data = y.data.squeeze() if y.data.ndim > 1 else y.data
-            return r2_score(y_data, mu_data.data)
+            unit_score = r2_score(y_data, mu_data.data)
+            scores[f"{unit_name}_r2"] = unit_score["r2"]
+            scores[f"{unit_name}_r2_std"] = unit_score["r2_std"]
+
+        return pd.Series(scores)
 
     def calculate_impact(
         self, y_true: xr.DataArray, y_pred: az.InferenceData

causalpy/tests/test_pymc_models.py

@@ -120,6 +120,9 @@ def test_fit_predict(self, coords, rng) -> None:
         ).shape == (20, 2 * 2)
         assert isinstance(score, pd.Series)
         assert score.shape == (2,)
+        # Test that the score follows the new standardized format
+        assert "unit_r2" in score.index
+        assert "unit_r2_std" in score.index
         assert isinstance(predictions, az.InferenceData)