Allow for indexing of apply parameters (#1847)

Author: williambdean

pymc_marketing/mmm/components/base.py

@@ -88,6 +88,27 @@ def __init__(self) -> None:
         super().__init__(msg)
 
 
+def index_variable(var, dims, idx) -> TensorVariable:
+    """Index a variable based on the provided dimensions and index.
+
+    Parameters
+    ----------
+    var : TensorVariable
+        The variable to index.
+    dims : tuple[str, ...]
+        The dims of the variable.
+    idx : dict[str, pt.TensorLike]
+        The index to use for the variable.
+
+    Returns
+    -------
+    TensorVariable
+        The indexed variable.
+
+    """
+    return var[tuple(idx[dim] if dim in idx else slice(None) for dim in dims)]
+
+
 class Transformation:
     """Base class for adstock and saturation functions.
 
@@ -338,17 +359,42 @@ def _infer_output_core_dims(self) -> tuple[str, ...]:
         return tuple(list({str(dim): None for dims in parameter_dims for dim in dims}))
 
     def _create_distributions(
-        self, dims: Dims | None = None
+        self,
+        dims: Dims | None = None,
+        idx: dict[str, pt.TensorLike] | None = None,
     ) -> dict[str, TensorVariable]:
-        dim_handler = create_dim_handler(dims or self._infer_output_core_dims())
+        if isinstance(dims, str):
+            dims = (dims,)
+
+        dims = dims or self.combined_dims
+        if idx is not None:
+            n_idx_dims = len(idx)
+            dummy_dims = tuple(f"DUMMY_{i}" for i in range(n_idx_dims))
+            if len(dummy_dims) > 1:
+                raise NotImplementedError(
+                    "The indexing with multiple dimensions is not supported yet."
+                )
+
+            dims = (*dummy_dims, *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"):
                 return dist
 
             var = dist.create_variable(variable_name)
-            return dim_handler(var, dist.dims)
+
+            dist_dims = dist.dims
+            if idx is not None:
+                var = index_variable(var, dist.dims, idx)
+
+                dist_dims = tuple(
+                    [(dim if dim not in idx else "DUMMY_0") for dim in dist.dims]
+                )
+
+            return dim_handler(var, dist_dims)
 
         return {
             parameter_name: create_variable(parameter_name, variable_name)
@@ -566,7 +612,12 @@ def plot_curve_hdi(
             hdi_kwargs=hdi_kwargs,
         )
 
-    def apply(self, x: pt.TensorLike, dims: Dims | None = None) -> TensorVariable:
+    def apply(
+        self,
+        x: pt.TensorLike,
+        dims: Dims | None = None,
+        idx: dict[str, pt.TensorLike] | None = None,
+    ) -> TensorVariable:
         """Call within a model context.
 
         Used internally of the MMM to apply the transformation to the data.
@@ -599,7 +650,7 @@ def apply(self, x: pt.TensorLike, dims: Dims | None = None) -> TensorVariable:
                 transformed_data = transformation.apply(data, dims="channel")
 
         """
-        kwargs = self._create_distributions(dims=dims)
+        kwargs = self._create_distributions(dims=dims, idx=idx)
         return self.function(x, **kwargs)
 
 

tests/mmm/components/test_base.py

@@ -26,6 +26,7 @@
     ParameterPriorException,
     Transformation,
     create_registration_meta,
+    index_variable,
 )
 from pymc_marketing.mmm.components.saturation import TanhSaturation
 
@@ -477,3 +478,99 @@ def create_variable(self, name: str):
 
     curve = saturation.sample_curve(prior, 10)
     assert curve.dims == ("chain", "draw", "x", "dim_a")
+
+
+@pytest.mark.parametrize(
+    "var, dims, idx, expected",
+    [
+        (
+            np.array([[1, 2, 3], [4, 5, 6]]),
+            ("geo", "channel"),
+            {"geo": [0, 0, 1, 1]},
+            np.array(
+                [
+                    [1, 2, 3],
+                    [1, 2, 3],
+                    [4, 5, 6],
+                    [4, 5, 6],
+                ]
+            ),
+        ),
+    ],
+)
+def test_index_variable(var, dims, idx, expected) -> None:
+    result = index_variable(var, dims=dims, idx=idx)
+    if isinstance(result, TensorVariable):
+        result = result.eval()
+
+    np.testing.assert_allclose(result, expected)
+
+
+def test_apply_idx(new_transformation_class) -> None:
+    instance = new_transformation_class(
+        priors={
+            "a": Prior(
+                "HalfNormal",
+                dims="geo",
+            ),
+            "b": Prior(
+                "HalfNormal",
+                dims="channel",
+            ),
+        }
+    )
+
+    X = np.array(
+        [
+            [0, 0, 0],
+            [1, 1, 1],
+            [2, 2, 2],
+            [0, 0, 0],
+            [1, 1, 1],
+            [2, 2, 2],
+        ]
+    )
+
+    coords = {"geo": ["A", "B"], "channel": ["TV", "Radio", "Online"]}
+    with pm.Model(coords=coords) as model:
+        idx = [0, 0, 0, 1, 1, 1]
+        Y = instance.apply(X, idx={"geo": idx}, dims="channel")
+
+        expected = instance.function(
+            X,
+            a=model["new_a"][idx, None],
+            b=model["new_b"],
+        )
+
+    np.testing.assert_allclose(
+        Y.eval(),
+        expected.eval(),
+    )
+
+
+def test_apply_index_too_many(new_transformation_class) -> None:
+    instance = new_transformation_class(
+        priors={
+            "a": Prior(
+                "HalfNormal",
+                dims=("geo", "product"),
+            ),
+            "b": Prior(
+                "HalfNormal",
+                dims="channel",
+            ),
+        }
+    )
+
+    coords = {
+        "geo": ["A", "B"],
+        "product": ["X", "Y", "Z"],
+        "channel": ["TV", "Radio", "Online"],
+    }
+    with pm.Model(coords=coords):
+        idx = {
+            "geo": [0, 0, 0, 1, 1, 1],
+            "product": [0, 1, 2, 0, 1, 2],
+        }
+        with pytest.raises(NotImplementedError, match="The indexing"):
+            instance.apply(None, idx=idx, dims="channel")