Add functionality: allow sample_posterior_predictive to test MMM performance with test data

[nheusch-se]

Hi team,

since this is my first post, just let me say upfront: I love working with pymc-marketing; you've done an amazing job- it's a great package!

I've been trying to look at the predictive performance of an estimated MMM on a test set. In other words - after estimating the model on training data (spends/covariates and sales for 2020-2023), how good are the model's predicted sales for 2024 (when I give it the spends/covariates for 2024)? I know this is not necessarily the goal of causal modelling (after all, I'm after the "true" parameters for adstock and saturation), but in practice decent predictive performance can also be helpful to convince users of the MMM.

BaseMMM.sample_posterior_predictive already provides this functionality: to generate the predicted y's for my training period (2020-2023), I can use
y_train_pred = mmm.sample_posterior_predictive(X_pred=X_train, extend_idata=True, include_last_observations=True)
This is also very handy, because it adds the posterior_predictive to the idata object. That's useful because I save the model and with it the idata object (and can hence retrieve everything again later on).

Now the problem: to generate the predicted y's for my test period (2024), I can use
y_test_pred = mmm.sample_posterior_predictive(X_pred=X_test, extend_idata=True, include_last_observations=True)
However, this overwrites the posterior predictive (for the training period) in the idata object. Hence, when saving my model - I now have to choose whether I want the posterior predictions for either the training or the test period in my idata object. The alternative is to use extend_idata=False; then my saved model has the data from posterior predictive for the training data, but nothing for the test period at all.

This is a little annoying, because pymc actually allows pm.sample_posterior_predictive(... , predictions=True). In this case, the test data (X's for the test period and the posterior predictive) get added to the idata object, without overwriting anything, as idata.predictions and idata.predictions_constant_data. I could now save my model, and would have everything included.

It would be great to allow mmm.sample_posterior_predictive() to use the option predictions=True. Currently, I can already pass it as **sample_posterior_predictive_kwargs, which then gets passed to pymc.sample_posterior_predictive. However, this leads to an error, because when passing predictions=True to mmm.sample_posterior_predictive(), it still tries to extract the "posterior_predictive" from idata. However, in this case it should extract idata.predictions. To get the correct behaviour, predictions=True would hence need to be a 'native' parameter of mmm.sample_posterior_predictive, which then also leads to the extraction of "predictions" from idata.

I've gotten around this with some awkward monkey patching (creating a new function sample_posterior_predictive_test; sorry, it's a little verbose and would have been easier to just change the function sample_posterior_predictive to natively accept the predictions=True parameter, but this was my first go), but it would be great to add this as a native feature of pymc-marketing. I'm happy to help - I hope this description makes sense to you in the first place.

# write a function, which calculate the posterior predictive for the test set and adds it to the idata object
from xarray import DataArray
def sample_posterior_predictive_test(
    self,
    X_pred,
    extend_idata: bool = True,
    combined: bool = True,
    include_last_observations: bool = False,
    original_scale: bool = True,
    **sample_posterior_predictive_kwargs,
) -> DataArray:
    """Sample from the model's posterior predictive distribution.

    Parameters
    ----------
    X_pred : array, shape (n_pred, n_features)
        The input data used for prediction.
    extend_idata : bool, optional
        Boolean determining whether the predictions should be added to inference data object. Defaults to True.
    combined: bool, optional
        Combine chain and draw dims into sample. Won't work if a dim named sample already exists. Defaults to True.
    include_last_observations: bool, optional
        Boolean determining whether to include the last observations of the training data in order to carry over
        costs with the adstock transformation. Assumes that X_pred are the next predictions following the
        training data.Defaults to False.
    original_scale: bool, optional
        Boolean determining whether to return the predictions in the original scale of the target variable.
        Defaults to True.
    **sample_posterior_predictive_kwargs
        Additional arguments to pass to pymc.sample_posterior_predictive

    Returns
    -------
    posterior_predictive_samples : DataArray, shape (n_pred, samples)
        Posterior predictive samples for each input X_pred

    """

    import pymc as pm
    import arviz as az
    from pymc_marketing.mmm.utils import apply_sklearn_transformer_across_dim

    if include_last_observations:
        X_pred = pd.concat(
            [self.X.iloc[-self.adstock.l_max :, :], X_pred], axis=0
        ).sort_values(by=self.date_column)

    self._data_setter(X_pred)
    print(type(pm))

    with self.model:  # sample with new input data
        post_pred = pm.sample_posterior_predictive(
            self.idata, predictions=True, **sample_posterior_predictive_kwargs
        )
        if extend_idata:
            self.idata.extend(post_pred, join="right")  # type: ignore

    posterior_predictive_samples = az.extract(
        post_pred, "predictions", combined=combined
    )

    if include_last_observations:
        posterior_predictive_samples = posterior_predictive_samples.isel(
            date=slice(self.adstock.l_max, None)
        )

    if original_scale:
        posterior_predictive_samples = apply_sklearn_transformer_across_dim(
            data=posterior_predictive_samples,
            func=self.get_target_transformer().inverse_transform,
            dim_name="date",
        )

    return posterior_predictive_samples

# add this functionality to the class
import types
mmm.sample_posterior_predictive_test = types.MethodType(sample_posterior_predictive_test, mmm)

# using the newly added functionality, add the predictions to the idata object.
# Generate predictions on the test set
y_test_pred = mmm.sample_posterior_predictive_test(
    X_pred=X_test, extend_idata=True, include_last_observations=True, 
    original_scale=True
)

To get - later on - the posterior predictive in the original scale, I can use:

from pymc_marketing.mmm.utils import apply_sklearn_transformer_across_dim

ytest_pred_originalscale = apply_sklearn_transformer_across_dim(
                data=mmm.idata.predictions,
                func=mmm.get_target_transformer().inverse_transform,
                dim_name="date"
            )

[wd60622]

Hi, the sample_posterior_predictive is to be used even for X_test. That is why there is the include_last_observations in order to continue the lagged effects from the training data set.

Though sample_posterior_predictive will return a new DataArray, that would be the same one that would be in predictions group

I will create an issue for us to have a better clarify this and work with the predictions=True kwarg which can already be passed.

Let me know if I didn't answer your question!

[nheusch-se]

yes, exactly!

[wd60622]

Great. Thanks for reporting then. I opened up a PR to fix it!

[wd60622]

Note, that the current behavior will still add to the InferenceData despite raising error. You can catch the error in program if you'd like

[nheusch-se]

Thank you so much - this was lightning fast!

[wd60622]

At your service 😄
Thank you for the feedback