Support HSGP instance for time_varying_media in MMM (#1881)

* Support HSGP instance for time_varying_media in MMM

Allow passing an HSGPBase instance (e.g., SoftPlusHSGP) to the time_varying_media argument in MMM, enabling custom latent process dimensions and priors. Update logic to handle both boolean and HSGPBase types, and add tests to verify correct behavior and dimension broadcasting for both single- and multi-dimensional cases.

* Changing tests

* update model version

Author: cetagostini

pymc_marketing/mmm/multidimensional.py

@@ -47,6 +47,7 @@
 )
 from pymc_marketing.mmm.events import EventEffect
 from pymc_marketing.mmm.fourier import YearlyFourier
+from pymc_marketing.mmm.hsgp import HSGPBase
 from pymc_marketing.mmm.lift_test import (
     add_lift_measurements_to_likelihood_from_saturation,
     scale_lift_measurements,
@@ -115,7 +116,7 @@ class MMM(ModelBuilder):
     """
 
     _model_type: str = "MMMM (Multi-Dimensional Marketing Mix Model)"
-    version: str = "0.0.1"
+    version: str = "0.0.2"
 
     @validate_call
     def __init__(
@@ -137,8 +138,13 @@ def __init__(
             Field(strict=True, description="Whether to use a time-varying intercept"),
         ] = False,
         time_varying_media: Annotated[
-            bool,
-            Field(strict=True, description="Whether to use time-varying media effects"),
+            bool | InstanceOf[HSGPBase],
+            Field(
+                description=(
+                    "Whether to use time-varying media effects, or pass an HSGP instance "
+                    "(e.g., SoftPlusHSGP) specifying dims and priors."
+                ),
+            ),
         ] = False,
         dims: tuple[str, ...] | None = Field(
             None, description="Additional dimensions for the model."
@@ -757,7 +763,7 @@ def _generate_and_preprocess_model_data(
             for dim in self.xarray_dataset.coords.dims
         }
 
-        if self.time_varying_intercept | self.time_varying_media:
+        if bool(self.time_varying_intercept) or bool(self.time_varying_media):
             self._time_index = np.arange(0, X[self.date_column].unique().shape[0])
             self._time_index_mid = X[self.date_column].unique().shape[0] // 2
             self._time_resolution = (
@@ -1036,7 +1042,7 @@ def build_model(
             for mu_effect in self.mu_effects:
                 mu_effect.create_data(self)
 
-            if self.time_varying_intercept | self.time_varying_media:
+            if bool(self.time_varying_intercept) or bool(self.time_varying_media):
                 time_index = pm.Data(
                     name="time_index",
                     value=self._time_index,
@@ -1066,7 +1072,7 @@ def build_model(
                 )
 
             # Add media logic
-            if self.time_varying_media:
+            if isinstance(self.time_varying_media, bool) and self.time_varying_media:
                 baseline_channel_contribution = pm.Deterministic(
                     name="baseline_channel_contribution",
                     var=self.forward_pass(
@@ -1079,13 +1085,44 @@ def build_model(
                     X=time_index,
                     dims=("date", *self.dims),
                     **self.model_config["media_tvp_config"],
-                ).create_variable("media_latent_process")
+                ).create_variable("media_temporal_latent_multiplier")
 
                 channel_contribution = pm.Deterministic(
                     name="channel_contribution",
                     var=baseline_channel_contribution * media_latent_process[..., None],
                     dims=("date", *self.dims, "channel"),
                 )
+            elif isinstance(self.time_varying_media, HSGPBase):
+                baseline_channel_contribution = self.forward_pass(
+                    x=channel_data_, dims=(*self.dims, "channel")
+                )
+                baseline_channel_contribution.name = "baseline_channel_contribution"
+                baseline_channel_contribution.dims = (
+                    "date",
+                    *self.dims,
+                    "channel",
+                )
+
+                # Register internal time index and build latent process
+                self.time_varying_media.register_data(time_index)
+                media_latent_process = self.time_varying_media.create_variable(
+                    "media_temporal_latent_multiplier"
+                )
+
+                # Determine broadcasting over channel axis
+                media_dims = pm.modelcontext(None).named_vars_to_dims[
+                    media_latent_process.name
+                ]
+                if "channel" in media_dims:
+                    media_broadcast = media_latent_process
+                else:
+                    media_broadcast = media_latent_process[..., None]
+
+                channel_contribution = pm.Deterministic(
+                    name="channel_contribution",
+                    var=baseline_channel_contribution * media_broadcast,
+                    dims=("date", *self.dims, "channel"),
+                )
             else:
                 channel_contribution = pm.Deterministic(
                     name="channel_contribution",
@@ -1681,7 +1718,7 @@ def add_lift_test_measurements(
         # This is coupled with the name of the
         # latent process Deterministic
         time_varying_var_name = (
-            "media_latent_process" if self.time_varying_media else None
+            "media_temporal_latent_multiplier" if self.time_varying_media else None
         )
         add_lift_measurements_to_likelihood_from_saturation(
             df_lift_test=df_lift_test_scaled,

tests/mmm/test_multidimensional.py

@@ -25,7 +25,7 @@
 from pytensor.tensor.basic import TensorVariable
 from scipy.optimize import OptimizeResult
 
-from pymc_marketing.mmm import GeometricAdstock, LogisticSaturation
+from pymc_marketing.mmm import GeometricAdstock, LogisticSaturation, SoftPlusHSGP
 from pymc_marketing.mmm.additive_effect import EventAdditiveEffect, LinearTrendEffect
 from pymc_marketing.mmm.events import EventEffect, GaussianBasis
 from pymc_marketing.mmm.lift_test import _swap_columns_and_last_index_level
@@ -265,7 +265,7 @@ def unstack(data, name):
     if time_varying_intercept:
         assert "intercept_latent_process" in var_names
     if time_varying_media:
-        assert "media_latent_process" in var_names
+        assert "media_temporal_latent_multiplier" in var_names
     if yearly_seasonality is not None:
         assert "fourier_contribution" in var_names
 
@@ -501,6 +501,88 @@ def test_sample_posterior_predictive_partial_overlap_with_include_last_observati
         )
 
 
+@pytest.mark.parametrize(
+    "hsgp_dims",
+    [
+        pytest.param(("date",), id="hsgp-dims=date"),
+        pytest.param(("date", "channel"), id="hsgp-dims=date,channel"),
+    ],
+)
+def test_time_varying_media_with_custom_hsgp_single_dim(single_dim_data, hsgp_dims):
+    """Ensure passing an HSGP instance to time_varying_media works (single-dim)."""
+    X, y = single_dim_data
+
+    # Build HSGP using the new API
+    hsgp = SoftPlusHSGP.parameterize_from_data(
+        X=np.arange(X.shape[0]),
+        dims=hsgp_dims,
+    )
+
+    mmm = MMM(
+        date_column="date",
+        target_column="target",
+        channel_columns=["channel_1", "channel_2", "channel_3"],
+        adstock=GeometricAdstock(l_max=2),
+        saturation=LogisticSaturation(),
+        time_varying_media=hsgp,
+    )
+
+    mmm.build_model(X, y)
+
+    # Check latent multiplier exists with the expected dims
+    var_name = "media_temporal_latent_multiplier"
+    assert var_name in mmm.model.named_vars
+    latent_dims = mmm.model.named_vars_to_dims[var_name]
+    assert latent_dims == hsgp_dims
+
+    # Channel contribution should always be date x channel
+    assert mmm.model.named_vars_to_dims["channel_contribution"] == ("date", "channel")
+
+
+@pytest.mark.parametrize(
+    "hsgp_dims",
+    [
+        pytest.param(("date", "country"), id="hsgp-dims=date,country"),
+        pytest.param(
+            ("date", "country", "channel"), id="hsgp-dims=date,country,channel"
+        ),
+    ],
+)
+def test_time_varying_media_with_custom_hsgp_multi_dim(df, hsgp_dims):
+    """Ensure passing an HSGP instance to time_varying_media works (multi-dim)."""
+    X = df.drop(columns=["y"])
+    y = df["y"]
+
+    hsgp = SoftPlusHSGP.parameterize_from_data(
+        X=np.arange(X.shape[0]),
+        dims=hsgp_dims,
+    )
+
+    mmm = MMM(
+        date_column="date",
+        channel_columns=["C1", "C2"],
+        target_column="y",
+        dims=("country",),
+        adstock=GeometricAdstock(l_max=2),
+        saturation=LogisticSaturation(),
+        time_varying_media=hsgp,
+    )
+
+    mmm.build_model(X, y)
+
+    var_name = "media_temporal_latent_multiplier"
+    assert var_name in mmm.model.named_vars
+    latent_dims = mmm.model.named_vars_to_dims[var_name]
+    assert latent_dims == hsgp_dims
+
+    # Channel contribution should always be date x country x channel
+    assert mmm.model.named_vars_to_dims["channel_contribution"] == (
+        "date",
+        "country",
+        "channel",
+    )
+
+
 def test_sample_posterior_predictive_no_overlap_with_include_last_observations(
     single_dim_data, mock_pymc_sample
 ):