Sensitivity analysis and marginal effects (#1673)

* initial stab at CounterfactualSweep class + associated example notebook

* attempt to add the new notebook to the examples gallery

* delete commented code

* fix example in docs and re-run notebook with some hidden inputs/outputs

* add some TODO's to the notebook

* Update pymc_marketing/mmm/marginal_effects.py

Co-authored-by: Will Dean <[email protected]>

* improve type hinting

* update docstring of plot_marginal_effects method

* Use Literal in type hint

* change to use pymc_marketing.mmm.multidimensional.MMM

* scaling of the marginal effects plot to not put undue emphasis on numerical imprecision

* Results now returned as self contained xr.Dataset. Plot methods are now static methods

* X no longer required as an input to CounterfactualSweep

* remove redundant sweep_values index

* rename to SensivityAnalysis

* compute gradient with xarray instead of numpy

* add MMM.sensitivity_analysis as wrapper to call SensitivityAnalysis

* formatting

* rename notebook

* remove commented code in notebook

* fix scaling + add crosshairs on plots

* combine into a single plot function

* api change, results now stored in idata, and fix crosshairs

* minor tweaks

* better sweep values for additive sweep example

* move plot_sensitivity_analysis into MMMPlotSuite

* rename example in the gallery view. Docs updated

* add functionality to plot y-axis in percentage terms

* add a check for presence of idata.sensitivity_analysis

* update API according to Carlos' suggestions

* predictors -> var_names

* Add tests for SensitivityAnalysis class

* Add tests for plot_sensitivity_analysis in sensitivity analysis

* more tests to increase code coverage of plot code

* update notebook

---------

Co-authored-by: Carlos Trujillo <[email protected]>
Co-authored-by: Will Dean <[email protected]>
Co-authored-by: Juan Orduz <[email protected]>

Author: 4 people

docs/source/gallery/gallery.md

@@ -30,6 +30,11 @@ Welcome to the PyMC-Marketing example gallery! This gallery provides visual navi
 :img-top: ../gallery/images/mmm_multidimensional_example.png
 :link: ../notebooks/mmm/mmm_multidimensional_example.html
 :::
+
+:::{grid-item-card} Sensitivity Analysis and Marginal Effects
+:img-top: ../gallery/images/mmm_sensitivity_analysis.png
+:link: ../notebooks/mmm/mmm_sensitivity_analysis.html
+:::
 ::::
 
 ### Budget Allocation
@@ -110,6 +115,7 @@ Welcome to the PyMC-Marketing example gallery! This gallery provides visual navi
 :img-top: ../gallery/images/mmm_counterfactuals.png
 :link: ../notebooks/mmm/mmm_counterfactuals.html
 :::
+
 ::::
 
 ### Case Studies

docs/source/notebooks/mmm/mmm_sensitivity_analysis.ipynb

@@ -62,6 +62,7 @@
     preprocessing_method_X,
     preprocessing_method_y,
 )
+from pymc_marketing.mmm.sensitivity_analysis import SensitivityAnalysis
 from pymc_marketing.mmm.validating import validation_method_X, validation_method_y
 
 __all__ = [
@@ -90,6 +91,7 @@
     "PeriodicCovFunc",
     "RootSaturation",
     "SaturationTransformation",
+    "SensitivityAnalysis",
     "SoftPlusHSGP",
     "TanhSaturation",
     "TanhSaturationBaselined",

pymc_marketing/mmm/__init__.py

@@ -53,6 +53,7 @@
 )
 from pymc_marketing.mmm.plot import MMMPlotSuite
 from pymc_marketing.mmm.scaling import Scaling, VariableScaling
