allow for arbitrary additive effects in MVITS model (#1861)

* allow specifying date dim name and change name

* add docstring example

* implement arbitrary mu_effects

* add simple run through test

Author: williambdean

pymc_marketing/customer_choice/mv_its.py

@@ -26,6 +26,7 @@
 from pymc_extras.prior import Prior
 from xarray import DataArray
 
+from pymc_marketing.mmm.additive_effect import MuEffect
 from pymc_marketing.model_builder import ModelBuilder, create_idata_accessor
 from pymc_marketing.model_config import parse_model_config
 
@@ -69,6 +70,10 @@ def __init__(
         model_config = model_config or {}
         model_config = parse_model_config(model_config)
 
+        self.mu_effects: list[MuEffect] = []
+        # extras dims for the likelihood
+        self.dims = ("existing_product",)
+
         super().__init__(model_config=model_config, sampler_config=sampler_config)
 
         self._distribution_checks()
@@ -264,7 +269,7 @@ def build_model(
         """
         self._generate_and_preprocess_model_data(X, y)  # type: ignore
 
-        with pm.Model(coords=self.coords) as model:
+        with pm.Model(coords=self.coords) as self.model:
             # data
             _existing_sales = pm.Data(
                 "existing_sales",
@@ -277,6 +282,9 @@ def build_model(
                 dims="time",
             )
 
+            for mu_effect in self.mu_effects:
+                mu_effect.create_data(self)
+
             # priors
             intercept = self.model_config["intercept"].create_variable(name="intercept")
 
@@ -301,11 +309,13 @@ def build_model(
                 pm.Deterministic("new sales", beta_all[-1])
 
             # expectation
-            mu = pm.Deterministic(
-                "mu",
-                intercept[None, :] - y[:, None] * beta[None, :],
-                dims=("time", "existing_product"),
-            )
+
+            mu = intercept[None, :] - y[:, None] * beta[None, :]
+
+            for mu_effect in self.mu_effects:
+                mu += mu_effect.create_effect(self)
+
+            mu = pm.Deterministic("mu", mu, dims=("time", "existing_product"))
 
             # likelihood
             self.model_config["likelihood"].create_likelihood_variable(
@@ -314,8 +324,6 @@ def build_model(
                 observed=_existing_sales,
             )
 
-        self.model = model
-
     def _data_setter(
         self,
         X: np.ndarray | pd.DataFrame,

pymc_marketing/customer_choice/synthetic_data.py

@@ -57,6 +57,33 @@ def generate_saturated_data(
     data: pd.DataFrame
         The synthetic data generated.
 
+
+    Examples
+    --------
+    Generate some synthetic data for the MVITS model:
+
+    .. code-block:: python
+
+        import numpy as np
+
+        from pymc_marketing.customer_choice import generate_saturated_data
+
+        seed = sum(map(ord, "Saturated Market Data"))
+        rng = np.random.default_rng(seed)
+
+        scenario = {
+            "total_sales_mu": 1_000,
+            "total_sales_sigma": 5,
+            "treatment_time": 40,
+            "n_observations": 100,
+            "market_shares_before": [[0.7, 0.3, 0]],
+            "market_shares_after": [[0.65, 0.25, 0.1]],
+            "market_share_labels": ["competitor", "own", "new"],
+            "random_seed": rng,
+        }
+
+        data = generate_saturated_data(**scenario)
+
     """
     rng: np.random.Generator = (
         random_seed

pymc_marketing/mmm/additive_effect.py

@@ -28,7 +28,7 @@
 from pymc_marketing.mmm.utils import create_index
 
 
-class MMM(Protocol):
+class Model(Protocol):
     """Protocol MMM."""
 
     @property
@@ -43,30 +43,34 @@ def model(self) -> pm.Model:
 class MuEffect(Protocol):
     """Protocol for arbitrary additive mu effect."""
 
-    def create_data(self, mmm: MMM) -> None:
+    def create_data(self, mmm: Model) -> None:
         """Create the required data in the model."""
 
-    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+    def create_effect(self, mmm: Model) -> pt.TensorVariable:
         """Create the additive effect in the model."""
 
-    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
+    def set_data(self, mmm: Model, model: pm.Model, X: xr.Dataset) -> None:
         """Set the data for new predictions."""
 
 
 class FourierEffect:
     """Fourier seasonality additive effect for MMM."""
 
-    def __init__(self, fourier: FourierBase):
+    def __init__(self, fourier: FourierBase, date_dim_name: str = "date"):
         """Initialize the Fourier effect.
 
         Parameters
         ----------
         fourier : FourierBase
+            The FourierBase instance to use for the effect.
+        date_dim_name : str, optional
+            The name of the date dimension in the model, by default "date".
 
         """
         self.fourier = fourier
+        self.date_dim_name: str = date_dim_name
 
-    def create_data(self, mmm: MMM) -> None:
+    def create_data(self, mmm: Model) -> None:
         """Create the required data in the model.
 
         Parameters
@@ -77,16 +81,16 @@ def create_data(self, mmm: MMM) -> None:
         model = mmm.model
 
         # Get dates from model coordinates
-        dates = pd.to_datetime(model.coords["date"])
+        dates = pd.to_datetime(model.coords[self.date_dim_name])
 
         # Add weekday data to the model
         pm.Data(
             f"{self.fourier.prefix}_day",
             self.fourier._get_days_in_period(dates).to_numpy(),
-            dims="date",
+            dims=self.date_dim_name,
         )
 
-    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+    def create_effect(self, mmm: Model) -> pt.TensorVariable:
         """Create the Fourier effect in the model.
 
         Parameters
@@ -107,18 +111,18 @@ def create_effect(self, mmm: MMM) -> pt.TensorVariable:
 
         # Create a deterministic variable for the effect
         dims = (dim for dim in mmm.dims if dim in self.fourier.prior.dims)
-        fourier_dims = ("date", *dims)
+        fourier_dims = (self.date_dim_name, *dims)
         fourier_effect_det = pm.Deterministic(
             f"{self.fourier.prefix}_effect",
             fourier_effect,
             dims=fourier_dims,
         )
 
         # Handle dimensions for the MMM model
-        dim_handler = create_dim_handler(("date", *mmm.dims))
+        dim_handler = create_dim_handler((self.date_dim_name, *mmm.dims))
         return dim_handler(fourier_effect_det, fourier_dims)
 
-    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
+    def set_data(self, mmm: Model, model: pm.Model, X: xr.Dataset) -> None:
         """Set the data for new predictions.
 
         Parameters
@@ -131,7 +135,7 @@ def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
             The dataset for prediction
         """
         # Get dates from the new dataset
-        new_dates = pd.to_datetime(model.coords["date"])
+        new_dates = pd.to_datetime(model.coords[self.date_dim_name])
 
         # Update the data
         new_data = {
@@ -243,12 +247,13 @@ class MockMMM:
 
     """
 
-    def __init__(self, trend: LinearTrend, prefix: str):
+    def __init__(self, trend: LinearTrend, prefix: str, date_dim_name: str = "date"):
         self.trend = trend
         self.prefix = prefix
         self.linear_trend_first_date: pd.Timestamp
+        self.date_dim_name: str = date_dim_name
 
-    def create_data(self, mmm: MMM) -> None:
+    def create_data(self, mmm: Model) -> None:
         """Create the required data in the model.
 
         Parameters
@@ -259,13 +264,13 @@ def create_data(self, mmm: MMM) -> None:
         model: pm.Model = mmm.model
 
         # Create time index data (normalized between 0 and 1)
-        dates = pd.to_datetime(model.coords["date"])
+        dates = pd.to_datetime(model.coords[self.date_dim_name])
         self.linear_trend_first_date = dates[0]
         t = (dates - self.linear_trend_first_date).days.astype(float)
 
-        pm.Data(f"{self.prefix}_t", t, dims="date")
+        pm.Data(f"{self.prefix}_t", t, dims=self.date_dim_name)
 
-    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+    def create_effect(self, mmm: Model) -> pt.TensorVariable:
         """Create the trend effect in the model.
 
         Parameters
@@ -289,19 +294,22 @@ def create_effect(self, mmm: MMM) -> pt.TensorVariable:
         trend_effect = self.trend.apply(t)
 
         # Create deterministic for the trend effect
-        trend_dims = ("date", *self.trend.dims)  # type: ignore
-        trend_non_broadcastable_dims = ("date", *self.trend.non_broadcastable_dims)
+        trend_dims = (self.date_dim_name, *self.trend.dims)  # type: ignore
+        trend_non_broadcastable_dims = (
+            self.date_dim_name,
+            *self.trend.non_broadcastable_dims,
+        )
         trend_effect = pm.Deterministic(
             f"{self.prefix}_effect_contribution",
             trend_effect[create_index(trend_dims, trend_non_broadcastable_dims)],
             dims=trend_non_broadcastable_dims,
         )
 
         # Return the trend effect
-        dim_handler = create_dim_handler(("date", *mmm.dims))
+        dim_handler = create_dim_handler((self.date_dim_name, *mmm.dims))
         return dim_handler(trend_effect, trend_non_broadcastable_dims)
 
-    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
+    def set_data(self, mmm: Model, model: pm.Model, X: xr.Dataset) -> None:
         """Set the data for new predictions.
 
         Parameters
@@ -314,7 +322,7 @@ def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
             The dataset for prediction.
         """
         # Create normalized time index for new data
-        new_dates = pd.to_datetime(model.coords["date"])
+        new_dates = pd.to_datetime(model.coords[self.date_dim_name])
         t = (new_dates - self.linear_trend_first_date).days.astype(float)
 
         # Update the data
@@ -338,13 +346,16 @@ class EventAdditiveEffect(BaseModel):
     reference_date : str
         The arbitrary reference date to calculate distance from events in days. Default
         is "2025-01-01".
+    date_dim_name : str
+        The name of the date dimension in the model. Default is "date".
 
     """
 
     df_events: InstanceOf[pd.DataFrame]
     prefix: str
     effect: EventEffect
     reference_date: str = "2025-01-01"
+    date_dim_name: str = "date"
 
     def model_post_init(self, context: Any, /) -> None:
         """Post initialization of the model."""
@@ -365,7 +376,7 @@ def end_dates(self) -> pd.Series:
         """The end dates of the events."""
         return pd.to_datetime(self.df_events["end_date"])
 
-    def create_data(self, mmm: MMM) -> None:
+    def create_data(self, mmm: Model) -> None:
         """Create the required data in the model.
 
         Parameters
@@ -376,15 +387,15 @@ def create_data(self, mmm: MMM) -> None:
         """
         model: pm.Model = mmm.model
 
-        model_dates = pd.to_datetime(model.coords["date"])
+        model_dates = pd.to_datetime(model.coords[self.date_dim_name])
 
         model.add_coord(self.prefix, self.df_events["name"].to_numpy())
 
         if "days" not in model:
             pm.Data(
                 "days",
                 days_from_reference(model_dates, self.reference_date),
-                dims="date",
+                dims=self.date_dim_name,
             )
 
         pm.Data(
@@ -398,7 +409,7 @@ def create_data(self, mmm: MMM) -> None:
             dims=self.prefix,
         )
 
-    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+    def create_effect(self, mmm: Model) -> pt.TensorVariable:
         """Create the event effect in the model.
 
         Parameters
@@ -430,15 +441,15 @@ def create_basis_matrix(start_ref, end_ref):
         total_effect = pm.Deterministic(
             f"{self.prefix}_total_effect",
             event_effect.sum(axis=1),
-            dims="date",
+            dims=self.date_dim_name,
         )
 
-        dim_handler = create_dim_handler(("date", *mmm.dims))
-        return dim_handler(total_effect, "date")
+        dim_handler = create_dim_handler((self.date_dim_name, *mmm.dims))
+        return dim_handler(total_effect, self.date_dim_name)
 
-    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
+    def set_data(self, mmm: Model, model: pm.Model, X: xr.Dataset) -> None:
         """Set the data for new predictions."""
-        new_dates = pd.to_datetime(model.coords["date"])
+        new_dates = pd.to_datetime(model.coords[self.date_dim_name])
 
         new_data = {
             "days": days_from_reference(new_dates, self.reference_date),

tests/customer_choice/test_mv_its.py

@@ -27,6 +27,8 @@
     generate_saturated_data,
     generate_unsaturated_data,
 )
+from pymc_marketing.mmm.additive_effect import FourierEffect
+from pymc_marketing.mmm.fourier import YearlyFourier
 
 seed = sum(map(ord, "CustomerChoice"))
 rng = np.random.default_rng(seed)
@@ -286,3 +288,43 @@ def test_calculate_counterfactual_raises() -> None:
     match = "Call the 'fit' method first."
     with pytest.raises(RuntimeError, match=match):
         model.calculate_counterfactual()
+
+
+def test_support_for_mu_effects(saturated_data, mock_pymc_sample) -> None:
+    model = MVITS(existing_sales=["competitor", "own"])
+
+    n_order = 5
+    fourier = YearlyFourier(
+        n_order=n_order,
+        prior=Prior(
+            "Laplace",
+            mu=0,
+            b=1,
+            dims=("fourier", "existing_product"),
+        ),
+    )
+    effect = FourierEffect(fourier=fourier, date_dim_name="time")
+
+    model.mu_effects.append(effect)
+    model.sample(
+        saturated_data.loc[:, ["competitor", "own"]],
+        saturated_data["new"],
+        random_seed=rng,
+        sample_prior_predictive_kwargs={"samples": 10},
+    )
+
+    n_time = len(saturated_data)
+    n_existing_products = 2
+    posterior_size = {"chain": 1, "draw": 10}
+
+    assert model.posterior["fourier_effect"].sizes == {
+        **posterior_size,
+        "time": n_time,
+        "existing_product": n_existing_products,
+    }
+    assert model.posterior["fourier_beta"].sizes == {
+        **posterior_size,
+        "fourier": n_order * 2,
+        "existing_product": n_existing_products,
+    }
+    assert model.posterior["fourier"].sizes == {"fourier": n_order * 2}