rename primary_unit_name -> treated_unit + revert to no "Unit" title

Author: drbenvincent

causalpy/experiments/synthetic_control.py

@@ -231,22 +231,22 @@ def _bayesian_plot(
         # pre-intervention period
 
         # Get treated unit name - default to first unit if None
-        primary_unit_name = (
+        treated_unit = (
             treated_unit if treated_unit is not None else self.treated_units[0]
         )
 
-        if primary_unit_name not in self.treated_units:
+        if treated_unit not in self.treated_units:
             raise ValueError(
-                f"treated_unit '{primary_unit_name}' not found. Available units: {self.treated_units}"
+                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=primary_unit_name
+                treated_units=treated_unit
             )
             post_pred_plot = self.post_pred["posterior_predictive"].mu.sel(
-                treated_units=primary_unit_name
+                treated_units=treated_unit
             )
         else:
             pre_pred_plot = self.pre_pred["posterior_predictive"].mu
@@ -264,7 +264,7 @@ def _bayesian_plot(
         # Plot observations for primary treated unit
         (h,) = ax[0].plot(
             self.datapre.index,
-            self.datapre_treated.sel(treated_units=primary_unit_name),
+            self.datapre_treated.sel(treated_units=treated_unit),
             "k.",
             label="Observations",
         )
@@ -283,42 +283,40 @@ def _bayesian_plot(
 
         ax[0].plot(
             self.datapost.index,
-            self.datapost_treated.sel(treated_units=primary_unit_name),
+            self.datapost_treated.sel(treated_units=treated_unit),
             "k.",
         )
         # Shaded causal effect for primary treated unit
         h = ax[0].fill_between(
             self.datapost.index,
             y1=post_pred_plot.mean(dim=["chain", "draw"]).values,
-            y2=self.datapost_treated.sel(treated_units=primary_unit_name).values,
+            y2=self.datapost_treated.sel(treated_units=treated_unit).values,
             color="C2",
             alpha=0.25,
             label="Causal impact",
         )
         handles.append(h)
         labels.append("Causal impact")
 
-        ax[0].set(title=f"{self._get_score_title(round_to)}\nUnit")
+        ax[0].set(title=f"{self._get_score_title(round_to)}")
 
         # MIDDLE PLOT -----------------------------------------------
         plot_xY(
             self.datapre.index,
-            self.pre_impact.sel(treated_units=primary_unit_name),
+            self.pre_impact.sel(treated_units=treated_unit),
             ax=ax[1],
             plot_hdi_kwargs={"color": "C0"},
         )
         plot_xY(
             self.datapost.index,
-            self.post_impact.sel(treated_units=primary_unit_name),
+            self.post_impact.sel(treated_units=treated_unit),
             ax=ax[1],
             plot_hdi_kwargs={"color": "C1"},
         )
         ax[1].axhline(y=0, c="k")
         ax[1].fill_between(
             self.datapost.index,
-            y1=self.post_impact.mean(["chain", "draw"]).sel(
-                treated_units=primary_unit_name
-            ),
+            y1=self.post_impact.mean(["chain", "draw"]).sel(treated_units=treated_unit),
             color="C0",
             alpha=0.25,
             label="Causal impact",
@@ -329,7 +327,7 @@ def _bayesian_plot(
         ax[2].set(title="Cumulative Causal Impact")
         plot_xY(
             self.datapost.index,
-            self.post_impact_cumulative.sel(treated_units=primary_unit_name),
+            self.post_impact_cumulative.sel(treated_units=treated_unit),
             ax=ax[2],
             plot_hdi_kwargs={"color": "C1"},
         )
@@ -385,25 +383,25 @@ def _ols_plot(
         counterfactual_label = "Counterfactual"
 
         # Get treated unit name - default to first unit if None
-        primary_unit_name = (
+        treated_unit = (
             treated_unit if treated_unit is not None else self.treated_units[0]
         )
 
-        if primary_unit_name not in self.treated_units:
+        if treated_unit not in self.treated_units:
             raise ValueError(
-                f"treated_unit '{primary_unit_name}' not found. Available units: {self.treated_units}"
+                f"treated_unit '{treated_unit}' not found. Available units: {self.treated_units}"
             )
 
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
         ax[0].plot(
             self.datapre_treated["obs_ind"],
-            self.datapre_treated.sel(treated_units=primary_unit_name),
+            self.datapre_treated.sel(treated_units=treated_unit),
             "k.",
         )
         ax[0].plot(
             self.datapost_treated["obs_ind"],
-            self.datapost_treated.sel(treated_units=primary_unit_name),
+            self.datapost_treated.sel(treated_units=treated_unit),
             "k.",
         )
 
@@ -415,7 +413,7 @@ def _ols_plot(
             ls=":",
             c="k",
         )
-        ax[0].set(title=f"{self._get_score_title(round_to)}\nUnit")
+        ax[0].set(title=f"{self._get_score_title(round_to)}")
         # Shaded causal effect - handle different prediction formats
         try:
             # For OLS, predictions might be simple arrays
@@ -435,9 +433,7 @@ def _ols_plot(
         ax[0].fill_between(
             self.datapost.index,
             y1=post_pred_values,
-            y2=np.squeeze(
-                self.datapost_treated.sel(treated_units=primary_unit_name).data
-            ),
+            y2=np.squeeze(self.datapost_treated.sel(treated_units=treated_unit).data),
             color="C0",
             alpha=0.25,
             label="Causal impact",
@@ -521,13 +517,13 @@ def get_plot_data_bayesian(
         post_data = self.datapost.copy()
 
         # Get treated unit name - default to first unit if None
-        primary_unit_name = (
+        treated_unit = (
             treated_unit if treated_unit is not None else self.treated_units[0]
         )
 
-        if primary_unit_name not in self.treated_units:
+        if treated_unit not in self.treated_units:
             raise ValueError(
-                f"treated_unit '{primary_unit_name}' not found. Available units: {self.treated_units}"
+                f"treated_unit '{treated_unit}' not found. Available units: {self.treated_units}"
             )
 
         # Extract predictions - handle multi-unit case
@@ -541,10 +537,10 @@ def get_plot_data_bayesian(
         if len(self.treated_units) > 1:
             # Multi-unit case: extract primary unit
             pre_data["prediction"] = pre_pred_vals.sel(
-                treated_units=primary_unit_name
+                treated_units=treated_unit
             ).values
             post_data["prediction"] = post_pred_vals.sel(
-                treated_units=primary_unit_name
+                treated_units=treated_unit
             ).values
         else:
             # Single unit case
@@ -555,13 +551,13 @@ def get_plot_data_bayesian(
         if len(self.treated_units) > 1:
             pre_hdi = get_hdi_to_df(
                 self.pre_pred["posterior_predictive"].mu.sel(
-                    treated_units=primary_unit_name
+                    treated_units=treated_unit
                 ),
                 hdi_prob=hdi_prob,
             )
             post_hdi = get_hdi_to_df(
                 self.post_pred["posterior_predictive"].mu.sel(
-                    treated_units=primary_unit_name
+                    treated_units=treated_unit
                 ),
                 hdi_prob=hdi_prob,
             )
@@ -583,21 +579,21 @@ def get_plot_data_bayesian(
         # Impact data - always use primary unit for main dataframe
         pre_data["impact"] = (
             self.pre_impact.mean(dim=["chain", "draw"])
-            .sel(treated_units=primary_unit_name)
+            .sel(treated_units=treated_unit)
             .values
         )
         post_data["impact"] = (
             self.post_impact.mean(dim=["chain", "draw"])
-            .sel(treated_units=primary_unit_name)
+            .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=primary_unit_name), hdi_prob=hdi_prob
+                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=primary_unit_name), hdi_prob=hdi_prob
+                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)