MMM: Fix Scaling Intercept (#1845)

* patch intercept scaling

* fix condition

* test init

* undo

* fix tests

Author: juanitorduz

pymc_marketing/mmm/mmm.py

@@ -1983,13 +1983,35 @@ def sample_posterior_predictive(
                 date=slice(self.adstock.l_max, None)
             )
 
+        intercept_condition = (
+            "intercept" in sample_posterior_predictive_kwargs.get("var_names", [])
+            and not self.time_varying_intercept
+        )
+
         if original_scale:
+            # We need to expand the intercept to the date dimension
+            # because the target transformer is applied to the date dimension
+            if intercept_condition:
+                posterior_predictive_samples["intercept"] = (
+                    posterior_predictive_samples["intercept"]
+                    .expand_dims(
+                        dim={"date": posterior_predictive_samples["date"]}, axis=0
+                    )
+                    .rename("intercept")
+                )
+
             posterior_predictive_samples = apply_sklearn_transformer_across_dim(
                 data=posterior_predictive_samples,
                 func=self.get_target_transformer().inverse_transform,
                 dim_name="date",
             )
 
+            # We need to remove the date dimension after the inverse transform
+            if intercept_condition:
+                posterior_predictive_samples["intercept"] = (
+                    posterior_predictive_samples["intercept"].isel(date=0)
+                )
+
         return posterior_predictive_samples
 
     def add_lift_test_measurements(

tests/mmm/test_mmm.py

@@ -1000,32 +1000,52 @@ def new_date_ranges_to_test():
     "new_dates",
     new_date_ranges_to_test(),
 )
-@pytest.mark.parametrize("combined", [True, False])
-@pytest.mark.parametrize("original_scale", [True, False])
+@pytest.mark.parametrize(
+    argnames="combined",
+    argvalues=[True, False],
+    ids=["combined", "not_combined"],
+)
+@pytest.mark.parametrize(
+    argnames="original_scale",
+    argvalues=[True, False],
+    ids=["original_scale", "scaled"],
+)
+@pytest.mark.parametrize(
+    argnames="var_names",
+    argvalues=[None, ["mu", "y_sigma", "channel_contribution"], ["mu", "intercept"]],
+    ids=["no_var_names", "var_names", "var_names_with_intercept"],
+)
 def test_new_data_sample_posterior_predictive_method(
     generate_data,
     toy_X,
     model_name: str,
     new_dates: pd.DatetimeIndex,
     combined: bool,
     original_scale: bool,
+    var_names: list[str] | None,
     request,
 ) -> None:
     """This is the method that is used in all the other methods that generate predictions."""
     mmm = request.getfixturevalue(model_name)
     X = generate_data(new_dates)
 
+    kwargs = {"var_names": var_names} if var_names is not None else {}
+
     posterior_predictive = mmm.sample_posterior_predictive(
         X=X,
         extend_idata=False,
         combined=combined,
         original_scale=original_scale,
+        **kwargs,
     )
     pd.testing.assert_index_equal(
         pd.DatetimeIndex(posterior_predictive.coords["date"]),
         new_dates,
     )
 
+    if var_names is not None:
+        assert var_names == list(posterior_predictive.data_vars)
+
 
 @pytest.mark.parametrize(
     "predictions",