Out-of-Box MultiDimensional MMM

[cetagostini]

Description

Creating an API to support multiple dims.

Related Issue

• Closes #
• Related to Hierarchical Market Mix Modeling #497 Draft [Creating a Hierarchical MMM - Model] #460 MultiDimensional Media Mix Model #1004

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Modules affected

• MMM
• CLV

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc-marketing--1036.org.readthedocs.build/en/1036/

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

Attention: Patch coverage is 63.90135% with 161 lines in your changes missing coverage. Please review.

Project coverage is 92.48%. Comparing base (542a85b) to head (58c2b18).
Report is 102 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc_marketing/mmm/plot.py	13.76%	119 Missing ⚠️
pymc_marketing/mmm/multidimensional.py	86.36%	42 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1036      +/-   ##
==========================================
- Coverage   94.73%   92.48%   -2.25%     
==========================================
  Files          51       52       +1     
  Lines        5505     5947     +442     
==========================================
+ Hits         5215     5500     +285     
- Misses        290      447     +157     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[tim-mcwilliams]

@cetagostini I've been experimenting with this new feature and came across a potential bug. When trying to use partial pulling across a geo dim, I am running into a broadcasting error.

Looking like its throwing that error when creating the channel_contributions var, specifically in the forward_pass function. Right now the function dims are set to only look at the "channel"

return second.apply(x=first.apply(x=x, dims="channel"), dims="channel")

However, modifying that to include the dims being passed to the VanillaMultiDimensionalMMM class like so

return second.apply(x=first.apply(x=x, dims=(*self.dims,"channel")), dims=(*self.dims,"channel"))

fixes the broadcasting error and I was able to fit the model from there.

Here's the full traceback - @wd60622

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[150], line 1
----> 1 mmm_fit = mmm.fit(
      2     X=region_model_data.drop(columns="units"),
      3     y=region_model_data.drop(columns=Xs)
      4 )

