Media transformation class and different transformations based on sub…

…sets of channels (pymc-labs#968)
radiokosmos · Sep 1, 2024 · e7cc642 · e7cc642
1 parent 9db32bf
commit e7cc642
Show file tree

Hide file tree

Showing 6 changed files with 628 additions and 0 deletions.
diff --git a/docs/source/uml/classes_mmm.png b/docs/source/uml/classes_mmm.png
diff --git a/docs/source/uml/packages_mmm.png b/docs/source/uml/packages_mmm.png
diff --git a/pymc_marketing/mmm/__init__.py b/pymc_marketing/mmm/__init__.py
@@ -38,6 +38,11 @@
     saturation_from_dict,
 )
 from pymc_marketing.mmm.fourier import MonthlyFourier, YearlyFourier
+from pymc_marketing.mmm.media_transformation import (
+    MediaConfig,
+    MediaConfigList,
+    MediaTransformation,
+)
 from pymc_marketing.mmm.mmm import MMM
 from pymc_marketing.mmm.preprocessing import (
     preprocessing_method_X,
@@ -46,6 +51,9 @@
 from pymc_marketing.mmm.validating import validation_method_X, validation_method_y
 
 __all__ = [
+    "MediaTransformation",
+    "MediaConfigList",
+    "MediaConfig",
     "AdstockTransformation",
     "BaseValidateMMM",
     "DelayedAdstock",

diff --git a/pymc_marketing/mmm/components/base.py b/pymc_marketing/mmm/components/base.py
@@ -128,6 +128,23 @@ def __repr__(self) -> str:
             ")"
         )
 
+    def set_dims_for_all_priors(self, dims: Dims):
+        """Convinience method to set the dims for all priors.
+
+        Parameters
+        ----------
+        dims : Dims
+            The dims for the priors.
+
+        Returns
+        -------
+        Transformation
+        """
+        for prior in self.function_priors.values():
+            prior.dims = dims
+
+        return self
+
     def to_dict(self) -> dict[str, Any]:
         """Convert the transformation to a dictionary.