Repository: pymc-labs/pymc-marketing
Pull request #1675: Change the event function to class
Author: login: williambdean
Labels: name: MMM
name: tests
Created: 2025-05-07T16:47:55Z | Merged: 2025-05-08T02:51:45Z

<!-- !! Thank your for opening a PR !! -->

<!--- Provide a self-contained summary of your changes in the Title above -->
<!--- This is what will be shown in the automatic release notes: https://github.com/pymc-labs/pymc-marketing/releases -->

## Description
<!--- Describe your changes in detail -->

The other additive effects are represented with class. This migrates the function to a class.

## Related Issue
<!--- It is good practice to first open an issue explaining the bug / new feature that is addressed by this PR -->
<!--- Please type an `x` in one of the boxes below and provide the issue number after the # sign: -->
- [ ] Closes #
- [ ] Related to #

## Checklist
<!--- Make sure you have completed the following steps before submitting your PR -->
<!--- Feel free to type an `x` in all the boxes below to let us know you have completed the steps: -->
- [ ] Checked that [the pre-commit linting/style checks pass](https://www.pymc-marketing.io/en/latest/contributing/index.html). Feel free to comment [`pre-commit.ci autofix` to auto-fix](https://pre-commit.ci/#configuration-autofix_prs).
- [ ] Included tests that prove the fix is effective or that the new feature works
- [ ] Added necessary documentation (docstrings and/or example notebooks) using [numpydoc format](https://numpydoc.readthedocs.io/en/latest/format.html).
- [ ] If you are a pro: each commit corresponds to a [relevant logical change](https://wiki.openstack.org/wiki/GitCommitMessages#Structural_split_of_changes)
<!--- You may find this guide helpful: https://mainmatter.com/blog/2021/05/26/keeping-a-clean-git-history/ -->


<!-- readthedocs-preview pymc-marketing start -->
----
📚 Documentation preview 📚: https://pymc-marketing--1675.org.readthedocs.build/en/1675/

<!-- readthedocs-preview pymc-marketing end -->


Diff:
diff --git a/pymc_marketing/mmm/additive_effect.py b/pymc_marketing/mmm/additive_effect.py
index 8f97e19b6..d2fb1ef92 100644
--- a/pymc_marketing/mmm/additive_effect.py
+++ b/pymc_marketing/mmm/additive_effect.py
@@ -13,11 +13,12 @@
 #   limitations under the License.
 """Additive effects for the multidimensional Marketing Mix Model."""
 
-from typing import Protocol
+from typing import Any, Protocol
 
 import pandas as pd
 import pymc as pm
 import xarray as xr
+from pydantic import BaseModel, InstanceOf
 from pytensor import tensor as pt
 
 from pymc_marketing.mmm.events import EventEffect, days_from_reference
@@ -318,14 +319,8 @@ def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
         pm.set_data({f"{self.prefix}_t": t}, model=model)
 
 
-def create_event_mu_effect(
-    df_events: pd.DataFrame,
-    prefix: str,
-    effect: EventEffect,
-) -> MuEffect:
-    """Create an event effect for the MMM.
-
-    This class has the ability to create data and mean effects for the MMM model.
+class EventAdditiveEffect(BaseModel):
+    """Event effect class for the MMM.
 
     Parameters
     ----------
@@ -338,105 +333,112 @@ def create_event_mu_effect(
         The prefix to use for the event effect and associated variables.
     effect : EventEffect
         The event effect to apply.
-
-    Returns
-    -------
-    MuEffect
-        The event effect which is used in the MMM.
+    reference_date : str
+        The arbitrary reference date to calculate distance from events in days. Default
+        is "2025-01-01".
 
     """
-    if missing_columns := set(["start_date", "end_date", "name"]).difference(
-        df_events.columns,
-    ):
-        raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    effect.basis.prefix = prefix
+    df_events: InstanceOf[pd.DataFrame]
+    prefix: str
+    effect: EventEffect
+    reference_date: str = "2025-01-01"
 
-    reference_date = "2025-01-01"
-    start_dates = pd.to_datetime(df_events["start_date"])
-    end_dates = pd.to_datetime(df_events["end_date"])
+    def model_post_init(self, context: Any, /) -> None:
+        """Post initialization of the model."""
+        if missing_columns := set(["start_date", "end_date", "name"]).difference(
+            self.df_events.columns
+        ):
+            raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    class Effect:
-        """Event effect class for the MMM."""
+        self.effect.basis.prefix = self.prefix
 
-        def create_data(self, mmm: MMM) -> None:
-            """Create the required data in the model.
+    @property
+    def start_dates(self) -> pd.Series:
+        """The start dates of the events."""
+        return pd.to_datetime(self.df_events["start_date"])
 
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+    @property
+    def end_dates(self) -> pd.Series:
+        """The end dates of the events."""
+        return pd.to_datetime(self.df_events["end_date"])
 
-            """
-            model: pm.Model = mmm.model
+    def create_data(self, mmm: MMM) -> None:
+        """Create the required data in the model.
 
-            model_dates = pd.to_datetime(model.coords["date"])
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            model.add_coord(prefix, df_events["name"].to_numpy())
+        """
+        model: pm.Model = mmm.model
 
-            if "days" not in model:
-                pm.Data(
-                    "days",
-                    days_from_reference(model_dates, reference_date),
-                    dims="date",
-                )
+        model_dates = pd.to_datetime(model.coords["date"])
 
+        model.add_coord(self.prefix, self.df_events["name"].to_numpy())
+
+        if "days" not in model:
             pm.Data(
-                f"{prefix}_start_diff",
-                days_from_reference(start_dates, reference_date),
-                dims=prefix,
-            )
-            pm.Data(
-                f"{prefix}_end_diff",
-                days_from_reference(end_dates, reference_date),
-                dims=prefix,
+                "days",
+                days_from_reference(model_dates, self.reference_date),
+                dims="date",
             )
 
-        def create_effect(self, mmm: MMM) -> pt.TensorVariable:
-            """Create the event effect in the model.
-
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+        pm.Data(
+            f"{self.prefix}_start_diff",
+            days_from_reference(self.start_dates, self.reference_date),
+            dims=self.prefix,
+        )
+        pm.Data(
+            f"{self.prefix}_end_diff",
+            days_from_reference(self.end_dates, self.reference_date),
+            dims=self.prefix,
+        )
 
-            Returns
-            -------
-            pt.TensorVariable
-                The average event effect in the model.
+    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+        """Create the event effect in the model.
 
-            """
-            model: pm.Model = mmm.model
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            s_ref = model["days"][:, None] - model[f"{prefix}_start_diff"]
-            e_ref = model["days"][:, None] - model[f"{prefix}_end_diff"]
+        Returns
+        -------
+        pt.TensorVariable
+            The average event effect in the model.
 
-            def create_basis_matrix(s_ref, e_ref):
-                return pt.where(
-                    (s_ref >= 0) & (e_ref <= 0),
-                    0,
-                    pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
-                )
+        """
+        model: pm.Model = mmm.model
 
-            X = create_basis_matrix(s_ref, e_ref)
-            event_effect = effect.apply(X, name=prefix)
+        start_ref = model["days"][:, None] - model[f"{self.prefix}_start_diff"]
+        end_ref = model["days"][:, None] - model[f"{self.prefix}_end_diff"]
 
-            total_effect = pm.Deterministic(
-                f"{prefix}_total_effect",
-                event_effect.sum(axis=1),
-                dims="date",
+        def create_basis_matrix(start_ref, end_ref):
+            return pt.where(
+                (start_ref >= 0) & (end_ref <= 0),
+                0,
+                pt.where(pt.abs(start_ref) < pt.abs(end_ref), start_ref, end_ref),
             )
 
-            dim_handler = create_dim_handler(("date", *mmm.dims))
-            return dim_handler(total_effect, "date")
+        X = create_basis_matrix(start_ref, end_ref)
+        event_effect = self.effect.apply(X, name=self.prefix)
 
-        def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
-            """Set the data for new predictions."""
-            new_dates = pd.to_datetime(model.coords["date"])
+        total_effect = pm.Deterministic(
+            f"{self.prefix}_total_effect",
+            event_effect.sum(axis=1),
+            dims="date",
+        )
 
-            new_data = {
-                "days": days_from_reference(new_dates, reference_date),
-            }
-            pm.set_data(new_data=new_data, model=model)
+        dim_handler = create_dim_handler(("date", *mmm.dims))
+        return dim_handler(total_effect, "date")
+
+    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
+        """Set the data for new predictions."""
+        new_dates = pd.to_datetime(model.coords["date"])
 
-    return Effect()
+        new_data = {
+            "days": days_from_reference(new_dates, self.reference_date),
+        }
+        pm.set_data(new_data=new_data, model=model)
diff --git a/pymc_marketing/mmm/events.py b/pymc_marketing/mmm/events.py
index 934bc41c0..40df3386a 100644
--- a/pymc_marketing/mmm/events.py
+++ b/pymc_marketing/mmm/events.py
@@ -53,13 +53,13 @@ def create_basis_matrix(df_events: pd.DataFrame, model_dates: np.ndarray):
         start_dates = df_events["start_date"]
         end_dates = df_events["end_date"]
 
-        s_ref = difference_in_days(model_dates, start_dates)
-        e_ref = difference_in_days(model_dates, end_dates)
+        start_ref = difference_in_days(model_dates, start_dates)
+        end_ref = difference_in_days(model_dates, end_dates)
 
         return np.where(
-            (s_ref >= 0) & (e_ref <= 0),
+            (start_ref >= 0) & (end_ref <= 0),
             0,
-            np.where(np.abs(s_ref) < np.abs(e_ref), s_ref, e_ref),
+            np.where(np.abs(start_ref) < np.abs(end_ref), start_ref, end_ref),
         )
 
 
diff --git a/pymc_marketing/mmm/multidimensional.py b/pymc_marketing/mmm/multidimensional.py
index 28e174a75..56e9729cb 100644
--- a/pymc_marketing/mmm/multidimensional.py
+++ b/pymc_marketing/mmm/multidimensional.py
@@ -33,7 +33,7 @@
 from scipy.optimize import OptimizeResult
 
 from pymc_marketing.mmm import SoftPlusHSGP
-from pymc_marketing.mmm.additive_effect import MuEffect, create_event_mu_effect
+from pymc_marketing.mmm.additive_effect import EventAdditiveEffect, MuEffect
 from pymc_marketing.mmm.budget_optimizer import OptimizerCompatibleModelWrapper
 from pymc_marketing.mmm.components.adstock import (
     AdstockTransformation,
@@ -247,7 +247,11 @@ def add_events(
                 f"Event effect dims {effect.dims} must contain {prefix} and {self.dims}"
             )
 
-        event_effect = create_event_mu_effect(df_events, prefix, effect)
+        event_effect = EventAdditiveEffect(
+            df_events=df_events,
+            prefix=prefix,
+            effect=effect,
+        )
         self.mu_effects.append(event_effect)
 
     @property
diff --git a/tests/mmm/test_additive_effect.py b/tests/mmm/test_additive_effect.py
index 42530ff46..9767ef90a 100644
--- a/tests/mmm/test_additive_effect.py
+++ b/tests/mmm/test_additive_effect.py
@@ -16,7 +16,10 @@
 import pymc as pm
 import pytest
 
-from pymc_marketing.mmm.additive_effect import FourierEffect, LinearTrendEffect
+from pymc_marketing.mmm.additive_effect import (
+    FourierEffect,
+    LinearTrendEffect,
+)
 from pymc_marketing.mmm.fourier import MonthlyFourier, WeeklyFourier, YearlyFourier
 from pymc_marketing.mmm.linear_trend import LinearTrend
 from pymc_marketing.prior import Prior
diff --git a/tests/mmm/test_multidimensional.py b/tests/mmm/test_multidimensional.py
index 548d110cc..5309c9da7 100644
--- a/tests/mmm/test_multidimensional.py
+++ b/tests/mmm/test_multidimensional.py
@@ -24,11 +24,11 @@
 from scipy.optimize import OptimizeResult
 
 from pymc_marketing.mmm import GeometricAdstock, LogisticSaturation
+from pymc_marketing.mmm.additive_effect import EventAdditiveEffect
 from pymc_marketing.mmm.events import EventEffect, GaussianBasis
 from pymc_marketing.mmm.multidimensional import (
     MMM,
     MultiDimensionalBudgetOptimizerWrapper,
-    create_event_mu_effect,
 )
 from pymc_marketing.mmm.scaling import Scaling, VariableScaling
 from pymc_marketing.prior import Prior
@@ -469,7 +469,11 @@ def test_create_effect_mu_effect(
     df_events,
     event_effect,
 ) -> None:
-    effect = create_event_mu_effect(df_events, prefix="holiday", effect=event_effect)
+    effect = EventAdditiveEffect(
+        df_events=df_events,
+        prefix="holiday",
+        effect=event_effect,
+    )
 
     with mock_mmm.model:
         effect.create_data(mock_mmm)


Reviews:
user: login: williambdean
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T16:49:07Z
user: login: cetagostini
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:02:34Z
user: login: PabloRoque
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:33:12Z
user: login: PabloRoque
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:33:14Z
user: login: PabloRoque
body: Neat!

One minor pedantic comment on var names, and one around `def set_data`
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:39:23Z
user: login: williambdean
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:59:40Z
user: login: williambdean
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T00:27:25Z

Review comments:
diff_hunk: @@ -318,125 +319,109 @@ def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
         pm.set_data({f"{self.prefix}_t": t}, model=model)
 
 
-def create_event_mu_effect(
-    df_events: pd.DataFrame,
-    prefix: str,
-    effect: EventEffect,
-) -> MuEffect:
-    """Create an event effect for the MMM.
+class EventAdditiveEffect(BaseModel):
path: pymc_marketing/mmm/additive_effect.py
user: login: williambdean
body: Unfortunatly, `EventEffect` is class name already ...
created_at: 2025-05-07T16:49:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -318,125 +319,109 @@ def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
         pm.set_data({f"{self.prefix}_t": t}, model=model)
 
 
-def create_event_mu_effect(
-    df_events: pd.DataFrame,
-    prefix: str,
-    effect: EventEffect,
-) -> MuEffect:
-    """Create an event effect for the MMM.
+class EventAdditiveEffect(BaseModel):
path: pymc_marketing/mmm/additive_effect.py
user: login: cetagostini
body: This is fine, once we move to the new class we can remove the previous class and keep this one only with the short name!
created_at: 2025-05-07T17:02:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -338,105 +333,112 @@ def create_event_mu_effect(
         The prefix to use for the event effect and associated variables.
     effect : EventEffect
         The event effect to apply.
-
-    Returns
-    -------
-    MuEffect
-        The event effect which is used in the MMM.
+    reference_date : str
+        The arbitrary reference date to calculate distance from events in days. Default
+        is "2025-01-01".
 
     """
-    if missing_columns := set(["start_date", "end_date", "name"]).difference(
-        df_events.columns,
-    ):
-        raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    effect.basis.prefix = prefix
+    df_events: InstanceOf[pd.DataFrame]
+    prefix: str
+    effect: EventEffect
+    reference_date: str = "2025-01-01"
 
-    reference_date = "2025-01-01"
-    start_dates = pd.to_datetime(df_events["start_date"])
-    end_dates = pd.to_datetime(df_events["end_date"])
+    def model_post_init(self, context: Any, /) -> None:
+        """Post initialization of the model."""
+        if missing_columns := set(["start_date", "end_date", "name"]).difference(
+            self.df_events.columns
+        ):
+            raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    class Effect:
-        """Event effect class for the MMM."""
+        self.effect.basis.prefix = self.prefix
 
-        def create_data(self, mmm: MMM) -> None:
-            """Create the required data in the model.
+    @property
+    def start_dates(self) -> pd.Series:
+        """The start dates of the events."""
+        return pd.to_datetime(self.df_events["start_date"])
 
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+    @property
+    def end_dates(self) -> pd.Series:
+        """The end dates of the events."""
+        return pd.to_datetime(self.df_events["end_date"])
 
-            """
-            model: pm.Model = mmm.model
+    def create_data(self, mmm: MMM) -> None:
+        """Create the required data in the model.
 
-            model_dates = pd.to_datetime(model.coords["date"])
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            model.add_coord(prefix, df_events["name"].to_numpy())
+        """
+        model: pm.Model = mmm.model
 
-            if "days" not in model:
-                pm.Data(
-                    "days",
-                    days_from_reference(model_dates, reference_date),
-                    dims="date",
-                )
+        model_dates = pd.to_datetime(model.coords["date"])
 
+        model.add_coord(self.prefix, self.df_events["name"].to_numpy())
+
+        if "days" not in model:
             pm.Data(
-                f"{prefix}_start_diff",
-                days_from_reference(start_dates, reference_date),
-                dims=prefix,
-            )
-            pm.Data(
-                f"{prefix}_end_diff",
-                days_from_reference(end_dates, reference_date),
-                dims=prefix,
+                "days",
+                days_from_reference(model_dates, self.reference_date),
+                dims="date",
             )
 
-        def create_effect(self, mmm: MMM) -> pt.TensorVariable:
-            """Create the event effect in the model.
-
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+        pm.Data(
+            f"{self.prefix}_start_diff",
+            days_from_reference(self.start_dates, self.reference_date),
+            dims=self.prefix,
+        )
+        pm.Data(
+            f"{self.prefix}_end_diff",
+            days_from_reference(self.end_dates, self.reference_date),
+            dims=self.prefix,
+        )
 
-            Returns
-            -------
-            pt.TensorVariable
-                The average event effect in the model.
+    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+        """Create the event effect in the model.
 
-            """
-            model: pm.Model = mmm.model
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            s_ref = model["days"][:, None] - model[f"{prefix}_start_diff"]
-            e_ref = model["days"][:, None] - model[f"{prefix}_end_diff"]
+        Returns
+        -------
+        pt.TensorVariable
+            The average event effect in the model.
 
-            def create_basis_matrix(s_ref, e_ref):
-                return pt.where(
-                    (s_ref >= 0) & (e_ref <= 0),
-                    0,
-                    pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
-                )
+        """
+        model: pm.Model = mmm.model
 
-            X = create_basis_matrix(s_ref, e_ref)
-            event_effect = effect.apply(X, name=prefix)
+        s_ref = model["days"][:, None] - model[f"{self.prefix}_start_diff"]
+        e_ref = model["days"][:, None] - model[f"{self.prefix}_end_diff"]
path: pymc_marketing/mmm/additive_effect.py
user: login: PabloRoque
body: perhaps the more explicit `start_ref`, `end_ref`?
created_at: 2025-05-07T17:33:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -338,105 +333,112 @@ def create_event_mu_effect(
         The prefix to use for the event effect and associated variables.
     effect : EventEffect
         The event effect to apply.
-
-    Returns
-    -------
-    MuEffect
-        The event effect which is used in the MMM.
+    reference_date : str
+        The arbitrary reference date to calculate distance from events in days. Default
+        is "2025-01-01".
 
     """
-    if missing_columns := set(["start_date", "end_date", "name"]).difference(
-        df_events.columns,
-    ):
-        raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    effect.basis.prefix = prefix
+    df_events: InstanceOf[pd.DataFrame]
+    prefix: str
+    effect: EventEffect
+    reference_date: str = "2025-01-01"
 
-    reference_date = "2025-01-01"
-    start_dates = pd.to_datetime(df_events["start_date"])
-    end_dates = pd.to_datetime(df_events["end_date"])
+    def model_post_init(self, context: Any, /) -> None:
+        """Post initialization of the model."""
+        if missing_columns := set(["start_date", "end_date", "name"]).difference(
+            self.df_events.columns
+        ):
+            raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    class Effect:
-        """Event effect class for the MMM."""
+        self.effect.basis.prefix = self.prefix
 
-        def create_data(self, mmm: MMM) -> None:
-            """Create the required data in the model.
+    @property
+    def start_dates(self) -> pd.Series:
+        """The start dates of the events."""
+        return pd.to_datetime(self.df_events["start_date"])
 
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+    @property
+    def end_dates(self) -> pd.Series:
+        """The end dates of the events."""
+        return pd.to_datetime(self.df_events["end_date"])
 
-            """
-            model: pm.Model = mmm.model
+    def create_data(self, mmm: MMM) -> None:
+        """Create the required data in the model.
 
-            model_dates = pd.to_datetime(model.coords["date"])
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            model.add_coord(prefix, df_events["name"].to_numpy())
+        """
+        model: pm.Model = mmm.model
 
-            if "days" not in model:
-                pm.Data(
-                    "days",
-                    days_from_reference(model_dates, reference_date),
-                    dims="date",
-                )
+        model_dates = pd.to_datetime(model.coords["date"])
 
+        model.add_coord(self.prefix, self.df_events["name"].to_numpy())
+
+        if "days" not in model:
             pm.Data(
-                f"{prefix}_start_diff",
-                days_from_reference(start_dates, reference_date),
-                dims=prefix,
-            )
-            pm.Data(
-                f"{prefix}_end_diff",
-                days_from_reference(end_dates, reference_date),
-                dims=prefix,
+                "days",
+                days_from_reference(model_dates, self.reference_date),
+                dims="date",
             )
 
-        def create_effect(self, mmm: MMM) -> pt.TensorVariable:
-            """Create the event effect in the model.
-
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+        pm.Data(
+            f"{self.prefix}_start_diff",
+            days_from_reference(self.start_dates, self.reference_date),
+            dims=self.prefix,
+        )
+        pm.Data(
+            f"{self.prefix}_end_diff",
+            days_from_reference(self.end_dates, self.reference_date),
+            dims=self.prefix,
+        )
 
-            Returns
-            -------
-            pt.TensorVariable
-                The average event effect in the model.
+    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+        """Create the event effect in the model.
 
-            """
-            model: pm.Model = mmm.model
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            s_ref = model["days"][:, None] - model[f"{prefix}_start_diff"]
-            e_ref = model["days"][:, None] - model[f"{prefix}_end_diff"]
+        Returns
+        -------
+        pt.TensorVariable
+            The average event effect in the model.
 
-            def create_basis_matrix(s_ref, e_ref):
-                return pt.where(
-                    (s_ref >= 0) & (e_ref <= 0),
-                    0,
-                    pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
-                )
+        """
+        model: pm.Model = mmm.model
 
-            X = create_basis_matrix(s_ref, e_ref)
-            event_effect = effect.apply(X, name=prefix)
+        s_ref = model["days"][:, None] - model[f"{self.prefix}_start_diff"]
+        e_ref = model["days"][:, None] - model[f"{self.prefix}_end_diff"]
 
-            total_effect = pm.Deterministic(
-                f"{prefix}_total_effect",
-                event_effect.sum(axis=1),
-                dims="date",
+        def create_basis_matrix(s_ref, e_ref):
path: pymc_marketing/mmm/additive_effect.py
user: login: PabloRoque
body: Accordingly, same as last comment
created_at: 2025-05-07T17:33:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -338,105 +333,112 @@ def create_event_mu_effect(
         The prefix to use for the event effect and associated variables.
     effect : EventEffect
         The event effect to apply.
-
-    Returns
-    -------
-    MuEffect
-        The event effect which is used in the MMM.
+    reference_date : str
+        The arbitrary reference date to calculate distance from events in days. Default
+        is "2025-01-01".
 
     """
-    if missing_columns := set(["start_date", "end_date", "name"]).difference(
-        df_events.columns,
-    ):
-        raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    effect.basis.prefix = prefix
+    df_events: InstanceOf[pd.DataFrame]
+    prefix: str
+    effect: EventEffect
+    reference_date: str = "2025-01-01"
 
-    reference_date = "2025-01-01"
-    start_dates = pd.to_datetime(df_events["start_date"])
-    end_dates = pd.to_datetime(df_events["end_date"])
+    def model_post_init(self, context: Any, /) -> None:
+        """Post initialization of the model."""
+        if missing_columns := set(["start_date", "end_date", "name"]).difference(
+            self.df_events.columns
+        ):
+            raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    class Effect:
-        """Event effect class for the MMM."""
+        self.effect.basis.prefix = self.prefix
 
-        def create_data(self, mmm: MMM) -> None:
-            """Create the required data in the model.
+    @property
+    def start_dates(self) -> pd.Series:
+        """The start dates of the events."""
+        return pd.to_datetime(self.df_events["start_date"])
 
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+    @property
+    def end_dates(self) -> pd.Series:
+        """The end dates of the events."""
+        return pd.to_datetime(self.df_events["end_date"])
 
-            """
-            model: pm.Model = mmm.model
+    def create_data(self, mmm: MMM) -> None:
+        """Create the required data in the model.
 
-            model_dates = pd.to_datetime(model.coords["date"])
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            model.add_coord(prefix, df_events["name"].to_numpy())
+        """
+        model: pm.Model = mmm.model
 
-            if "days" not in model:
-                pm.Data(
-                    "days",
-                    days_from_reference(model_dates, reference_date),
-                    dims="date",
-                )
+        model_dates = pd.to_datetime(model.coords["date"])
 
+        model.add_coord(self.prefix, self.df_events["name"].to_numpy())
+
+        if "days" not in model:
             pm.Data(
-                f"{prefix}_start_diff",
-                days_from_reference(start_dates, reference_date),
-                dims=prefix,
-            )
-            pm.Data(
-                f"{prefix}_end_diff",
-                days_from_reference(end_dates, reference_date),
-                dims=prefix,
+                "days",
+                days_from_reference(model_dates, self.reference_date),
+                dims="date",
             )
 
-        def create_effect(self, mmm: MMM) -> pt.TensorVariable:
-            """Create the event effect in the model.
-
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+        pm.Data(
+            f"{self.prefix}_start_diff",
+            days_from_reference(self.start_dates, self.reference_date),
+            dims=self.prefix,
+        )
+        pm.Data(
+            f"{self.prefix}_end_diff",
+            days_from_reference(self.end_dates, self.reference_date),
+            dims=self.prefix,
+        )
 
-            Returns
-            -------
-            pt.TensorVariable
-                The average event effect in the model.
+    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+        """Create the event effect in the model.
 
-            """
-            model: pm.Model = mmm.model
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            s_ref = model["days"][:, None] - model[f"{prefix}_start_diff"]
-            e_ref = model["days"][:, None] - model[f"{prefix}_end_diff"]
+        Returns
+        -------
+        pt.TensorVariable
+            The average event effect in the model.
 
-            def create_basis_matrix(s_ref, e_ref):
-                return pt.where(
-                    (s_ref >= 0) & (e_ref <= 0),
-                    0,
-                    pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
-                )
+        """
+        model: pm.Model = mmm.model
 
-            X = create_basis_matrix(s_ref, e_ref)
-            event_effect = effect.apply(X, name=prefix)
+        s_ref = model["days"][:, None] - model[f"{self.prefix}_start_diff"]
+        e_ref = model["days"][:, None] - model[f"{self.prefix}_end_diff"]
 
-            total_effect = pm.Deterministic(
-                f"{prefix}_total_effect",
-                event_effect.sum(axis=1),
-                dims="date",
+        def create_basis_matrix(s_ref, e_ref):
+            return pt.where(
+                (s_ref >= 0) & (e_ref <= 0),
+                0,
+                pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
             )
 
-            dim_handler = create_dim_handler(("date", *mmm.dims))
-            return dim_handler(total_effect, "date")
+        X = create_basis_matrix(s_ref, e_ref)
+        event_effect = self.effect.apply(X, name=self.prefix)
 
-        def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
-            """Set the data for new predictions."""
-            new_dates = pd.to_datetime(model.coords["date"])
+        total_effect = pm.Deterministic(
+            f"{self.prefix}_total_effect",
+            event_effect.sum(axis=1),
+            dims="date",
+        )
 
-            new_data = {
-                "days": days_from_reference(new_dates, reference_date),
-            }
-            pm.set_data(new_data=new_data, model=model)
+        dim_handler = create_dim_handler(("date", *mmm.dims))
+        return dim_handler(total_effect, "date")
+
+    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
path: pymc_marketing/mmm/additive_effect.py
user: login: PabloRoque
body: Do we need both `mmm` and `model`? Can't we use `mmm.model`?
created_at: 2025-05-07T17:38:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -338,105 +333,112 @@ def create_event_mu_effect(
         The prefix to use for the event effect and associated variables.
     effect : EventEffect
         The event effect to apply.
-
-    Returns
-    -------
-    MuEffect
-        The event effect which is used in the MMM.
+    reference_date : str
+        The arbitrary reference date to calculate distance from events in days. Default
+        is "2025-01-01".
 
     """
-    if missing_columns := set(["start_date", "end_date", "name"]).difference(
-        df_events.columns,
-    ):
-        raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    effect.basis.prefix = prefix
+    df_events: InstanceOf[pd.DataFrame]
+    prefix: str
+    effect: EventEffect
+    reference_date: str = "2025-01-01"
 
-    reference_date = "2025-01-01"
-    start_dates = pd.to_datetime(df_events["start_date"])
-    end_dates = pd.to_datetime(df_events["end_date"])
+    def model_post_init(self, context: Any, /) -> None:
+        """Post initialization of the model."""
+        if missing_columns := set(["start_date", "end_date", "name"]).difference(
+            self.df_events.columns
+        ):
+            raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    class Effect:
-        """Event effect class for the MMM."""
+        self.effect.basis.prefix = self.prefix
 
-        def create_data(self, mmm: MMM) -> None:
-            """Create the required data in the model.
+    @property
+    def start_dates(self) -> pd.Series:
+        """The start dates of the events."""
+        return pd.to_datetime(self.df_events["start_date"])
 
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+    @property
+    def end_dates(self) -> pd.Series:
+        """The end dates of the events."""
+        return pd.to_datetime(self.df_events["end_date"])
 
-            """
-            model: pm.Model = mmm.model
+    def create_data(self, mmm: MMM) -> None:
+        """Create the required data in the model.
 
-            model_dates = pd.to_datetime(model.coords["date"])
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            model.add_coord(prefix, df_events["name"].to_numpy())
+        """
+        model: pm.Model = mmm.model
 
-            if "days" not in model:
-                pm.Data(
-                    "days",
-                    days_from_reference(model_dates, reference_date),
-                    dims="date",
-                )
+        model_dates = pd.to_datetime(model.coords["date"])
 
+        model.add_coord(self.prefix, self.df_events["name"].to_numpy())
+
+        if "days" not in model:
             pm.Data(
-                f"{prefix}_start_diff",
-                days_from_reference(start_dates, reference_date),
-                dims=prefix,
-            )
-            pm.Data(
-                f"{prefix}_end_diff",
-                days_from_reference(end_dates, reference_date),
-                dims=prefix,
+                "days",
+                days_from_reference(model_dates, self.reference_date),
+                dims="date",
             )
 
-        def create_effect(self, mmm: MMM) -> pt.TensorVariable:
-            """Create the event effect in the model.
-
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+        pm.Data(
+            f"{self.prefix}_start_diff",
+            days_from_reference(self.start_dates, self.reference_date),
+            dims=self.prefix,
+        )
+        pm.Data(
+            f"{self.prefix}_end_diff",
+            days_from_reference(self.end_dates, self.reference_date),
+            dims=self.prefix,
+        )
 
-            Returns
-            -------
-            pt.TensorVariable
-                The average event effect in the model.
+    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+        """Create the event effect in the model.
 
-            """
-            model: pm.Model = mmm.model
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            s_ref = model["days"][:, None] - model[f"{prefix}_start_diff"]
-            e_ref = model["days"][:, None] - model[f"{prefix}_end_diff"]
+        Returns
+        -------
+        pt.TensorVariable
+            The average event effect in the model.
 
-            def create_basis_matrix(s_ref, e_ref):
-                return pt.where(
-                    (s_ref >= 0) & (e_ref <= 0),
-                    0,
-                    pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
-                )
+        """
+        model: pm.Model = mmm.model
 
-            X = create_basis_matrix(s_ref, e_ref)
-            event_effect = effect.apply(X, name=prefix)
+        s_ref = model["days"][:, None] - model[f"{self.prefix}_start_diff"]
+        e_ref = model["days"][:, None] - model[f"{self.prefix}_end_diff"]
 
-            total_effect = pm.Deterministic(
-                f"{prefix}_total_effect",
-                event_effect.sum(axis=1),
-                dims="date",
+        def create_basis_matrix(s_ref, e_ref):
+            return pt.where(
+                (s_ref >= 0) & (e_ref <= 0),
+                0,
+                pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
             )
 
-            dim_handler = create_dim_handler(("date", *mmm.dims))
-            return dim_handler(total_effect, "date")
+        X = create_basis_matrix(s_ref, e_ref)
+        event_effect = self.effect.apply(X, name=self.prefix)
 
-        def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
-            """Set the data for new predictions."""
-            new_dates = pd.to_datetime(model.coords["date"])
+        total_effect = pm.Deterministic(
+            f"{self.prefix}_total_effect",
+            event_effect.sum(axis=1),
+            dims="date",
+        )
 
-            new_data = {
-                "days": days_from_reference(new_dates, reference_date),
-            }
-            pm.set_data(new_data=new_data, model=model)
+        dim_handler = create_dim_handler(("date", *mmm.dims))
+        return dim_handler(total_effect, "date")
+
+    def set_data(self, mmm: MMM, model: pm.Model, X: xr.Dataset) -> None:
path: pymc_marketing/mmm/additive_effect.py
user: login: williambdean
body: The mmm class used a copy of model for the predictions which isnt stored
created_at: 2025-05-07T17:59:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -338,105 +333,112 @@ def create_event_mu_effect(
         The prefix to use for the event effect and associated variables.
     effect : EventEffect
         The event effect to apply.
-
-    Returns
-    -------
-    MuEffect
-        The event effect which is used in the MMM.
+    reference_date : str
+        The arbitrary reference date to calculate distance from events in days. Default
+        is "2025-01-01".
 
     """
-    if missing_columns := set(["start_date", "end_date", "name"]).difference(
-        df_events.columns,
-    ):
-        raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    effect.basis.prefix = prefix
+    df_events: InstanceOf[pd.DataFrame]
+    prefix: str
+    effect: EventEffect
+    reference_date: str = "2025-01-01"
 
-    reference_date = "2025-01-01"
-    start_dates = pd.to_datetime(df_events["start_date"])
-    end_dates = pd.to_datetime(df_events["end_date"])
+    def model_post_init(self, context: Any, /) -> None:
+        """Post initialization of the model."""
+        if missing_columns := set(["start_date", "end_date", "name"]).difference(
+            self.df_events.columns
+        ):
+            raise ValueError(f"Columns {missing_columns} are missing in df_events.")
 
-    class Effect:
-        """Event effect class for the MMM."""
+        self.effect.basis.prefix = self.prefix
 
-        def create_data(self, mmm: MMM) -> None:
-            """Create the required data in the model.
+    @property
+    def start_dates(self) -> pd.Series:
+        """The start dates of the events."""
+        return pd.to_datetime(self.df_events["start_date"])
 
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+    @property
+    def end_dates(self) -> pd.Series:
+        """The end dates of the events."""
+        return pd.to_datetime(self.df_events["end_date"])
 
-            """
-            model: pm.Model = mmm.model
+    def create_data(self, mmm: MMM) -> None:
+        """Create the required data in the model.
 
-            model_dates = pd.to_datetime(model.coords["date"])
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            model.add_coord(prefix, df_events["name"].to_numpy())
+        """
+        model: pm.Model = mmm.model
 
-            if "days" not in model:
-                pm.Data(
-                    "days",
-                    days_from_reference(model_dates, reference_date),
-                    dims="date",
-                )
+        model_dates = pd.to_datetime(model.coords["date"])
 
+        model.add_coord(self.prefix, self.df_events["name"].to_numpy())
+
+        if "days" not in model:
             pm.Data(
-                f"{prefix}_start_diff",
-                days_from_reference(start_dates, reference_date),
-                dims=prefix,
-            )
-            pm.Data(
-                f"{prefix}_end_diff",
-                days_from_reference(end_dates, reference_date),
-                dims=prefix,
+                "days",
+                days_from_reference(model_dates, self.reference_date),
+                dims="date",
             )
 
-        def create_effect(self, mmm: MMM) -> pt.TensorVariable:
-            """Create the event effect in the model.
-
-            Parameters
-            ----------
-            mmm : MMM
-                The MMM model instance.
+        pm.Data(
+            f"{self.prefix}_start_diff",
+            days_from_reference(self.start_dates, self.reference_date),
+            dims=self.prefix,
+        )
+        pm.Data(
+            f"{self.prefix}_end_diff",
+            days_from_reference(self.end_dates, self.reference_date),
+            dims=self.prefix,
+        )
 
-            Returns
-            -------
-            pt.TensorVariable
-                The average event effect in the model.
+    def create_effect(self, mmm: MMM) -> pt.TensorVariable:
+        """Create the event effect in the model.
 
-            """
-            model: pm.Model = mmm.model
+        Parameters
+        ----------
+        mmm : MMM
+            The MMM model instance.
 
-            s_ref = model["days"][:, None] - model[f"{prefix}_start_diff"]
-            e_ref = model["days"][:, None] - model[f"{prefix}_end_diff"]
+        Returns
+        -------
+        pt.TensorVariable
+            The average event effect in the model.
 
-            def create_basis_matrix(s_ref, e_ref):
-                return pt.where(
-                    (s_ref >= 0) & (e_ref <= 0),
-                    0,
-                    pt.where(pt.abs(s_ref) < pt.abs(e_ref), s_ref, e_ref),
-                )
+        """
+        model: pm.Model = mmm.model
 
-            X = create_basis_matrix(s_ref, e_ref)
-            event_effect = effect.apply(X, name=prefix)
+        s_ref = model["days"][:, None] - model[f"{self.prefix}_start_diff"]
+        e_ref = model["days"][:, None] - model[f"{self.prefix}_end_diff"]
 
-            total_effect = pm.Deterministic(
-                f"{prefix}_total_effect",
-                event_effect.sum(axis=1),
-                dims="date",
+        def create_basis_matrix(s_ref, e_ref):
path: pymc_marketing/mmm/additive_effect.py
user: login: williambdean
body: done. Good suggestion
created_at: 2025-05-08T00:27:25Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-05-07T16:57:10Z
body: ## [Codecov](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs) Report
:x: Patch coverage is `39.02439%` with `25 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 69.53%. Comparing base ([`51c75c1`](https://app.codecov.io/gh/pymc-labs/pymc-marketing/commit/51c75c188bdad1e0eee3a05cb0a779ca080fb202?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs)) to head ([`6bd990f`](https://app.codecov.io/gh/pymc-labs/pymc-marketing/commit/6bd990fb107693bb24c88e285279d113717e9210?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs)).
:warning: Report is 94 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs) | Patch % | Lines |
|---|---|---|
| [pymc\_marketing/mmm/additive\_effect.py](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?src=pr&el=tree&filepath=pymc_marketing%2Fmmm%2Fadditive_effect.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs#diff-cHltY19tYXJrZXRpbmcvbW1tL2FkZGl0aXZlX2VmZmVjdC5weQ==) | 38.46% | [24 Missing :warning: ](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs) |
| [pymc\_marketing/mmm/multidimensional.py](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?src=pr&el=tree&filepath=pymc_marketing%2Fmmm%2Fmultidimensional.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs#diff-cHltY19tYXJrZXRpbmcvbW1tL211bHRpZGltZW5zaW9uYWwucHk=) | 50.00% | [1 Missing :warning: ](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs) |

<details><summary>Additional details and impacted files</summary>



```diff
@@            Coverage Diff             @@
##             main    #1675      +/-   ##
==========================================
+ Coverage   69.40%   69.53%   +0.12%     
==========================================
  Files          56       56              
  Lines        6623     6630       +7     
==========================================
+ Hits         4597     4610      +13     
+ Misses       2026     2020       -6     
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pymc-labs/pymc-marketing/pull/1675?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-labs).
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: williambdean
created_at: 2025-05-08T00:28:33Z
body: Thanks for the review and feedback @PabloRoque 

Happy to change the signature in a future review if possible. But it was needed for the out of sample due to the copy of PyMC model being used.
author_association: CONTRIBUTOR