Merge branch 'main' into dependabot/github_actions/actions-0c2b297cc4

Author: drbenvincent

causalpy/experiments/panel_regression.py

@@ -557,7 +557,7 @@ def get_plot_data_bayesian(self, **kwargs: Any) -> pd.DataFrame:
         """
         # Get posterior predictions
         if isinstance(self.model, PyMCModel):
-            mu = self.model.idata.posterior["mu"]  # type: ignore[attr-defined]
+            mu = self.model.idata.posterior["mu"]  # type: ignore[union-attr]
             pred_mean = mu.mean(dim=["chain", "draw"]).values.flatten()
             pred_lower = mu.quantile(0.025, dim=["chain", "draw"]).values.flatten()
             pred_upper = mu.quantile(0.975, dim=["chain", "draw"]).values.flatten()
@@ -673,14 +673,17 @@ def plot_unit_effects(
 
         if isinstance(self.model, PyMCModel):
             # Bayesian: get posterior means
-            beta = self.model.idata.posterior["beta"]  # type: ignore[attr-defined]
+            beta = self.model.idata.posterior["beta"]  # type: ignore[union-attr]
             unit_fe_indices = [self.labels.index(name) for name in unit_fe_names]
 
             # Get mean and std for each unit FE
             fe_means = []
             for idx in unit_fe_indices:
                 fe_means.append(
-                    float(beta.sel(coeffs=self.labels[idx]).mean(dim=["chain", "draw"]))
+                    beta.sel(coeffs=self.labels[idx])
+                    .mean(dim=["chain", "draw"])
+                    .squeeze("treated_units", drop=True)
+                    .item()
                 )
 
             ax.hist(

causalpy/tests/test_panel_regression.py

@@ -290,6 +290,9 @@ def test_panel_regression_plot_coefficients(mock_pymc_sample, small_panel_data):
 
 
 @pytest.mark.integration
+@pytest.mark.filterwarnings(
+    "error:Conversion of an array with ndim > 0 to a scalar:DeprecationWarning"
+)
 def test_panel_regression_plot_unit_effects(mock_pymc_sample, small_panel_data):
     """Test plot_unit_effects method."""
     result = cp.PanelRegression(