Allow for media scaling settings (#1573)

* specify the channel scaling

* implement the channel scaling specification

* add test for the scaling

* add additional checks

* set nans to zero

* add nans to the test

Author: williambdean

pymc_marketing/mmm/multidimensional.py

@@ -17,6 +17,7 @@
 
 import json
 import warnings
+from copy import deepcopy
 from typing import Any, Literal, Protocol
 
 import arviz as az
@@ -267,17 +268,30 @@ def __init__(
         self.dims = dims
 
         if isinstance(scaling, dict):
+            scaling = deepcopy(scaling)
+
+            if "channel" not in scaling:
+                scaling["channel"] = VariableScaling(method="max", dims=self.dims)
+            if "target" not in scaling:
+                scaling["target"] = VariableScaling(method="max", dims=self.dims)
+
             scaling = Scaling(**scaling)
 
         self.scaling: Scaling = scaling or Scaling(
-            target=VariableScaling(method="max", dims=self.dims)
+            target=VariableScaling(method="max", dims=self.dims),
+            channel=VariableScaling(method="max", dims=self.dims),
         )
 
         if set(self.scaling.target.dims).difference([*self.dims, "date"]):
             raise ValueError(
                 f"Target scaling dims {self.scaling.target.dims} must contain {self.dims} and 'date'"
             )
 
+        if set(self.scaling.channel.dims).difference([*self.dims, "channel", "date"]):
+            raise ValueError(
+                f"Channel scaling dims {self.scaling.channel.dims} must contain {self.dims}, 'channel', and 'date'"
+            )
+
         model_config = model_config if model_config is not None else {}
         sampler_config = sampler_config
         model_config = parse_model_config(
@@ -864,11 +878,21 @@ def forward_pass(
 
     def _compute_scales(self) -> None:
         """Compute and save scaling factors for channels and target."""
-        method = getattr(self.xarray_dataset, self.scaling.target.method)
-        self.scalers = method(dim=("date", *self.dims))
-        self.scalers["_target"] = method(dim=("date", *self.scaling.target.dims))[
-            "_target"
-        ]
+        self.scalers = xr.Dataset()
+
+        channel_method = getattr(
+            self.xarray_dataset["_channel"],
+            self.scaling.channel.method,
+        )
+        self.scalers["_channel"] = channel_method(
+            dim=("date", *self.scaling.channel.dims)
+        )
+
+        target_method = getattr(
+            self.xarray_dataset["_target"],
+            self.scaling.target.method,
+        )
+        self.scalers["_target"] = target_method(dim=("date", *self.scaling.target.dims))
 
     def get_scales_as_xarray(self) -> dict[str, xr.DataArray]:
         """Return the saved scaling factors as xarray DataArrays.
@@ -1023,7 +1047,7 @@ def build_model(
             _channel_scale = pm.Data(
                 "channel_scale",
                 self.scalers._channel.values,
-                dims="channel",
+                dims=self.scalers._channel.dims,
             )
             _target_scale = pm.Data(
                 "target_scale",
@@ -1048,7 +1072,12 @@ def build_model(
             )
 
             # Scale `channel_data` and `target`
-            channel_data_ = _channel_data / _channel_scale
+            channel_dim_handler = create_dim_handler(("date", *self.dims, "channel"))
+            channel_data_ = _channel_data / channel_dim_handler(
+                _channel_scale,
+                self.scalers._channel.dims,
+            )
+            channel_data_ = pt.switch(pt.isnan(channel_data_), 0.0, channel_data_)
             channel_data_.name = "channel_data_scaled"
             channel_data_.dims = ("date", *self.dims, "channel")
 

pymc_marketing/mmm/scaling.py

@@ -71,3 +71,7 @@ class Scaling(BaseModel):
         ...,
         description="The scaling for the target variable.",
     )
+    channel: VariableScaling = Field(
+        ...,
+        description="The scaling for the channel variable.",
+    )

tests/mmm/test_multidimensional.py

@@ -28,7 +28,7 @@
     MMM,
     create_event_mu_effect,
 )
-from pymc_marketing.mmm.scaling import VariableScaling
+from pymc_marketing.mmm.scaling import Scaling, VariableScaling
 from pymc_marketing.prior import Prior
 
 
@@ -123,12 +123,14 @@ def single_dim_data():
     # Generate random channel data
     channel_1 = np.random.randint(100, 500, size=len(date_range))
     channel_2 = np.random.randint(100, 500, size=len(date_range))
+    channel_3 = np.nan
 
     df = pd.DataFrame(
         {
             "date": date_range,
             "channel_1": channel_1,
             "channel_2": channel_2,
+            "channel_3": channel_3,
         }
     )
     # Target is sum of channels with noise
@@ -137,7 +139,7 @@ def single_dim_data():
         + df["channel_2"]
         + np.random.randint(100, 300, size=len(date_range))
     )
-    X = df[["date", "channel_1", "channel_2"]].copy()
+    X = df[["date", "channel_1", "channel_2", "channel_3"]].copy()
 
     return X, df.set_index(["date"])["target"].copy()
 
@@ -161,14 +163,16 @@ def multi_dim_data():
         for date in date_range:
             channel_1 = np.random.randint(100, 500)
             channel_2 = np.random.randint(100, 500)
+            channel_3 = np.nan
             target = channel_1 + channel_2 + np.random.randint(50, 150)
-            records.append((date, country, channel_1, channel_2, target))
+            records.append((date, country, channel_1, channel_2, channel_3, target))
 
     df = pd.DataFrame(
-        records, columns=["date", "country", "channel_1", "channel_2", "target"]
+        records,
+        columns=["date", "country", "channel_1", "channel_2", "channel_3", "target"],
     )
 
-    X = df[["date", "country", "channel_1", "channel_2"]].copy()
+    X = df[["date", "country", "channel_1", "channel_2", "channel_3"]].copy()
 
     return X, df["target"].copy()
 
@@ -208,7 +212,7 @@ def test_fit(
     mmm = MMM(
         date_column="date",
         target_column="target",
-        channel_columns=["channel_1", "channel_2"],
+        channel_columns=["channel_1", "channel_2", "channel_3"],
         dims=dims,
         adstock=adstock,
         saturation=saturation,
@@ -296,7 +300,7 @@ def test_sample_posterior_predictive_new_data(single_dim_data, mock_pymc_sample)
     mmm = MMM(
         date_column="date",
         target_column="target",
-        channel_columns=["channel_1", "channel_2"],
+        channel_columns=["channel_1", "channel_2", "channel_3"],
         adstock=adstock,
         saturation=saturation,
     )
@@ -307,6 +311,11 @@ def test_sample_posterior_predictive_new_data(single_dim_data, mock_pymc_sample)
 
     mmm.sample_posterior_predictive(X_train, extend_idata=True, random_seed=42)
 
+    def no_null_values(ds):
+        return ds.y.isnull().mean()
+
+    np.testing.assert_allclose(no_null_values(mmm.idata.posterior_predictive), 0)
+
     # Sample posterior predictive on new data
     out_of_sample_idata = mmm.sample_posterior_predictive(
         X_new, extend_idata=False, random_seed=42
@@ -318,6 +327,8 @@ def test_sample_posterior_predictive_new_data(single_dim_data, mock_pymc_sample)
         "there should be a 'posterior_predictive' group in the inference data."
     )
 
+    np.testing.assert_allclose(no_null_values(out_of_sample_idata), 0)
+
     # Check the shape of that group. We expect the new date dimension to match X_new length
     # plus no addition if we didn't set include_last_observations (which is False by default).
     assert "date" in out_of_sample_idata.dims, (
@@ -349,7 +360,7 @@ def test_sample_posterior_predictive_same_data(single_dim_data, mock_pymc_sample
     mmm = MMM(
         date_column="date",
         target_column="target",
-        channel_columns=["channel_1", "channel_2"],
+        channel_columns=["channel_1", "channel_2", "channel_3"],
         adstock=adstock,
         saturation=saturation,
     )
@@ -587,7 +598,7 @@ def test_check_for_incompatible_dims(adstock, saturation, dims) -> None:
     kwargs = dict(
         date_column="date",
         target_column="target",
-        channel_columns=["channel_1", "channel_2"],
+        channel_columns=["channel_1", "channel_2", "channel_3"],
     )
     with pytest.raises(ValueError):
         MMM(
@@ -608,7 +619,7 @@ def test_different_target_scaling(method, multi_dim_data, mock_pymc_sample) -> N
         scaling=scaling,
         date_column="date",
         target_column="target",
-        channel_columns=["channel_1", "channel_2"],
+        channel_columns=["channel_1", "channel_2", "channel_3"],
         dims=("country",),
     )
     assert mmm.scaling.target == VariableScaling(method=method, dims=())
@@ -645,7 +656,7 @@ def test_target_scaling_raises() -> None:
             scaling=scaling,
             date_column="date",
             target_column="target",
-            channel_columns=["channel_1", "channel_2"],
+            channel_columns=["channel_1", "channel_2", "channel_3"],
         )
 
 
@@ -664,7 +675,7 @@ def test_target_scaling_and_contributions(
         scaling=scaling,
         date_column="date",
         target_column="target",
-        channel_columns=["channel_1", "channel_2"],
+        channel_columns=["channel_1", "channel_2", "channel_3"],
         dims=("country",),
     )
 
@@ -680,3 +691,51 @@ def test_target_scaling_and_contributions(
         mmm.fit(X, y)
     except Exception as e:
         pytest.fail(f"Unexpected error: {e}")
+
+
+@pytest.mark.parametrize(
+    "dims, expected_dims",
+    [
+        ((), ("country", "channel")),
+        (("country",), ("channel",)),
+        (("channel",), ("country",)),
+    ],
+    ids=["country-channel", "country", "channel"],
+)
+def test_channel_scaling(multi_dim_data, dims, expected_dims, mock_pymc_sample) -> None:
+    X, y = multi_dim_data
+
+    scaling = {"channel": {"method": "mean", "dims": dims}}
+    mmm = MMM(
+        adstock=GeometricAdstock(l_max=2),
+        saturation=LogisticSaturation(),
+        scaling=scaling,
+        date_column="date",
+        target_column="target",
+        channel_columns=["channel_1", "channel_2", "channel_3"],
+        dims=("country",),
+    )
+
+    mmm.fit(X, y)
+
+    assert mmm.scalers._channel.dims == expected_dims
+
+
+def test_scaling_dict_doesnt_mutate() -> None:
+    scaling = {}
+    dims = ("country",)
+    mmm = MMM(
+        adstock=GeometricAdstock(l_max=2),
+        saturation=LogisticSaturation(),
+        scaling=scaling,
+        date_column="date",
+        target_column="target",
+        channel_columns=["channel_1", "channel_2", "channel_3"],
+        dims=dims,
+    )
+
+    assert scaling == {}
+    assert mmm.scaling == Scaling(
+        target=VariableScaling(method="max", dims=dims),
+        channel=VariableScaling(method="max", dims=dims),
+    )