Time varying intercept

docs/source/notebooks/index.md

@@ -9,6 +9,7 @@
 mmm/mmm_example
 mmm/mmm_budget_allocation_example
 mmm/mmm_lift_test
+mmm/mmm_tvp_example
 :::
 
 :::{toctree}

pymc_marketing/constants.py

@@ -0,0 +1 @@
+DAYS_IN_YEAR = 365.25

pymc_marketing/mmm/base.py

@@ -19,6 +19,7 @@
 import pandas as pd
 import pymc as pm
 import seaborn as sns
+from numpy.typing import NDArray
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import FunctionTransformer
 from xarray import DataArray, Dataset
@@ -31,6 +32,7 @@
     find_sigmoid_inflection_point,
     sigmoid_saturation,
     standardize_scenarios_dict_keys,
+    transform_1d_array,
 )
 from pymc_marketing.mmm.validating import (
     ValidateChannelColumns,
@@ -55,13 +57,19 @@ def __init__(
         sampler_config: dict | None = None,
         **kwargs,
     ) -> None:
-        self.X: pd.DataFrame | None = None
-        self.y: pd.Series | np.ndarray | None = None
         self.date_column: str = date_column
         self.channel_columns: list[str] | tuple[str] = channel_columns
+
         self.n_channel: int = len(channel_columns)
-        self._fit_result: az.InferenceData | None = None
-        self._posterior_predictive: az.InferenceData | None = None
+
+        self.X: pd.DataFrame
+        self.y: pd.Series | np.ndarray
+
+        self._time_resolution: int
+        self._time_index: NDArray[np.int_]
+        self._time_index_mid: int
+        self._fit_result: az.InferenceData
+        self._posterior_predictive: az.InferenceData
         super().__init__(model_config=model_config, sampler_config=sampler_config)
 
     @property
@@ -314,7 +322,7 @@ def plot_prior_predictive(
         return fig
 
     def plot_posterior_predictive(
-        self, original_scale: bool = False, **plt_kwargs: Any
+        self, original_scale: bool = False, ax: plt.Axes = None, **plt_kwargs: Any
     ) -> plt.Figure:
         posterior_predictive_data: Dataset = self.posterior_predictive
         likelihood_hdi_94: DataArray = az.hdi(
@@ -332,10 +340,14 @@ def plot_posterior_predictive(
                 Xt=likelihood_hdi_50
             )
 
-        fig, ax = plt.subplots(**plt_kwargs)
+        if ax is None:
+            fig, ax = plt.subplots(**plt_kwargs)
+        else:
+            fig = ax.figure
+
         if self.X is not None and self.y is not None:
             ax.fill_between(
-                x=self.X[self.date_column],
+                x=posterior_predictive_data.date,
                 y1=likelihood_hdi_94[:, 0],
                 y2=likelihood_hdi_94[:, 1],
                 color="C0",
@@ -344,19 +356,29 @@ def plot_posterior_predictive(
             )
 
             ax.fill_between(
-                x=self.X[self.date_column],
+                x=posterior_predictive_data.date,
                 y1=likelihood_hdi_50[:, 0],
                 y2=likelihood_hdi_50[:, 1],
                 color="C0",
                 alpha=0.3,
                 label="$50\%$ HDI",  # noqa: W605
             )
 
-            target_to_plot: np.ndarray = np.asarray(
-                self.y if original_scale else self.preprocessed_data["y"]  # type: ignore
+            target_to_plot = np.asarray(
+                self.y
+                if original_scale
+                else transform_1d_array(self.get_target_transformer().transform, self.y)
             )
+
+            if len(target_to_plot) != len(posterior_predictive_data.date):
+                raise ValueError(
+                    "The length of the target variable doesn't match the length of the date column. "
+                    "If you are predicting out-of-sample, please overwrite `self.y` with the "
+                    "corresponding (non-transformed) target variable."
+                )
+
             ax.plot(
-                np.asarray(self.X[self.date_column]),
+                np.asarray(posterior_predictive_data.date),
                 target_to_plot,
                 color="black",
             )
@@ -435,11 +457,18 @@ def plot_components_contributions(self, **plt_kwargs: Any) -> plt.Figure:
             intercept = az.extract(
                 self.fit_result, var_names=["intercept"], combined=False
             )
-            intercept_hdi = np.repeat(
-                a=az.hdi(intercept).intercept.data[None, ...],
-                repeats=self.X[self.date_column].shape[0],
-                axis=0,
-            )
+
+            if intercept.ndim == 2:
+                # Intercept has a stationary prior
+                intercept_hdi = np.repeat(
+                    a=az.hdi(intercept).intercept.data[None, ...],
+                    repeats=self.X[self.date_column].shape[0],
+                    axis=0,
+                )
+            elif intercept.ndim == 3:
+                # Intercept has a time-varying prior
+                intercept_hdi = az.hdi(intercept).intercept.data
+
             ax.plot(
                 np.asarray(self.X[self.date_column]),
                 np.full(len(self.X[self.date_column]), intercept.mean().data),
@@ -1028,6 +1057,7 @@ def label_func(channel):
 
             def legend_title_func(channel):
                 return "Legend"
+
         else:
             nrows = len(channels_to_plot)
             figsize = (12, 4 * len(channels_to_plot))

pymc_marketing/mmm/delayed_saturated_mmm.py

@@ -14,13 +14,15 @@
 from pytensor.tensor import TensorVariable
 from xarray import DataArray, Dataset
 
+from pymc_marketing.constants import DAYS_IN_YEAR
 from pymc_marketing.mmm.base import MMM
 from pymc_marketing.mmm.lift_test import (
     add_logistic_empirical_lift_measurements_to_likelihood,
     scale_lift_measurements,
 )
 from pymc_marketing.mmm.preprocessing import MaxAbsScaleChannels, MaxAbsScaleTarget
 from pymc_marketing.mmm.transformers import geometric_adstock, logistic_saturation
+from pymc_marketing.mmm.tvp import create_time_varying_intercept, infer_time_index
 from pymc_marketing.mmm.utils import (
     apply_sklearn_transformer_across_dim,
     create_new_spend_data,
@@ -47,6 +49,7 @@ def __init__(
         date_column: str,
         channel_columns: list[str],
         adstock_max_lag: int,
+        time_varying_intercept: bool = False,
         model_config: dict | None = None,
         sampler_config: dict | None = None,
         validate_data: bool = True,
@@ -62,6 +65,10 @@ def __init__(
             Column name of the date variable.
         channel_columns : List[str]
             Column names of the media channel variables.
+        adstock_max_lag : int
+            Number of lags to consider in the adstock transformation.
+        time_varying_intercept : bool, optional
+            Whether to consider time-varying intercept, by default False.
         model_config : Dictionary, optional
             dictionary of parameters that initialise model configuration.
             Class-default defined by the user default_model_config method.
@@ -79,6 +86,7 @@ def __init__(
         """
         self.control_columns = control_columns
         self.adstock_max_lag = adstock_max_lag
+        self.time_varying_intercept = time_varying_intercept
         self.yearly_seasonality = yearly_seasonality
         self.date_column = date_column
         self.validate_data = validate_data
@@ -112,6 +120,24 @@ def _generate_and_preprocess_model_data(  # type: ignore
         ----------
         X : Union[pd.DataFrame, pd.Series], shape (n_obs, n_features)
         y : Union[pd.Series, np.ndarray], shape (n_obs,)
+
+        Sets
+        ----
+        preprocessed_data : Dict[str, Union[pd.DataFrame, pd.Series]]
+            Preprocessed data for the model.
+        X : pd.DataFrame
+            A filtered version of the input `X`, such that it is guaranteed that
+            it contains only the `date_column`, the columns that are specified
+            in the `channel_columns` and `control_columns`, and fourier features
+            if `yearly_seasonality=True`.
+        y : Union[pd.Series, np.ndarray]
+            The target variable for the model (as provided).
+        _time_index : np.ndarray
+            The index of the date column. Used by TVP
+        _time_index_mid : int
+            The middle index of the date index. Used by TVP.
+        _time_resolution: int
+            The time resolution of the date index. Used by TVP.
         """
         date_data = X[self.date_column]
         channel_data = X[self.channel_columns]
@@ -152,6 +178,13 @@ def _generate_and_preprocess_model_data(  # type: ignore
         self.X: pd.DataFrame = X_data
         self.y: pd.Series | np.ndarray = y
 
+        if self.time_varying_intercept:
+            self._time_index = np.arange(0, X.shape[0])
+            self._time_index_mid = X.shape[0] // 2
+            self._time_resolution = (
+                self.X[self.date_column].iloc[1] - self.X[self.date_column].iloc[0]
+            ).days
+
     def _save_input_params(self, idata) -> None:
         """Saves input parameters to the attrs of idata."""
         idata.attrs["date_column"] = json.dumps(self.date_column)
@@ -355,9 +388,23 @@ def build_model(
                 dims="date",
             )
 
-            intercept = self.intercept_dist(
-                name="intercept", **self.model_config["intercept"]["kwargs"]
-            )
+            if self.time_varying_intercept:
+                time_index = pm.Data(
+                    "time_index",
+                    self._time_index,
+                    dims="date",
+                )
+                intercept = create_time_varying_intercept(
+                    time_index,
+                    self._time_index_mid,
+                    self._time_resolution,
+                    self.intercept_dist,
+                    self.model_config,
+                )
+            else:
+                intercept = self.intercept_dist(
+                    name="intercept", **self.model_config["intercept"]["kwargs"]
+                )
 
             beta_channel = self.beta_channel_dist(
                 name="beta_channel",
@@ -391,9 +438,11 @@ def build_model(
                 var=logistic_saturation(x=channel_adstock, lam=lam),
                 dims=("date", "channel"),
             )
+
+            channel_contributions_var = channel_adstock_saturated * beta_channel
             channel_contributions = pm.Deterministic(
                 name="channel_contributions",
-                var=channel_adstock_saturated * beta_channel,
+                var=channel_contributions_var,
                 dims=("date", "channel"),
             )
 
@@ -468,7 +517,10 @@ def build_model(
     @property
     def default_model_config(self) -> dict:
         return {
-            "intercept": {"dist": "Normal", "kwargs": {"mu": 0, "sigma": 2}},
+            "intercept": {
+                "dist": "Normal",
+                "kwargs": {"mu": 0, "sigma": 2},
+            },
             "beta_channel": {"dist": "HalfNormal", "kwargs": {"sigma": 2}},
             "alpha": {"dist": "Beta", "kwargs": {"alpha": 1, "beta": 3}},
             "lam": {"dist": "Gamma", "kwargs": {"alpha": 3, "beta": 1}},
@@ -480,6 +532,14 @@ def default_model_config(self) -> dict:
             },
             "gamma_control": {"dist": "Normal", "kwargs": {"mu": 0, "sigma": 2}},
             "gamma_fourier": {"dist": "Laplace", "kwargs": {"mu": 0, "b": 1}},
+            "intercept_tvp_kwargs": {
+                "m": 200,
+                "L": None,
+                "eta_lam": 1,
+                "ls_mu": None,
+                "ls_sigma": 10,
+                "cov_func": None,
+            },
         }
 
     def _get_fourier_models_data(self, X) -> pd.DataFrame:
@@ -494,7 +554,9 @@ def _get_fourier_models_data(self, X) -> pd.DataFrame:
         date_data: pd.Series = pd.to_datetime(
             arg=X[self.date_column], format="%Y-%m-%d"
         )
-        periods: npt.NDArray[np.float_] = date_data.dt.dayofyear.to_numpy() / 365.25
+        periods: npt.NDArray[np.float_] = (
+            date_data.dt.dayofyear.to_numpy() / DAYS_IN_YEAR
+        )
         return generate_fourier_modes(
             periods=periods,
             n_order=self.yearly_seasonality,
@@ -678,6 +740,11 @@ def identity(x):
         if hasattr(self, "fourier_columns"):
             data["fourier_data"] = self._get_fourier_models_data(X)
 
+        if self.time_varying_intercept:
+            data["time_index"] = infer_time_index(
+                X[self.date_column], self.X[self.date_column], self._time_resolution
+            )
+
         if y is not None:
             if isinstance(y, pd.Series):
                 data["target"] = (