Change the event function to class (#1675)

* turn function into class

* use the class in the tests

* add docstrings

* remove the original function version

* add to the docstrig

* rename the variables

* change variable name in other documentation location

Author: williambdean

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)

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),
         )
 
 

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

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

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)