Increased computation time of mmm.sample_posterior_predictive with specified var_names when original_scale=True

[OlesyaNori]

When working with the model, we needed to repeatedly request predictions of the model's factor contributions using sample_posterior_predictive. The exact code used was:

    samples = model.sample_posterior_predictive(
                X_pred=X_test,
                extend_idata=False,
                include_last_observations=True,
                original_scale=True,
                var_names=["channel_contributions", "control_contributions", "yearly_seasonality_contribution"],
                progressbar=False,
            )

Each calculation took approximately 1.5 minutes, despite X_test containing only 52 rows. Attempts to speed up the transformations using compile_kwargs={'mode':'NUMBA'} or compile_kwargs={'mode':'JAX'} didn't result in any performance improvement.

However, when we modified the code to:

    samples = model.sample_posterior_predictive(
        X_pred=X_test,
        extend_idata=False,
        include_last_observations=True,
        original_scale=False,
        var_names=["channel_contributions", "control_contributions", "yearly_seasonality_contribution"],
        progressbar=False,
    )   
    contributions_over_time_test = pd.DataFrame(samples["channel_contributions"].mean(dim=["sample"]).values * model.get_target_transformer()["scaler"].scale_,  columns=channel_columns, index=X_test.date)

The computation time was reduced to just 2 seconds. The experiments were conducted on versions 0.14.0 and 0.15.1

[juanitorduz]

Thanks for reporting this 👀 . Maybe xarray vectorization is not as fast as we expected ... we can try to use your method.

[OriolAbril]

If this comes from this else in sample_posterior_predictive https://github.com/pymc-labs/pymc-marketing/blob/main/pymc_marketing/mmm/mmm.py#L2251 which calls apply_sklearn_transformer_across_dim we can definitely look at speeding this up. I see vectorize=True is hardcoded and from apply_ufunc docs:

vectorize (bool, optional) – If True, then assume func only takes arrays defined over core dimensions as input and vectorize it automatically with numpy.vectorize(). This option exists for convenience, but is almost always slower than supplying a pre-vectorized function.