Allow Lists and NumPy Arrays as Transformation Priors (#1879)

* adds wider functionality for priors

* adding tests for this change

---------

Co-authored-by: Juan Orduz <[email protected]>

Author: timbo112711

pymc_marketing/mmm/components/base.py

@@ -24,7 +24,7 @@
 from collections.abc import Iterable
 from copy import deepcopy
 from inspect import signature
-from typing import Any
+from typing import Any, TypeAlias
 
 import numpy as np
 import numpy.typing as npt
@@ -49,8 +49,13 @@
 # "x" for saturation, "time since exposure" for adstock
 NON_GRID_NAMES: frozenset[str] = frozenset({"x", "time since exposure"})
 
-SupportedPrior = (
-    InstanceOf[Prior] | float | InstanceOf[TensorVariable] | InstanceOf[VariableFactory]
+SupportedPrior: TypeAlias = (
+    InstanceOf[Prior]
+    | float
+    | InstanceOf[TensorVariable]
+    | InstanceOf[VariableFactory]
+    | list
+    | InstanceOf[npt.NDArray[np.floating]]
 )
 
 
@@ -127,7 +132,7 @@ class Transformation:
 
     Parameters
     ----------
-    priors : dict[str, Prior | float | TensorVariable | VariableFactory], optional
+    priors : dict[str, Prior | float | TensorVariable | VariableFactory | list  | numpy array], optional
         Dictionary with the priors for the parameters of the function. The keys should be the
         parameter names and the values the priors. If not provided, it will use the default
         priors from the subclass.
@@ -144,8 +149,7 @@ class Transformation:
 
     def __init__(
         self,
-        priors: dict[str, Prior | float | TensorVariable | VariableFactory]
-        | None = None,
+        priors: dict[str, SupportedPrior] | None = None,
         prefix: str | None = None,
     ) -> None:
         self._checks()

tests/mmm/components/test_base.py

@@ -11,14 +11,16 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+from unittest.mock import Mock
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pymc as pm
 import pytensor.tensor as pt
 import pytest
 import xarray as xr
 from pymc_extras.prior import Prior, VariableFactory
-from pytensor.tensor import TensorVariable
+from pytensor.tensor import TensorVariable, scalar
 
 from pymc_marketing.mmm.components.base import (
     DuplicatedTransformationError,
@@ -616,3 +618,40 @@ def test_apply_idx_more_dims(new_transformation_class) -> None:
         Y.eval(),
         expected.eval(),
     )
+
+
+class DummyTransformation(Transformation):
+    @staticmethod
+    def function(data, x):
+        return data + x
+
+    prefix = "dummy"
+    lookup_name = "dummy"
+    default_priors = {"x": Prior("Normal", mu=0, sigma=1)}
+
+
+mock_variable_factory = Mock(spec=VariableFactory)
+
+
+@pytest.mark.parametrize(
+    "prior_value",
+    [
+        Prior("Normal", mu=0, sigma=1),
+        0.5,
+        scalar("x"),
+        mock_variable_factory,
+        [0.1, 0.2, 0.3],
+        np.array([0.1, 0.2, 0.3], dtype=float),
+    ],
+)
+def test_transformation_accepts_supported_priors(prior_value):
+    priors = {"x": prior_value}
+    tfm = DummyTransformation(priors=priors)
+
+    actual = tfm.function_priors["x"]
+
+    # Compare array-likes properly
+    if isinstance(prior_value, list | np.ndarray):
+        assert np.array_equal(actual, prior_value)
+    else:
+        assert actual == prior_value