support case with additional dimension (#1857)

Author: williambdean

pymc_marketing/mmm/components/base.py

@@ -368,14 +368,7 @@ def _create_distributions(
 
         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)
+            dims = ("N", *dims)
 
         dim_handler = create_dim_handler(dims)
 
@@ -387,12 +380,11 @@ def create_variable(parameter_name: str, variable_name: str) -> TensorVariable:
             var = dist.create_variable(variable_name)
 
             dist_dims = dist.dims
-            if idx is not None:
+            if idx is not None and any(dim in idx for dim in dist_dims):
                 var = index_variable(var, dist.dims, idx)
 
-                dist_dims = tuple(
-                    [(dim if dim not in idx else "DUMMY_0") for dim in dist.dims]
-                )
+                dist_dims = [dim for dim in dist_dims if dim not in idx]
+                dist_dims = ("N", *dist_dims)
 
             return dim_handler(var, dist_dims)
 

tests/mmm/components/test_base.py

@@ -483,6 +483,22 @@ def create_variable(self, name: str):
 @pytest.mark.parametrize(
     "var, dims, idx, expected",
     [
+        (
+            np.arange(2 * 3 * 4).reshape(2, 3, 4),
+            ("geo", "product", "channel"),
+            {
+                "geo": [0, 0, 1, 1],
+                "product": [0, 2, 1, 0],
+            },
+            np.array(
+                [
+                    [0, 1, 2, 3],
+                    [8, 9, 10, 11],
+                    [16, 17, 18, 19],
+                    [12, 13, 14, 15],
+                ]
+            ),
+        ),
         (
             np.array([[1, 2, 3], [4, 5, 6]]),
             ("geo", "channel"),
@@ -548,7 +564,7 @@ def test_apply_idx(new_transformation_class) -> None:
     )
 
 
-def test_apply_index_too_many(new_transformation_class) -> None:
+def test_apply_idx_more_dims(new_transformation_class) -> None:
     instance = new_transformation_class(
         priors={
             "a": Prior(
@@ -557,20 +573,46 @@ def test_apply_index_too_many(new_transformation_class) -> None:
             ),
             "b": Prior(
                 "HalfNormal",
-                dims="channel",
+                dims=("product", "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"],
         "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")
+    with pm.Model(coords=coords) as model:
+        geo_idx = [0, 0, 0, 1, 1, 1]
+        product_idx = [0, 2, 1, 0, 1, 0]
+        Y = instance.apply(
+            X,
+            idx={
+                "geo": geo_idx,
+                "product": product_idx,
+            },
+            dims="channel",
+        )
+
+        expected = instance.function(
+            X,
+            a=model["new_a"][geo_idx, product_idx, None],
+            b=model["new_b"][product_idx],
+        )
+
+    np.testing.assert_allclose(
+        Y.eval(),
+        expected.eval(),
+    )