POC implement MultiDimensionalMMM with dimmed variables

Author: ricardoV94

pymc_marketing/mmm/components/adstock.py

@@ -56,12 +56,14 @@ def function(self, x, alpha):
 from __future__ import annotations
 
 import numpy as np
-import pytensor.tensor as pt
 import xarray as xr
 from pydantic import Field, validate_call
 from pymc_extras.deserialize import deserialize, register_deserialization
 from pymc_extras.prior import Prior
+from pytensor.tensor import TensorVariable
+from pytensor.xtensor import as_xtensor
 
+# from pymc.dims.distributions import Weibull  # TODO
 from pymc_marketing.mmm.components.base import (
     SupportedPrior,
     Transformation,
@@ -198,10 +200,15 @@ class BinomialAdstock(AdstockTransformation):
 
     lookup_name = "binomial"
 
-    def function(self, x, alpha):
+    def function(self, x, alpha, *, dim: str):
         """Binomial adstock function."""
         return binomial_adstock(
-            x, alpha=alpha, l_max=self.l_max, normalize=self.normalize, mode=self.mode
+            x,
+            alpha=alpha,
+            l_max=self.l_max,
+            normalize=self.normalize,
+            mode=self.mode,
+            dim=dim,
         )
 
     default_priors = {"alpha": Prior("Beta", alpha=1, beta=3)}
@@ -231,10 +238,15 @@ class GeometricAdstock(AdstockTransformation):
 
     lookup_name = "geometric"
 
-    def function(self, x, alpha):
+    def function(self, x, alpha, *, dim: str):
         """Geometric adstock function."""
         return geometric_adstock(
-            x, alpha=alpha, l_max=self.l_max, normalize=self.normalize, mode=self.mode
+            x,
+            alpha=alpha,
+            l_max=self.l_max,
+            normalize=self.normalize,
+            mode=self.mode,
+            dim=dim,
         )
 
     default_priors = {"alpha": Prior("Beta", alpha=1, beta=3)}
@@ -264,7 +276,7 @@ class DelayedAdstock(AdstockTransformation):
 
     lookup_name = "delayed"
 
-    def function(self, x, alpha, theta):
+    def function(self, x, alpha, theta, *, dim: str):
         """Delayed adstock function."""
         return delayed_adstock(
             x,
@@ -273,6 +285,7 @@ def function(self, x, alpha, theta):
             l_max=self.l_max,
             normalize=self.normalize,
             mode=self.mode,
+            dim=dim,
         )
 
     default_priors = {
@@ -305,7 +318,7 @@ class WeibullPDFAdstock(AdstockTransformation):
 
     lookup_name = "weibull_pdf"
 
-    def function(self, x, lam, k):
+    def function(self, x, lam, k, *, dim: str):
         """Weibull adstock function."""
         return weibull_adstock(
             x=x,
@@ -315,6 +328,7 @@ def function(self, x, lam, k):
             mode=self.mode,
             type=WeibullType.PDF,
             normalize=self.normalize,
+            dim=dim,
         )
 
     default_priors = {
@@ -347,7 +361,7 @@ class WeibullCDFAdstock(AdstockTransformation):
 
     lookup_name = "weibull_cdf"
 
-    def function(self, x, lam, k):
+    def function(self, x, lam, k, *, dim: str):
         """Weibull adstock function."""
         return weibull_adstock(
             x=x,
@@ -357,6 +371,7 @@ def function(self, x, lam, k):
             mode=self.mode,
             type=WeibullType.CDF,
             normalize=self.normalize,
+            dim=dim,
         )
 
     default_priors = {
@@ -370,9 +385,10 @@ class NoAdstock(AdstockTransformation):
 
     lookup_name: str = "no_adstock"
 
-    def function(self, x):
+    def function(self, x, *, dim: str | None = None) -> TensorVariable:
         """No adstock function."""
-        return pt.as_tensor_variable(x)
+        x = as_xtensor(x)
+        return x
 
     default_priors = {}
 

pymc_marketing/mmm/components/base.py

@@ -29,15 +29,18 @@
 import numpy as np
 import numpy.typing as npt
 import pymc as pm
+import pymc.dims as pmd
 import xarray as xr
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 from pydantic import InstanceOf
 from pymc.distributions.shape_utils import Dims
-from pymc_extras.prior import Prior, VariableFactory, create_dim_handler
+from pymc_extras.prior import Prior, VariableFactory
 from pytensor import tensor as pt
 from pytensor.tensor.variable import TensorVariable
+from pytensor.xtensor import as_xtensor
 
+from pymc_marketing.mmm.dims import XPrior
 from pymc_marketing.model_config import parse_model_config
 from pymc_marketing.plot import (
     SelToString,
@@ -220,13 +223,13 @@ def priors(self) -> dict[str, SupportedPrior]:
         return self.function_priors
 
     @function_priors.setter  # type: ignore
-    def function_priors(self, priors: dict[str, Any | Prior] | None) -> None:
+    def function_priors(self, priors: dict[str, Any | XPrior] | None) -> None:
         priors = priors or {}
 
         non_distributions = [
             key
             for key, value in priors.items()
-            if not isinstance(value, Prior) and not isinstance(value, dict)
+            if not isinstance(value, XPrior) and not isinstance(value, dict)
         ]
 
         priors = parse_model_config(priors, non_distributions=non_distributions)
@@ -313,9 +316,10 @@ def _has_defaults_for_all_arguments(self) -> None:
         function_signature = signature(self.function)
 
         # Remove the first one as assumed to be the data
+        # And the last one, as assumed to be dim
         parameters_that_need_priors = set(
             list(function_signature.parameters.keys())[1:]
-        )
+        ) - {"dim"}
         parameters_with_priors = set(self.default_priors.keys())
 
         missing_priors = parameters_that_need_priors - parameters_with_priors
@@ -379,8 +383,6 @@ def _create_distributions(
         if idx is not None:
             dims = ("N", *dims)
 
-        dim_handler = create_dim_handler(dims)
-
         def create_variable(parameter_name: str, variable_name: str) -> TensorVariable:
             dist = self.function_priors[parameter_name]
             if not hasattr(dist, "create_variable"):
@@ -392,10 +394,7 @@ def create_variable(parameter_name: str, variable_name: str) -> TensorVariable:
             if idx is not None and any(dim in idx for dim in dist_dims):
                 var = index_variable(var, dist.dims, idx)
 
-                dist_dims = [dim for dim in dist_dims if dim not in idx]
-                dist_dims = ("N", *dist_dims)
-
-            return dim_handler(var, dist_dims)
+            return var
 
         return {
             parameter_name: create_variable(parameter_name, variable_name)
@@ -521,16 +520,19 @@ def _sample_curve(
             }
         )
 
+        x = as_xtensor(x, dims=(x_dim,))
+
         with pm.Model(coords=coords):
-            pm.Deterministic(
+            pmd.Deterministic(
                 var_name,
-                self.apply(x, dims=output_core_dims),
+                self.apply(x, dims=output_core_dims, core_dim=x_dim),
                 dims=(x_dim, *output_core_dims),
             )
 
             return pm.sample_posterior_predictive(
                 parameters,
                 var_names=[var_name],
+                progressbar=False,
             ).posterior_predictive[var_name]
 
     def plot_curve_samples(
@@ -616,7 +618,9 @@ def plot_curve_hdi(
     def apply(
         self,
         x: pt.TensorLike,
+        *,
         dims: Dims | None = None,
+        core_dim: str,
         idx: dict[str, pt.TensorLike] | None = None,
     ) -> TensorVariable:
         """Call within a model context.
@@ -646,13 +650,18 @@ def apply(
 
             transformation = ...
 
-            coords = {"channel": ["TV", "Radio", "Digital"]}
+            coords = {
+                "channel": ["TV", "Radio", "Digital"],
+                "date": range(10),
+            }
             with pm.Model(coords=coords):
-                transformed_data = transformation.apply(data, dims="channel")
+                transformed_data = transformation.apply(
+                    data, dims="channel", core_dim="date"
+                )
 
         """
         kwargs = self._create_distributions(dims=dims, idx=idx)
-        return self.function(x, **kwargs)
+        return self.function(x, dim=core_dim, **kwargs)
 
 
 def _serialize_value(value: Any) -> Any: