Add LinearRegression wrapper (#1713)

* implement the transformations

* wrapper around linear regression

* expose via the mmm module

* rename variable to beta

* test for linear regression

* change name to NoSaturation

Author: williambdean

pymc_marketing/mmm/__init__.py

@@ -19,6 +19,7 @@
     AdstockTransformation,
     DelayedAdstock,
     GeometricAdstock,
+    NoAdstock,
     WeibullCDFAdstock,
     WeibullPDFAdstock,
     adstock_from_dict,
@@ -29,6 +30,7 @@
     InverseScaledLogisticSaturation,
     LogisticSaturation,
     MichaelisMentenSaturation,
+    NoSaturation,
     RootSaturation,
     SaturationTransformation,
     TanhSaturation,
@@ -48,6 +50,7 @@
     create_eta_prior,
     create_m_and_L_recommendations,
 )
+from pymc_marketing.mmm.linear_regression import FancyLinearRegression
 from pymc_marketing.mmm.linear_trend import LinearTrend
 from pymc_marketing.mmm.media_transformation import (
     MediaConfig,
@@ -68,6 +71,7 @@
     "BaseValidateMMM",
     "CovFunc",
     "DelayedAdstock",
+    "FancyLinearRegression",
     "GeometricAdstock",
     "HSGPPeriodic",
     "HillSaturation",
@@ -81,6 +85,8 @@
     "MediaTransformation",
     "MichaelisMentenSaturation",
     "MonthlyFourier",
+    "NoAdstock",
+    "NoSaturation",
     "PeriodicCovFunc",
     "RootSaturation",
     "SaturationTransformation",

pymc_marketing/mmm/components/adstock.py

@@ -329,6 +329,22 @@ def function(self, x, lam, k):
     }
 
 
+class NoAdstock(AdstockTransformation):
+    """Wrapper around no adstock transformation."""
+
+    lookup_name: str = "no_adstock"
+
+    def function(self, x):
+        """No adstock function."""
+        return x
+
+    default_priors = {}
+
+    def update_priors(self, priors):
+        """Update priors for the no adstock transformation."""
+        return
+
+
 def adstock_from_dict(data: dict) -> AdstockTransformation:
     """Create an adstock transformation from a dictionary."""
     data = data.copy()

pymc_marketing/mmm/components/saturation.py

@@ -459,6 +459,37 @@ def function(self, x, alpha, beta):
     }
 
 
+class NoSaturation(SaturationTransformation):
+    """Wrapper around linear saturation function.
+
+    For more information, see :func:`pymc_marketing.mmm.transformers.linear_saturation`.
+
+    .. plot::
+        :context: close-figs
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        from pymc_marketing.mmm import NoSaturation
+
+        rng = np.random.default_rng(0)
+
+        saturation = NoSaturation()
+        prior = saturation.sample_prior(random_seed=rng)
+        curve = saturation.sample_curve(prior)
+        saturation.plot_curve(curve, random_seed=rng)
+        plt.show()
+
+    """
+
+    lookup_name = "linear"
+
+    def function(self, x, beta):
+        """Linear saturation function."""
+        return pt.as_tensor_variable(beta * x)
+
+    default_priors = {"beta": Prior("HalfNormal", sigma=1)}
+
+
 def saturation_from_dict(data: dict) -> SaturationTransformation:
     """Get a saturation function from a dictionary."""
     data = data.copy()

pymc_marketing/mmm/linear_regression.py

@@ -0,0 +1,53 @@
+#   Copyright 2022 - 2025 The PyMC Labs Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+"""Linear regression model implemented using the MMM class."""
+
+from pymc_marketing.mmm.components.adstock import NoAdstock
+from pymc_marketing.mmm.components.saturation import NoSaturation
+from pymc_marketing.mmm.mmm import MMM
+
+
+def FancyLinearRegression(
+    **mmm_kwargs,
+) -> MMM:
+    """Create wrapper around MMM for a linear regression model.
+
+    See :func:`pymc_marketing.mmm.mmm.MMM` for more details.
+
+    Parameters
+    ----------
+    mmm_kwargs
+        Keyword arguments to pass to the MMM constructor.
+
+    Returns
+    -------
+    MMM
+        An instance of the MMM class with linear regression settings.
+
+    Examples
+    --------
+    Load a saved MMM model with linear regression settings:
+
+    .. code-block:: python
+
+        from pymc_marketing.mmm import MMM
+
+        linear_regression = MMM.load("linear_regression_model.nc")
+
+    """
+    return MMM(
+        adstock=NoAdstock(l_max=1),
+        saturation=NoSaturation(),
+        **mmm_kwargs,
+    )

tests/mmm/test_linear_regression.py

@@ -0,0 +1,93 @@
+#   Copyright 2022 - 2025 The PyMC Labs Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+import numpy as np
+import pandas as pd
+import pytest
+
+from pymc_marketing.mmm.linear_regression import FancyLinearRegression
+from pymc_marketing.mmm.mmm import MMM
+
+seed: int = sum(map(ord, "pymc_marketing"))
+rng: np.random.Generator = np.random.default_rng(seed=seed)
+
+
+@pytest.fixture(scope="module")
+def generate_data():
+    def _generate_data(date_data: pd.DatetimeIndex) -> pd.DataFrame:
+        n: int = date_data.size
+
+        return pd.DataFrame(
+            data={
+                "date": date_data,
+                "channel_1": rng.integers(low=0, high=400, size=n),
+                "channel_2": rng.integers(low=0, high=50, size=n),
+                "control_1": rng.gamma(shape=1000, scale=500, size=n),
+                "control_2": rng.gamma(shape=100, scale=5, size=n),
+                "other_column_1": rng.integers(low=0, high=100, size=n),
+                "other_column_2": rng.normal(loc=0, scale=1, size=n),
+            }
+        )
+
+    return _generate_data
+
+
+@pytest.fixture(scope="module")
+def toy_X(generate_data) -> pd.DataFrame:
+    date_data: pd.DatetimeIndex = pd.date_range(
+        start="2019-06-01", end="2021-12-31", freq="W-MON"
+    )
+
+    return generate_data(date_data)
+
+
+@pytest.fixture(scope="module")
+def toy_y(toy_X: pd.DataFrame) -> pd.Series:
+    return pd.Series(data=rng.integers(low=0, high=100, size=toy_X.shape[0]))
+
+
+@pytest.fixture(scope="module")
+def linear_regression() -> MMM:
+    return FancyLinearRegression(
+        date_column="date",
+        channel_columns=["channel_1", "channel_2"],
+        control_columns=["control_1", "control_2"],
+    )
+
+
+def test_fancy_linear_regression(
+    linear_regression: MMM,
+    toy_X: pd.DataFrame,
+    toy_y: pd.Series,
+    mock_pymc_sample,
+) -> None:
+    """Test that FancyLinearRegression returns an instance of MMM."""
+    assert isinstance(linear_regression, MMM)
+
+    linear_regression.fit(toy_X, toy_y)
+
+    assert set(linear_regression.fit_result.variables) == {
+        "chain",
+        "draw",
+        "date",
+        "channel",
+        "channel_contributions",
+        "control",
+        "control_contributions",
+        "intercept",
+        "mu",
+        "gamma_control",
+        "saturation_beta",
+        "total_contributions",
+        "y_sigma",
+    }