Allow VariableFactory in FourierBase (#1304)

* Allow VariableFactory in FourierBase

* change docstrings and loosen serialization

* tests for changes

---------

Co-authored-by: Will Dean <[email protected]>
Co-authored-by: Will Dean <[email protected]>

Author: aseyboldt

pymc_marketing/mmm/fourier.py

@@ -223,7 +223,7 @@
 
 from pymc_marketing.constants import DAYS_IN_MONTH, DAYS_IN_YEAR
 from pymc_marketing.plot import SelToString, plot_curve, plot_hdi, plot_samples
-from pymc_marketing.prior import Prior, create_dim_handler
+from pymc_marketing.prior import Prior, VariableFactory, create_dim_handler
 
 X_NAME: str = "day"
 NON_GRID_NAMES: frozenset[str] = frozenset({X_NAME})
@@ -274,8 +274,8 @@ class FourierBase(BaseModel):
     prefix : str, optional
         Alternative prefix for the fourier seasonality, by default None or
         "fourier"
-    prior : Prior, optional
-        Prior distribution for the fourier seasonality beta parameters, by
+    prior : Prior | VariableFactory, optional
+        Prior distribution or VariableFactory for the fourier seasonality beta parameters, by
         default `Prior("Laplace", mu=0, b=1)`
     variable_name : str, optional
         Name of the variable that multiplies the fourier modes. By default None,
@@ -286,17 +286,21 @@ class FourierBase(BaseModel):
     n_order: int = Field(..., gt=0)
     days_in_period: float = Field(..., gt=0)
     prefix: str = Field("fourier")
-    prior: InstanceOf[Prior] = Field(Prior("Laplace", mu=0, b=1))
+    prior: InstanceOf[Prior] | InstanceOf[VariableFactory] = Field(
+        Prior("Laplace", mu=0, b=1)
+    )
     variable_name: str | None = Field(None)
 
     def model_post_init(self, __context: Any) -> None:
         """Model post initialization for a Pydantic model."""
         if self.variable_name is None:
             self.variable_name = f"{self.prefix}_beta"
 
-        if not self.prior.dims:
+        if not self.prior.dims and isinstance(self.prior, Prior):
             self.prior = self.prior.deepcopy()
             self.prior.dims = self.prefix
+        elif not self.prior.dims:
+            self.prior.dims = self.prefix
 
     @model_validator(mode="after")
     def _check_variable_name(self) -> Self:
@@ -311,7 +315,7 @@ def _check_prior_has_right_dimensions(self) -> Self:
         return self
 
     @field_serializer("prior", when_used="json")
-    def serialize_prior(prior: Prior) -> dict[str, Any]:
+    def serialize_prior(prior: Any) -> dict[str, Any]:
         """Serialize the prior distribution.
 
         Parameters
@@ -325,6 +329,9 @@ def serialize_prior(prior: Prior) -> dict[str, Any]:
             The serialized prior distribution.
 
         """
+        if hasattr(prior, "to_dict"):
+            return prior.to_dict()
+
         return prior.to_json()
 
     @property
@@ -718,8 +725,8 @@ class YearlyFourier(FourierBase):
     prefix : str, optional
         Alternative prefix for the fourier seasonality, by default None or
         "fourier"
-    prior : Prior, optional
-        Prior distribution for the fourier seasonality beta parameters, by
+    prior : Prior | VariableFactory, optional
+        Prior distribution or VariableFactory for the fourier seasonality beta parameters, by
         default `Prior("Laplace", mu=0, b=1)`
     name : str, optional
         Name of the variable that multiplies the fourier modes, by default None
@@ -774,8 +781,8 @@ class MonthlyFourier(FourierBase):
     prefix : str, optional
         Alternative prefix for the fourier seasonality, by default None or
         "fourier"
-    prior : Prior, optional
-        Prior distribution for the fourier seasonality beta parameters, by
+    prior : Prior | VariableFactory, optional
+        Prior distribution or VariableFactory for the fourier seasonality beta parameters, by
         default `Prior("Laplace", mu=0, b=1)`
     name : str, optional
         Name of the variable that multiplies the fourier modes, by default None

tests/mmm/test_fourier.py

@@ -358,3 +358,46 @@ def test_fourier_base_instantiation():
             prior=Prior("Laplace", mu=0, b=1, dims="fourier"),
         )
     assert "Can't instantiate abstract class FourierBase" in str(exc_info.value)
+
+
+class ArbitraryCode:
+    def __init__(self, dims: tuple[str, ...]) -> None:
+        self.dims = dims
+
+    def create_variable(self, name: str):
+        return pm.Normal(name, dims=self.dims)
+
+
+def test_fourier_arbitrary_prior() -> None:
+    prior = ArbitraryCode(dims=("fourier",))
+    fourier = YearlyFourier(n_order=4, prior=prior)
+
+    x = np.arange(10)
+    with pm.Model():
+        y = fourier.apply(x)
+
+    assert y.eval().shape == (10,)
+
+
+def test_fourier_dims_modified() -> None:
+    prior = ArbitraryCode(dims=())
+    YearlyFourier(n_order=4, prior=prior)
+    assert prior.dims == "fourier"
+
+
+class SerializableArbitraryCode(ArbitraryCode):
+    def to_dict(self):
+        return {"dims": self.dims, "msg": "Hello, World!"}
+
+
+def test_fourier_serializable_arbitrary_prior() -> None:
+    prior = SerializableArbitraryCode(dims=("fourier",))
+    fourier = YearlyFourier(n_order=4, prior=prior)
+
+    assert fourier.model_dump(mode="json") == {
+        "n_order": 4,
+        "days_in_period": 365.25,
+        "prefix": "fourier",
+        "prior": {"dims": ["fourier"], "msg": "Hello, World!"},
+        "variable_name": "fourier_beta",
+    }