code simplifications by always having a treated_units dimension

Author: drbenvincent

causalpy/experiments/synthetic_control.py

@@ -240,17 +240,12 @@ def _bayesian_plot(
                 f"treated_unit '{treated_unit}' not found. Available units: {self.treated_units}"
             )
 
-        # For multi-unit, select primary unit for main plot
-        if len(self.treated_units) > 1:
-            pre_pred_plot = self.pre_pred["posterior_predictive"].mu.sel(
-                treated_units=treated_unit
-            )
-            post_pred_plot = self.post_pred["posterior_predictive"].mu.sel(
-                treated_units=treated_unit
-            )
-        else:
-            pre_pred_plot = self.pre_pred["posterior_predictive"].mu
-            post_pred_plot = self.post_pred["posterior_predictive"].mu
+        pre_pred_plot = self.pre_pred["posterior_predictive"].mu.sel(
+            treated_units=treated_unit
+        )
+        post_pred_plot = self.post_pred["posterior_predictive"].mu.sel(
+            treated_units=treated_unit
+        )
 
         h_line, h_patch = plot_xY(
             self.datapre.index,
@@ -419,6 +414,7 @@ def _ols_plot(
             # For OLS, predictions might be simple arrays
             post_pred_values = np.squeeze(self.post_pred)
         except (TypeError, AttributeError):
+            # TODO: WILL THIS PATH EVERY BIT HIT?
             # For PyMC predictions (InferenceData)
             post_pred_values = (
                 az.extract(self.post_pred, group="posterior_predictive", var_names="mu")
@@ -534,40 +530,19 @@ def get_plot_data_bayesian(
             self.post_pred, group="posterior_predictive", var_names="mu"
         ).mean("sample")
 
-        if len(self.treated_units) > 1:
-            # Multi-unit case: extract primary unit
-            pre_data["prediction"] = pre_pred_vals.sel(
-                treated_units=treated_unit
-            ).values
-            post_data["prediction"] = post_pred_vals.sel(
-                treated_units=treated_unit
-            ).values
-        else:
-            # Single unit case
-            pre_data["prediction"] = pre_pred_vals.values
-            post_data["prediction"] = post_pred_vals.values
+        # Extract predictions for the specified treated unit (always has treated_units dimension)
+        pre_data["prediction"] = pre_pred_vals.sel(treated_units=treated_unit).values
+        post_data["prediction"] = post_pred_vals.sel(treated_units=treated_unit).values
 
-        # HDI intervals for predictions
-        if len(self.treated_units) > 1:
-            pre_hdi = get_hdi_to_df(
-                self.pre_pred["posterior_predictive"].mu.sel(
-                    treated_units=treated_unit
-                ),
-                hdi_prob=hdi_prob,
-            )
-            post_hdi = get_hdi_to_df(
-                self.post_pred["posterior_predictive"].mu.sel(
-                    treated_units=treated_unit
-                ),
-                hdi_prob=hdi_prob,
-            )
-        else:
-            pre_hdi = get_hdi_to_df(
-                self.pre_pred["posterior_predictive"].mu, hdi_prob=hdi_prob
-            )
-            post_hdi = get_hdi_to_df(
-                self.post_pred["posterior_predictive"].mu, hdi_prob=hdi_prob
-            )
+        # HDI intervals for predictions (always use treated_units dimension)
+        pre_hdi = get_hdi_to_df(
+            self.pre_pred["posterior_predictive"].mu.sel(treated_units=treated_unit),
+            hdi_prob=hdi_prob,
+        )
+        post_hdi = get_hdi_to_df(
+            self.post_pred["posterior_predictive"].mu.sel(treated_units=treated_unit),
+            hdi_prob=hdi_prob,
+        )
 
         # Extract only the lower and upper columns and ensure proper indexing
         pre_lower_upper = pre_hdi.iloc[:, [0, -1]].values  # Get first and last columns
@@ -587,17 +562,13 @@ def get_plot_data_bayesian(
             .sel(treated_units=treated_unit)
             .values
         )
-        # Impact HDI intervals - use primary unit
-        if len(self.treated_units) > 1:
-            pre_impact_hdi = get_hdi_to_df(
-                self.pre_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob
-            )
-            post_impact_hdi = get_hdi_to_df(
-                self.post_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob
-            )
-        else:
-            pre_impact_hdi = get_hdi_to_df(self.pre_impact, hdi_prob=hdi_prob)
-            post_impact_hdi = get_hdi_to_df(self.post_impact, hdi_prob=hdi_prob)
+        # Impact HDI intervals (always use treated_units dimension)
+        pre_impact_hdi = get_hdi_to_df(
+            self.pre_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob
+        )
+        post_impact_hdi = get_hdi_to_df(
+            self.post_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob
+        )
 
         # Extract only the lower and upper columns for impact HDI
         pre_impact_lower_upper = pre_impact_hdi.iloc[:, [0, -1]].values
@@ -614,7 +585,7 @@ def _get_score_title(self, round_to=None):
         """Generate appropriate score title based on model type and number of treated units"""
         if isinstance(self.model, PyMCModel):
             if isinstance(self.score, pd.Series):
-                # Check if it's multi-unit format (has unit-specific keys)
+                # Now consistently has unit-specific keys for all cases
                 if len(self.treated_units) > 1:
                     mean_r2 = self.score.filter(regex=r".*_r2$").mean()
                     mean_r2_std = self.score.filter(regex=r".*_r2_std$").mean()
@@ -623,10 +594,11 @@ def _get_score_title(self, round_to=None):
                     (avg std = {round_num(mean_r2_std, round_to)})
                     """
                 else:
-                    # Single treated unit - Series has 'r2' and 'r2_std' keys
+                    # Single treated unit - use unit-specific keys
+                    unit_name = self.treated_units[0]
                     return 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[f"{unit_name}_r2"], round_to)}
+                    (std = {round_num(self.score[f"{unit_name}_r2_std"], round_to)})
                     """
             else:
                 # Fallback for non-Series score (shouldn't happen with WeightedSumFitter)

causalpy/pymc_models.py

@@ -110,7 +110,7 @@ def _data_setter(self, X) -> None:
 
         with self:
             if has_treated_units:
-                # Multiple treated units - get the number from the model coordinates
+                # Get the number of treated units from the model coordinates
                 treated_units_coord = getattr(self, "coords", {}).get(
                     "treated_units", []
                 )
@@ -122,7 +122,7 @@ def _data_setter(self, X) -> None:
                     coords={"obs_ind": np.arange(new_no_of_observations)},
                 )
             else:
-                # Single treated unit case
+                # Legacy case - this shouldn't happen with new WeightedSumFitter
                 pm.set_data(
                     {"X": X, "y": np.zeros(new_no_of_observations)},
                     coords={"obs_ind": np.arange(new_no_of_observations)},
@@ -378,28 +378,29 @@ def build_model(self, X, y, coords):
             n_predictors = X.shape[1]
             X = pm.Data("X", X, dims=["obs_ind", "coeffs"])
 
-            # Check if we have multiple treated units
-            if y.ndim > 1 and y.shape[1] > 1:
-                # Multiple treated units case
-                y = pm.Data("y", y, dims=["obs_ind", "treated_units"])
-                beta = pm.Dirichlet(
-                    "beta", a=np.ones(n_predictors), dims=["treated_units", "coeffs"]
-                )
-                sigma = pm.HalfNormal("sigma", 1, dims="treated_units")
-                mu = pm.Deterministic(
-                    "mu", pt.dot(X, beta.T), dims=["obs_ind", "treated_units"]
-                )
-                pm.Normal(
-                    "y_hat", mu, sigma, observed=y, dims=["obs_ind", "treated_units"]
-                )
+            # Always use treated_units dimension for consistency
+            # Convert to numpy array if it's an xarray DataArray
+            if hasattr(y, "values"):
+                y_data = y.values
             else:
-                # Single treated unit case (backward compatibility)
-                y_data = y[:, 0] if y.ndim > 1 else y
-                y = pm.Data("y", y_data, dims="obs_ind")
-                beta = pm.Dirichlet("beta", a=np.ones(n_predictors), dims="coeffs")
-                sigma = pm.HalfNormal("sigma", 1)
-                mu = pm.Deterministic("mu", pt.dot(X, beta), dims="obs_ind")
-                pm.Normal("y_hat", mu, sigma, observed=y, dims="obs_ind")
+                y_data = np.asarray(y)
+
+            # Ensure y_data has treated_units dimension
+            if y_data.ndim == 1:
+                y_data = y_data.reshape(-1, 1)  # Add treated_units dimension
+            elif y_data.ndim > 1 and y_data.shape[1] == 1:
+                pass  # Already has correct shape
+            # If y_data.ndim > 1 and y_data.shape[1] > 1, it's multi-unit and already correct
+
+            y = pm.Data("y", y_data, dims=["obs_ind", "treated_units"])
+            beta = pm.Dirichlet(
+                "beta", a=np.ones(n_predictors), dims=["treated_units", "coeffs"]
+            )
+            sigma = pm.HalfNormal("sigma", 1, dims="treated_units")
+            mu = pm.Deterministic(
+                "mu", pt.dot(X, beta.T), dims=["obs_ind", "treated_units"]
+            )
+            pm.Normal("y_hat", mu, sigma, observed=y, dims=["obs_ind", "treated_units"])
 
 
 class InstrumentalVariableRegression(PyMCModel):

causalpy/tests/test_multi_unit_sc.py

@@ -210,8 +210,8 @@ def test_single_unit_backward_compatibility(self, single_unit_sc_data):
 
         # Score should still work
         assert isinstance(sc.score, pd.Series)
-        assert "r2" in sc.score.index
-        assert "r2_std" in sc.score.index
+        assert "treated_0_r2" in sc.score.index
+        assert "treated_0_r2_std" in sc.score.index
 
     def test_multi_unit_plotting(self, multi_unit_sc_data):
         """Test that plotting works with multiple treated units."""

causalpy/tests/test_multi_unit_wsf.py

@@ -233,10 +233,9 @@ def test_backward_compatibility_single_unit(self, single_treated_data):
         # Test prediction
         pred = wsf.predict(X)
 
-        # For single unit, should not have treated_units dimension in some places
-        # but should still work correctly
+        # Now always has treated_units dimension, even for single unit
         mu_shape = pred["posterior_predictive"]["mu"].shape
-        expected_shape = (sample_kwargs["chains"], sample_kwargs["draws"], len(X))
+        expected_shape = (sample_kwargs["chains"], sample_kwargs["draws"], len(X), 1)
         assert mu_shape == expected_shape
 
     def test_print_coefficients_multi_unit(self, synthetic_control_data, capsys):
@@ -297,14 +296,16 @@ def test_scoring_single_unit(self, single_treated_data):
         # Test scoring
         score = wsf.score(X, y)
 
-        # For single unit, should have the same format as before
+        # Now consistently uses treated unit name prefix even for single unit
         assert isinstance(score, pd.Series)
-        assert "r2" in score.index
-        assert "r2_std" in score.index
+        assert "treated_0_r2" in score.index
+        assert "treated_0_r2_std" in score.index
 
         # R2 should be reasonable
-        assert score["r2"] >= -1  # R2 can be negative for very bad fits
-        assert score["r2_std"] >= 0  # Standard deviation should be non-negative
+        assert score["treated_0_r2"] >= -1  # R2 can be negative for very bad fits
+        assert (
+            score["treated_0_r2_std"] >= 0
+        )  # Standard deviation should be non-negative
 
     def test_r2_scores_differ_across_units(self, rng):
         """Test that R² scores are different for different treated units.