Separate Weibull adstock into CDF & PDF

pymc_marketing/mmm/components/adstock.py

@@ -203,6 +203,88 @@
     }
 
 
+class WeibullPDFAdstock(AdstockTransformation):
+    """Wrapper around weibull adstock with PDF function.
+
+    For more information, see :func:`pymc_marketing.mmm.transformers.weibull_adstock`.
+
+    .. plot::
+        :context: close-figs
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        from pymc_marketing.mmm import WeibullPDFAdstock
+
+        rng = np.random.default_rng(0)
+
+        adstock = WeibullPDFAdstock(l_max=10)
+        prior = adstock.sample_prior(random_seed=rng)
+        curve = adstock.sample_curve(prior)
+        adstock.plot_curve(curve, sample_kwargs={"rng": rng})
+        plt.show()
+
+    """
+
+    lookup_name = "weibull_pdf"
+
+    def function(self, x, lam, k):
+        return weibull_adstock(
+            x=x,
+            lam=lam,
+            k=k,
+            l_max=self.l_max,
+            mode=self.mode,
+            type=WeibullType.PDF,
+            normalize=self.normalize,
+        )
+
+    default_priors = {
+        "lam": Prior("Gamma", mu=2, sigma=1),
+        "k": Prior("Gamma", mu=3, sigma=1),
+    }
+
+
+class WeibullCDFAdstock(AdstockTransformation):
+    """Wrapper around weibull adstock with CDF function.
+
+    For more information, see :func:`pymc_marketing.mmm.transformers.weibull_adstock`.
+
+    .. plot::
+        :context: close-figs
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        from pymc_marketing.mmm import WeibullCDFAdstock
+
+        rng = np.random.default_rng(0)
+
+        adstock = WeibullCDFAdstock(l_max=10)
+        prior = adstock.sample_prior(random_seed=rng)
+        curve = adstock.sample_curve(prior)
+        adstock.plot_curve(curve, sample_kwargs={"rng": rng})
+        plt.show()
+
+    """
+
+    lookup_name = "weibull_cdf"
+
+    def function(self, x, lam, k):
+        return weibull_adstock(
+            x=x,
+            lam=lam,
+            k=k,
+            l_max=self.l_max,
+            mode=self.mode,
+            type=WeibullType.CDF,
+            normalize=self.normalize,
+        )
+
+    default_priors = {
+        "lam": Prior("Gamma", mu=2, sigma=2.5),
+        "k": Prior("Gamma", mu=2, sigma=2.5),
+    }
+
+
 class WeibullAdstock(AdstockTransformation):
     """Wrapper around weibull adstock function.
 
@@ -242,6 +324,13 @@
             l_max=l_max, normalize=normalize, mode=mode, priors=priors, prefix=prefix
         )
 
+        msg = f"Use the Weibull{kind}Adstock class instead for better default priors."
+        warnings.warn(
+            msg,
+            UserWarning,
+            stacklevel=1,
+        )
+
     def function(self, x, lam, k):
         return weibull_adstock(
             x=x,
@@ -250,6 +339,7 @@
             l_max=self.l_max,
             mode=self.mode,
             type=self.kind,
+            normalize=self.normalize,
         )
 
     default_priors = {
@@ -260,7 +350,13 @@
 
 ADSTOCK_TRANSFORMATIONS: dict[str, type[AdstockTransformation]] = {
     cls.lookup_name: cls  # type: ignore
-    for cls in [GeometricAdstock, DelayedAdstock, WeibullAdstock]
+    for cls in [
+        GeometricAdstock,
+        DelayedAdstock,
+        WeibullAdstock,
+        WeibullPDFAdstock,
+        WeibullCDFAdstock,
+    ]
 }
 
 

pymc_marketing/mmm/transformers.py

@@ -327,6 +327,7 @@
     axis: int = 0,
     mode: ConvMode = ConvMode.After,
     type: WeibullType | str = WeibullType.PDF,
+    normalize: bool = False,
 ):
     R"""Weibull Adstocking Transformation.
 
@@ -404,6 +405,9 @@
             where the effect overlaps with both preceding and succeeding elements.
     type : WeibullType or str, by default WeibullType.PDF
         Type of Weibull adstock transformation to be applied (PDF or CDF).
+    normalize : bool, by default False
+        Whether to normalize the weights.
+
 
     Returns
     -------
@@ -427,6 +431,8 @@
         w = pt.cumprod(padded_w, axis=-1)
     else:
         raise ValueError(f"Wrong WeibullType: {type}, expected of WeibullType")
+
+    w = w / pt.sum(w, axis=-1, keepdims=True) if normalize else w
     return batched_convolution(x, w, axis=axis, mode=mode)
 
 

tests/mmm/components/test_adstock.py

@@ -22,6 +22,8 @@
     DelayedAdstock,
     GeometricAdstock,
     WeibullAdstock,
+    WeibullCDFAdstock,
+    WeibullPDFAdstock,
     _get_adstock_function,
 )
 
@@ -32,6 +34,8 @@ def adstocks() -> list[AdstockTransformation]:
         GeometricAdstock(l_max=10),
         WeibullAdstock(l_max=10, kind="PDF"),
         WeibullAdstock(l_max=10, kind="CDF"),
+        WeibullPDFAdstock(l_max=10),
+        WeibullCDFAdstock(l_max=10),
     ]