Get marginal_effects hdi estimates (#1885)

* Improve marginal_effects calculation

* Update sensitivity analysis notebook

* Improve MockMMM.idata

Author: PabloRoque

docs/source/notebooks/mmm/mmm_sensitivity_analysis.ipynb

@@ -80,7 +80,9 @@ def run_sweep(
         predictions = []
         for sweep_value in self.sweep_values:
             X_new = self.create_intervention(sweep_value)
-            counterfac = self.mmm.predict(X_new, extend_idata=False, progressbar=False)
+            counterfac = self.mmm.sample_posterior_predictive(
+                X_new, extend_idata=False, combined=False, progressbar=False
+            )
             uplift = counterfac - actual
             predictions.append(uplift)
 
@@ -92,12 +94,13 @@ def run_sweep(
 
         marginal_effects = self.compute_marginal_effects(results, self.sweep_values)
 
-        results = xr.Dataset(
-            {
-                "y": results,
-                "marginal_effects": marginal_effects,
-            }
-        )
+        results = xr.merge(
+            [
+                results,
+                marginal_effects.rename({"y": "marginal_effects"}),
+            ]
+        ).transpose(..., "sweep")
+
         # Add metadata to the results
         results.attrs["sweep_type"] = self.sweep_type
         results.attrs["var_names"] = self.var_names
@@ -129,9 +132,5 @@ def create_intervention(self, sweep_value: float) -> pd.DataFrame:
     def compute_marginal_effects(results, sweep_values) -> xr.DataArray:
         """Compute marginal effects via finite differences from the sweep results."""
         marginal_effects = results.differentiate(coord="sweep")
-        marginal_effects = xr.DataArray(
-            marginal_effects,
-            dims=results.dims,
-            coords=results.coords,
-        )
+
         return marginal_effects

pymc_marketing/mmm/sensitivity_analysis.py

@@ -45,17 +45,21 @@ def __init__(self):
                     n_chains, n_draws, n_dates = 2, 10, 3
 
                     # This is what the sensitivity analysis expects to find
-                    posterior_predictive_data = xr.DataArray(
-                        np.random.normal(size=(n_chains, n_draws, n_dates)),
-                        dims=["chain", "draw", "date"],
+                    posterior_predictive_data = xr.Dataset(
+                        dict(
+                            y=(
+                                ["chain", "draw", "date"],
+                                np.random.normal(size=(n_chains, n_draws, n_dates)),
+                            ),
+                        ),
                         coords={
                             "chain": np.arange(n_chains),
                             "draw": np.arange(n_draws),
                             "date": dates,  # Use the same dates as the DataFrame
                         },
                     )
 
-                    self.posterior_predictive = {"y": posterior_predictive_data}
+                    self.posterior_predictive = posterior_predictive_data
 
                 def __getitem__(self, key):
                     if key == "posterior_predictive":
@@ -81,16 +85,26 @@ def predict(self, X_new, extend_idata=False, progressbar=False):
                 size=(n_chains, n_draws, n_dates),
             )
 
-            return xr.DataArray(
-                data,
-                dims=["chain", "draw", "date"],
+            return xr.Dataset(
+                dict(
+                    y=(
+                        ["chain", "draw", "date"],
+                        data,
+                    )  # Use the same dimensions as expected
+                ),
                 coords={
                     "chain": np.arange(n_chains),
                     "draw": np.arange(n_draws),
                     "date": X_new.index,  # Use the DataFrame index as dates
                 },
             )
 
+        def sample_posterior_predictive(
+            self, X_new, extend_idata=False, combined=False, progressbar=False
+        ):
+            # Mock implementation of sample_posterior_predictive
+            return self.predict(X_new, extend_idata, progressbar)
+
     return MockMMM()