+from pymc_marketing.mmm.sensitivity_analysis import SensitivityAnalysis
 from pymc_marketing.mmm.tvp import infer_time_index
 from pymc_marketing.mmm.utility import UtilityFunctionType, average_response
 from pymc_marketing.mmm.utils import (
@@ -1433,6 +1434,28 @@ def sample_posterior_predictive(
 
         return posterior_predictive_samples
 
+    @property
+    def sensitivity(self) -> SensitivityAnalysis:
+        """Access sensitivity analysis functionality.
+
+        Returns a SensitivityAnalysis instance that can be used to run
+        counterfactual sweeps on the model.
+
+        Returns
+        -------
+        SensitivityAnalysis
+            An instance configured with this MMM model.
+
+        Examples
+        --------
+        >>> mmm.sensitivity.run_sweep(
+        ...     var_names=["channel_1", "channel_2"],
+        ...     sweep_values=np.linspace(0.5, 2.0, 10),
+        ...     sweep_type="multiplicative",
+        ... )
+        """
+        return SensitivityAnalysis(mmm=self)
+
     def _make_channel_transform(
         self, df_lift_test: pd.DataFrame
     ) -> Callable[[np.ndarray], np.ndarray]:

pymc_marketing/mmm/multidimensional.py

@@ -1243,3 +1243,115 @@ def allocated_contribution_by_channel_over_time(
 
         fig.tight_layout()
         return fig, axes
+
+    def plot_sensitivity_analysis(
+        self,
+        hdi_prob: float = 0.94,
+        ax: plt.Axes | None = None,
+        marginal: bool = False,
+        percentage: bool = False,
+    ) -> plt.Axes:
+        """
+        Plot the counterfactual uplift or marginal effects curve.
+
+        Parameters
+        ----------
+        results : xr.Dataset
+            The dataset containing the results of the sweep.
+        hdi_prob : float, optional
+            The probability for computing the highest density interval (HDI). Default is 0.94.
+        ax : Optional[plt.Axes], optional
+            An optional matplotlib Axes on which to plot. If None, a new Axes is created.
+        marginal : bool, optional
+            If True, plot marginal effects. If False (default), plot uplift.
+        percentage : bool, optional
+            If True, plot the results on the y-axis as percentages, instead of absolute
+            values. Default is False.
+
+        Returns
+        -------
+        plt.Axes
+            The Axes object with the plot.
+        """
+        if ax is None:
+            _, ax = plt.subplots(figsize=(10, 6))
+
+        if percentage and marginal:
+            raise ValueError("Not implemented marginal effects in percentage scale.")
+
+        # Check if sensitivity analysis results exist in idata
+        if not hasattr(self.idata, "sensitivity_analysis"):
+            raise ValueError(
+                "No sensitivity analysis results found in 'self.idata'. "
+                "Please run the sensitivity analysis first using 'mmm.sensitivity.run_sweep()' method."
+            )
+
+        # grab sensitivity analysis results from idata
+        results = self.idata.sensitivity_analysis
+
+        x = results.sweep.values
+        if marginal:
+            y = results.marginal_effects.mean(dim=["chain", "draw"]).sum(dim="date")
+            y_hdi = results.marginal_effects.sum(dim="date")
+            color = "C1"
+            label = "Posterior mean marginal effect"
+            title = "Marginal effects plot"
+            ylabel = r"Marginal effect, $\frac{d\mathbb{E}[Y]}{dX}$"
+        else:
+            if percentage:
+                actual = self.idata.posterior_predictive["y"]
+                y = results.y.mean(dim=["chain", "draw"]).sum(dim="date") / actual.mean(
+                    dim=["chain", "draw"]
+                ).sum(dim="date")
+                y_hdi = results.y.sum(dim="date") / actual.sum(dim="date")
+            else:
+                y = results.y.mean(dim=["chain", "draw"]).sum(dim="date")
+                y_hdi = results.y.sum(dim="date")
+            color = "C0"
+            label = "Posterior mean"
+            title = "Sensitivity analysis plot"
+            ylabel = "Total uplift (sum over dates)"
+
+        ax.plot(x, y, label=label, color=color)
+
+        az.plot_hdi(
+            x,
+            y_hdi,
+            hdi_prob=hdi_prob,
+            color=color,
+            fill_kwargs={"alpha": 0.5, "label": f"{hdi_prob * 100:.0f}% HDI"},
+            plot_kwargs={"color": color, "alpha": 0.5},
+            smooth=False,
+            ax=ax,
+        )
+
+        ax.set(title=title)
+        if results.sweep_type == "absolute":
+            ax.set_xlabel(f"Absolute value of: {results.var_names}")
+        else:
+            ax.set_xlabel(
+                f"{results.sweep_type.capitalize()} change of: {results.var_names}"
+            )
+        ax.set_ylabel(ylabel)
+        plt.legend()
+
+        # Set y-axis limits based on the sign of y values
+        y_values = y.values if hasattr(y, "values") else np.array(y)
+        if np.all(y_values < 0):
+            ax.set_ylim(top=0)
+        elif np.all(y_values > 0):
+            ax.set_ylim(bottom=0)
+
+        ax.yaxis.set_major_formatter(
+            plt.FuncFormatter(lambda x, _: f"{x:.1%}" if percentage else f"{x:,.1f}")
+        )
+
+        # Add reference lines
+        if results.sweep_type == "multiplicative":
+            ax.axvline(x=1, color="k", linestyle="--", alpha=0.5)
+            if not marginal:
+                ax.axhline(y=0, color="k", linestyle="--", alpha=0.5)
+        elif results.sweep_type == "additive":
+            ax.axvline(x=0, color="k", linestyle="--", alpha=0.5)
+
+        return ax

pymc_marketing/mmm/plot.py

@@ -0,0 +1,137 @@
+#   Copyright 2022 - 2025 The PyMC Labs Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+"""Counterfactual sweeps for Marketing Mix Models (MMM)."""
+
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+class SensitivityAnalysis:
+    """SensitivityAnalysis class is used to perform counterfactual analysis on MMM's."""
+
+    def __init__(self, mmm) -> None:
+        """
+        Initialize the SensitivityAnalysis with a reference to the MMM instance.
+
+        Parameters
+        ----------
+        mmm : MMM
+            The marketing mix model instance used for predictions.
+        """
+        self.mmm = mmm
+
+    def run_sweep(
+        self,
+        var_names: list[str],
+        sweep_values: np.ndarray,
+        sweep_type: Literal[
+            "multiplicative", "additive", "absolute"
+        ] = "multiplicative",
+    ) -> xr.Dataset:
+        """Run the model's predict function over the sweep grid and store results.
+
+        Parameters
+        ----------
+        var_names : list[str]
+            List of variable names to intervene on.
+        sweep_values : np.ndarray
+            Array of sweep values.
+        sweep_type : Literal["multiplicative", "additive", "absolute"], optional
+            Type of intervention to apply, by default "multiplicative".
+            - 'multiplicative': Multiply the original predictor values by each sweep value.
+            - 'additive': Add each sweep value to the original predictor values.
+            - 'absolute': Set the predictor values directly to each sweep value (ignoring original values).
+
+        Returns
+        -------
+        xr.Dataset
+            Dataset containing the sensitivity analysis results.
+        """
+        # Validate that idata exists
+        if not hasattr(self.mmm, "idata"):
+            raise ValueError("idata does not exist. Build the model first and fit.")
+
+        # Store parameters for this run
+        self.var_names = var_names
+        self.sweep_values = sweep_values
+        self.sweep_type = sweep_type
+
+        # TODO: Ideally we can use this --------------------------------------------
+        # actual = self.mmm._get_group_predictive_data(
+        #     group="posterior_predictive", original_scale=True
+        # )["y"]
+        actual = self.mmm.idata["posterior_predictive"]["y"]
+        # --------------------------------------------------------------------------
+        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)
+            uplift = counterfac - actual
+            predictions.append(uplift)
+
+        results = (
+            xr.concat(predictions, dim="sweep")
+            .assign_coords(sweep=self.sweep_values)
+            .transpose(..., "sweep")
+        )
+
+        marginal_effects = self.compute_marginal_effects(results, self.sweep_values)
+
+        results = xr.Dataset(
+            {
+                "y": results,
+                "marginal_effects": marginal_effects,
+            }
+        )
+        # Add metadata to the results
+        results.attrs["sweep_type"] = self.sweep_type
+        results.attrs["var_names"] = self.var_names
+
+        # Add results to the MMM's idata
+        if hasattr(self.mmm.idata, "sensitivity_analysis"):
+            delattr(self.mmm.idata, "sensitivity_analysis")
+        self.mmm.idata.add_groups({"sensitivity_analysis": results})  # type: ignore
+
+        return results
+
+    def create_intervention(self, sweep_value: float) -> pd.DataFrame:
+        """Apply the intervention to the predictors."""
+        X_new = self.mmm.X.copy()
+        if self.sweep_type == "multiplicative":
+            for var_name in self.var_names:
+                X_new[var_name] *= sweep_value
+        elif self.sweep_type == "additive":
+            for var_name in self.var_names:
+                X_new[var_name] += sweep_value
+        elif self.sweep_type == "absolute":
+            for var_name in self.var_names:
+                X_new[var_name] = sweep_value
+        else:
+            raise ValueError(f"Unsupported sweep_type: {self.sweep_type}")
+        return X_new
+
+    @staticmethod
+    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