File [~/code/pymc-marketing/docs/source/notebooks/mmm/MultiDimensionalMMM.py:548](http://localhost:8888/lab/tree/docs/source/notebooks/mmm/~/code/pymc-marketing/docs/source/notebooks/mmm/MultiDimensionalMMM.py#line=547), in VanillaMultiDimensionalMMM.fit(self, X, y, progressbar, predictor_names, random_seed, **kwargs)
    545     predictor_names = []
    547 if not hasattr(self, "model"):
--> 548     self.build_model(X, y)
    550 # sampler_kwargs = create_sample_kwargs(
    551 #     self.sampler_config,
    552 #     progressbar,
    553 #     random_seed,
    554 #     **kwargs,
    555 # )
    556 with self.model:

File [~/code/pymc-marketing/docs/source/notebooks/mmm/MultiDimensionalMMM.py:431](http://localhost:8888/lab/tree/docs/source/notebooks/mmm/~/code/pymc-marketing/docs/source/notebooks/mmm/MultiDimensionalMMM.py#line=430), in VanillaMultiDimensionalMMM.build_model(self, X, y, **kwargs)
    426     pass
    428 else:
    429     channel_contributions = pm.Deterministic(
    430         name="channel_contributions",
--> 431         var=self.forward_pass(x=channel_data_),
    432         dims=("date", *self.dims, "channel"),
    433     )
    435 mu_var = intercept + channel_contributions.sum(axis=-1)
    437 if (
    438     self.control_columns is not None
    439     and len(self.control_columns) > 0
    440     and all(column in X.columns for column in self.control_columns)
    441 ):

File [~/code/pymc-marketing/docs/source/notebooks/mmm/MultiDimensionalMMM.py:284](http://localhost:8888/lab/tree/docs/source/notebooks/mmm/~/code/pymc-marketing/docs/source/notebooks/mmm/MultiDimensionalMMM.py#line=283), in VanillaMultiDimensionalMMM.forward_pass(self, x)
    259 """Transform channel input into target contributions of each channel.
    260 
    261 This method handles the ordering of the adstock and saturation
   (...)
    276 
    277 """
    278 first, second = (
    279     (self.adstock, self.saturation)
    280     if self.adstock_first
    281     else (self.saturation, self.adstock)
    282 )
--> 284 return second.apply(x=first.apply(x=x, dims="channel"), dims="channel")

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/mmm/components/base.py:555](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/mmm/components/base.py#line=554), in Transformation.apply(self, x, dims)
    522 def apply(self, x: pt.TensorLike, dims: Dims | None = None) -> TensorVariable:
    523     """Call within a model context.
    524 
    525     Used internally of the MMM to apply the transformation to the data.
   (...)
    553 
    554     """
--> 555     kwargs = self._create_distributions(dims=dims)
    556     return self.function(x, **kwargs)

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/mmm/components/base.py:315](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/mmm/components/base.py#line=314), in Transformation._create_distributions(self, dims)
    311     var = dist.create_variable(variable_name)
    312     return dim_handler(var, dist.dims)
    314 return {
--> 315     parameter_name: create_variable(parameter_name, variable_name)
    316     for parameter_name, variable_name in self.variable_mapping.items()
    317 }

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/mmm/components/base.py:312](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/mmm/components/base.py#line=311), in Transformation._create_distributions.<locals>.create_variable(parameter_name, variable_name)
    310 dist = self.function_priors[parameter_name]
    311 var = dist.create_variable(variable_name)
--> 312 return dim_handler(var, dist.dims)

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/prior.py:192](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/prior.py#line=191), in create_dim_handler.<locals>.func(x, dims)
    191 def func(x: pt.TensorLike, dims: Dims) -> pt.TensorVariable:
--> 192     return handle_dims(x, dims, desired_dims)

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/prior.py:182](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pymc_marketing/prior.py#line=181), in handle_dims(x, dims, desired_dims)
    177 args = [
    178     "x" if missing else idx
    179     for (idx, missing) in zip(new_idx, missing_dims, strict=False)
    180 ]
    181 args = _remove_leading_xs(args)
--> 182 return x.dimshuffle(*args)

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pytensor/tensor/variable.py:347](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pytensor/tensor/variable.py#line=346), in _tensor_py_operators.dimshuffle(self, *pattern)
    345 if (len(pattern) == 1) and (isinstance(pattern[0], list | tuple)):
    346     pattern = pattern[0]
--> 347 op = pt.elemwise.DimShuffle(list(self.type.broadcastable), pattern)
    348 return op(self)

File [/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pytensor/tensor/elemwise.py:171](http://localhost:8888/opt/anaconda3/envs/marketing_env/lib/python3.12/site-packages/pytensor/tensor/elemwise.py#line=170), in DimShuffle.__init__(self, input_broadcastable, new_order)
    168             drop.append(i)
    169         else:
    170             # We cannot drop non-broadcastable dimensions
--> 171             raise ValueError(
    172                 "Cannot drop a non-broadcastable dimension: "
    173                 f"{input_broadcastable}, {new_order}"
    174             )
    176 # This is the list of the original dimensions that we keep
    177 self.shuffle = [x for x in new_order if x != "x"]

ValueError: Cannot drop a non-broadcastable dimension: [False, False], (0,)

[williambdean]

Good catch @tim-mcwilliams
So I am hearing you got it working with this fix, right?

@cetagostini You can do checks on the dims in the priors at initialization in order to catch errors earlier if needed

[tim-mcwilliams]

@wd60622 thanks! Correct, with the fix I was able to get the model working.

[luiztauffer]

hi guys, this looks really good and useful! Thanks for pointing me here, @cetagostini

I've read the example notebook and I have some questions:

• my impression is that what it is doing is basically training 4 separate models at the same time, is that correct?
• what would be the difference between this example using VanillaMultiDimensionalMMM and simply running 4 separate MMM?
• can VanillaMultiDimensionalMMM help making hierarchical assumptions? Let's say, we want the adstock_alpha prior to be modeled by a gaussian that's estimated iteratively for the group country

[cetagostini]

hi guys, this looks really good and useful! Thanks for pointing me here, @cetagostini

I've read the example notebook and I have some questions:

• my impression is that what it is doing is basically training 4 separate models at the same time, is that correct?
• what would be the difference between this example using VanillaMultiDimensionalMMM and simply running 4 separate MMM?
• can VanillaMultiDimensionalMMM help making hierarchical assumptions? Let's say, we want the adstock_alpha prior to be modeled by a gaussian that's estimated iteratively for the group country

Hey @luiztauffer Not exactly. Running 4 separate MMMs treats each dimension (e.g., country or channel) as independent, with no shared parameters or assumptions.

VanillaMultiDimensionalMMM (Currently only MMM class) uses a unified probabilistic framework where all dimensions (e.g., country, channel) are modeled together within a single Bayesian model. While it may seem like separate models are being trained because results are reported or analyzed per dimension, these dimensions are interdependent.

For example:

• Shared priors (like saturation_alpha, gamma_fourier) ensure that the model learns a common structure across dimensions. This means, that parameters for one dimension are informed by data from other dimensions.
• Any hierarchical structure or shared information (e.g., shared priors or parameters) allows for information pooling, improving generalization in smaller datasets. e.g: If one country has less data, the model can use the overall trend (from all countries) to regularize estimates.

So, while the output may look dimension-specific (e.g., country-level or channel-level effects), the model leverages cross-dimensional information during training 🙌🏻

[cetagostini]

@wd60622 feel free to jump if you want.

Play around today and get a head start on the API, for now we are able to properly train a mmm with all the features of the current MMM class. Basically we have TVPs and hierarchical structures, with the possibility of being multi-dimensional.

On the other hand, scaling occurs within the model, avoiding the need to scale inputs in plots or outputs in the model, everything is ready to use, reducing functions like sample_posterior_predictive and avoiding things like the optimizer needing to borrow scalers in the future.

• Work on basic plotsuite
• Integrate experiments.
• Clean API comments and tags
• Create test suite

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:25Z
----------------------------------------------------------------

Before me we merge we need to add a basic title and a short intro. I can refine and complete in a following pr

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:26Z
----------------------------------------------------------------

add text on the class type

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:27Z
----------------------------------------------------------------

Comment on the Prior structure

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:28Z
----------------------------------------------------------------

Explain how to specfy the model

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:29Z
----------------------------------------------------------------

Provide a summary on the model structure. In particular, which variables are hierarchical and so on ... Again, I can add more info on a subsequent pr

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:30Z
----------------------------------------------------------------

Explain this (do we need it?)

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:31Z
----------------------------------------------------------------

explain this plot

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:32Z
----------------------------------------------------------------

Explain all of these plots

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:35:32Z
----------------------------------------------------------------

explain the context of this model (and comments line at the end?)

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

juanitorduz commented on 2025-02-04T12:36:30Z
----------------------------------------------------------------

add watermark at the end

[juanitorduz]

@cetagostini @wd60622 this looks great! I left some minor comments on the code. Overal,l it looks almost good to go.

We need to add some basic text explanation to the notebook because we just have code. I can take over the notebook improvements in a following PR, but for this one we need the bare minimum ;)

This is going to be a game changer! Amazing job!

[williambdean]

All feedback should be included now @juanitorduz
Is there a DataFrame that you have in mind that would fit the outstanding comment? I can add a test case for it if it doesn't exist already

[juanitorduz]

All feedback should be included now @juanitorduz Is there a DataFrame that you have in mind that would fit the outstanding comment? I can add a test case for it if it doesn't exist already

Thank will! I usually do not work with pandas + indices wo maybe you know more when df.reset_index() could be a problem. Otherwise is all good ! Thanks!

[juanitorduz]

ahg, iths seems the precommit job is broken :(

[williambdean]

Thank will! I usually do not work with pandas + indices wo maybe you know more when df.reset_index() could be a problem. Otherwise is all good ! Thanks!

The API should work like before which didn't make use of an actual index / MultiIndex

[williambdean]

ahg, iths seems the precommit job is broken :(

They passed locally. Trust me