Making its work better

Author: cetagostini

causalpy/pymc_models.py

@@ -1298,9 +1298,9 @@ def _prepare_idata(self):
         new_idata = self.idata.copy()
         # Get smoothed posterior and sum over state dimension
         smoothed = self.conditional_idata.isel(observed_state=0).rename(
-            {"smoothed_posterior": "y_hat"}
+            {"smoothed_posterior_observed": "y_hat"}
         )
-        y_hat_summed = smoothed.y_hat.sum(dim="state")
+        y_hat_summed = smoothed.y_hat.copy()
 
         # Rename 'time' to 'obs_ind' to match CausalPy conventions
         if "time" in y_hat_summed.dims: