Make MMM Model component base

[thipokKub]

Hello, I really like this repo. It helped me a lot in a past few months. Currently all the MMM model development seems to focus on BaseDelayedSaturatedMMM class. Because MMM model is a GLM model, it can be decomposed into multiple components such as trend, seasonality, media, and control variables

However, the current model class fixed the method for trend to always be linear, media, and control variables must always have the same prior and such. This is somewhat limiting. I proposed the idea of implementing MMM as a component based model such that it can be customize more easily

The idea stem from #274

For example this is my proof of concept idea

Setup Code

import pymc as pm
import numpy as np
import pandas as pd
from enum import Enum
import pytensor.tensor as pt

class ComponentBase(pm.Model):
    pass

class TrendComponent(ComponentBase):
    def __init__(self, name, t):
        super().__init__(name)
        self._t = pm.MutableData("t", t)

class SeasonalityComponent(ComponentBase):
    def __init__(self, name, fourier_features):
        super().__init__(name)
        self._fourier_features = pm.MutableData("fourier_features", fourier_features)

Utils code

# Utils?
class Spline:
    # Enum
    class SplineType(Enum):
        LINEAR = "Linear"
        GAUSSIAN = "Gaussian"
        PARABOLIC = "Parabolic"
        
    # # Helpers
    @staticmethod
    def get_spline(t, t_refs, ktype: SplineType = SplineType.LINEAR, **kwargs):
        if ktype == Spline.SplineType.LINEAR:
            return Spline.get_linear_spline(t, t_refs, **kwargs)
        elif ktype == Spline.SplineType.GAUSSIAN:
            return Spline.get_gaussian_spline(t, t_refs, **kwargs)
        elif ktype == Spline.SplineType.PARABOLIC:
            return Spline.get_parabolic_spline(t, t_refs, **kwargs)
        else:
            raise NotImplementedError
    
    @staticmethod
    def get_linear_spline(t, t_refs):
        ker = np.zeros((t.shape[0], t_refs.shape[0]))
        ker[t < t_refs[0], 0] = 1
        for idx in range(t_refs.shape[0] - 1):
            valid_idxes = (t >= t_refs[idx]) & (t < t_refs[idx + 1])
            norm = t_refs[idx + 1] - t_refs[idx]
            ker[valid_idxes, idx] = np.abs(t[valid_idxes] - t_refs[idx + 1])/norm
            ker[valid_idxes, idx + 1] = np.abs(t[valid_idxes] - t_refs[idx])/norm
        ker[t >= t_refs[-1], -1] = 1
        return ker / np.sum(ker, axis=1, keepdims=True)
    
    @staticmethod
    def get_gaussian_spline(t, t_refs, rho: float = 0.1, alpha: float = 1, point_to_flatten: float = 1):
        ker = np.where(
            (t <= point_to_flatten).reshape(-1, 1),
            (alpha**2) * np.exp(-1 * np.power(t.reshape(-1, 1) - t_refs.reshape(1, -1), 2)/(2 * rho**2 )),
            (alpha**2) * np.exp(-1 * np.power(np.array([point_to_flatten] * t.shape[0]).reshape(-1, 1) - t_refs.reshape(1, -1), 2)/(2 * rho**2 ))
        )
        return ker / np.sum(ker, axis=1, keepdims=True)
    
    @staticmethod
    def get_parabolic_spline(t, t_refs):
        ker = np.zeros((t.shape[0], t_refs.shape[0]))
        ker[t < t_refs[0], 0] = 1
        for idx in range(t_refs.shape[0] - 1):
            valid_idxes = (t >= t_refs[idx]) & (t < t_refs[idx + 1])
            norm = t_refs[idx + 1] - t_refs[idx]
            ker[valid_idxes, idx] = 0.75 * (1 - ((t[valid_idxes] - t_refs[idx])/norm)**2)
            ker[valid_idxes, idx + 1] = 0.75 *(1 - ((t[valid_idxes] - t_refs[idx + 1])/norm)**2)
        ker[t >= t_refs[-1], -1] = 1
        return ker / np.sum(ker, axis=1, keepdims=True)
    
class Seasonal:
    @staticmethod
    def get_fourier_features(dates, n_order):
        periods = dates.day_of_year / 365.25
        return pd.DataFrame(
            {
                f"{func}_order_{order}": getattr(np, func)(2 * np.pi * periods * order)
                for order in range(1, n_order + 1)
                for func in ("sin", "cos")
            }
        )

Trend Components

class LinearTrend(TrendComponent):
    _name = "Trend.Linear"
    def __init__(self, t, **kwargs):
        super().__init__(self._name, t)
        zero_slope = getattr(kwargs, "zero_slope", False)
        zero_intercept = getattr(kwargs, "zero_intercept", False)
        
        self._slope = 0 if zero_slope else pm.Normal("slope", mu=0, sigma=1)
        self._intercept = 0 if zero_intercept else pm.Normal("intercept", mu=0, sigma=1)
        self._trend = pm.Deterministic("trend", var=(self._slope * self._t + self._intercept))
        
class LogisticTrend(TrendComponent):
    _name = "Trend.Logistic"
    def __init__(self, t, **kwargs):
        super().__init__(self._name, t)
        
        self._max_L = pm.HalfNormal("max_L", sigma=1)
        self._growth_rate = pm.Normal("growth_rate", mu=0, sigma=1)
        self._x_offset = pm.Normal("x_offset", mu=0, sigma=1)
        self._y_offset = pm.Normal("y_offset", mu=0, sigma=1)
        
        self._trend = pm.Deterministic("trend", var=(
            self._max_L/(1 + pm.math.exp(-self._growth_rate * (self._t - self._x_offset))) + self._y_offset
        ))

class SplineTrend(TrendComponent):
    _name = "Trend.Spline"
    def __init__(self, t, checkpoints, spline: Spline.SplineType = Spline.SplineType.LINEAR, **kwargs):
        super().__init__(self._name, t)
        if type(checkpoints) is int:
            checkpoints = np.linspace(t.min(), t.max(), checkpoints + 2)[1:-1]
        spline_kernel = Spline.get_spline(t, checkpoints, spline, **kwargs)
        
        self._sigma = pm.HalfNormal("sigma", sigma=1)
        self._anchors = pt.cumsum(pm.Laplace("anchors", mu=0, b=self._sigma, shape=(len(checkpoints), )))
        self._trend = pm.Deterministic("trend", var=(pm.math.dot(spline_kernel, self._anchors)))

Seasonality Components

class ConstantSeasonality(SeasonalityComponent):
    _name = "Seasonality.Constant"
    def __init__(self, fourier_features, **kwargs):
        super().__init__(self._name, fourier_features)
        size = fourier_features.shape[1]
        self._coeff = pm.Normal("coeff", mu=0, sigma=1, shape=(size,))
        self._seasonality = pm.Deterministic("seasonality", var=(self._coeff * self._fourier_features).sum(axis=1))

class SplineSeasonality(SeasonalityComponent):
    _name = "Seasonality.Spline"
    def __init__(
        self,
        fourier_features,
        t, checkpoints,
        spline: Spline.SplineType = Spline.SplineType.LINEAR,
        **kwargs
    ):
        super().__init__(self._name, fourier_features)
        if type(checkpoints) is int:
            checkpoints = np.linspace(t.min(), t.max(), checkpoints + 2)[1:-1]
        spline_kernel = Spline.get_spline(t, checkpoints, spline, **kwargs)
        size = fourier_features.shape[1]
        
        self._sigma = pm.HalfNormal("sigma", sigma=1)
        self._anchors = pt.cumsum(pm.Laplace("anchors", mu=0, b=self._sigma, shape=(len(checkpoints), size)), axis=0)
        self._seasonality = pm.Deterministic(
            "seasonality", var=(self._fourier_features * pm.math.dot(spline_kernel, self._anchors)).sum(axis=1)
        )

Construct model

dates = pd.date_range("2020-01-01", "2021-01-01")
t = np.linspace(0, 1, len(dates))
fourier_feats = Seasonal.get_fourier_features(dates, 3)

with pm.Model() as model:
    # trend = LinearTrend(t)
    # trend = LogisticTrend(t)
    trend = SplineTrend(t, 10, Spline.SplineType.LINEAR)
    
    # seasonality = ConstantSeasonality(fourier_feats)
    seasonality = SplineSeasonality(fourier_feats, t, 10)
    
    mu = pm.Deterministic("mu", var=(trend._trend + seasonality._seasonality))

display(model)
display(pm.model_to_graphviz(model))

It is now a lot easier to swap trend, and seasonality component. For uniformity, each specific component type should have the same output rv name such as trend component must have self._trend, and seasonality component must have self._seasonality. However this probably going to affect MMM diagnosis graph ploting in a major way. So I think it this require some discussion such as a standard interface, or data structure

[twiecki]

Damn, that's very cool! I can also see how time-varying parameters could easily fit into this framework.

[twiecki]

I also wonder if bambi could help in constructing more flexible MMMs, but not sure that's compatible with what you're talking about here.

[thipokKub]

I think bart can be apply as a a time varying coefficient for media, and controls

But I’m a bit skeptical on the control variables, as the flexibility might negate the influence from media variables

[BBDS-Colt]

I really like this modular, OOP proposal! Could go a long way towards making the MMM model more flexible & customizable.

[twiecki]

@thipokKub Seems like feedback to your proposal is quite positive, want to do a PR?

[thipokKub]

Sure! I probably going to open a pr on the weekends though

[geirrlod]

Do anyone know if this is moving forward? Seems like an excellent idea. Or if there is some other plans to make it a bit easier to customize than today.

[twiecki]

Haven't heard from @thipokKub. Do you want to take this up? Would be a great contribution.

[thipokKub]

Sorry for late reply, I did have a forked a few months ago. But due to the company I worked for shifted their focus, I no longer have time to develop it. But if you want my code I could push a pr

[twiecki]

@thipokKub Yes, any hacky/drafty code would be welcome!

[wd60622]

I'd be really interested in seeing this! and would be willing to help out here

I think that this should fit into the current classes but might be a good start to break down the workload.

For instance, refactoring and extracting the different contributions

• channel contributions
• control
• Seasonality contribution

def get_fourier_contributions(X_fourier: np.ndarray, gamma_fourier: Dict[str, Any]) -> pt.TensorVariable:
    """Seasonal contribution. Must be defined within a model"""
    fourier_data_ = pm.MutableData(
        name="fourier_data",
        value=X_fourier, 
        dims=("date", "fourier_mode"),
    )
    
    gamma_fourier = pm.Laplace(
        name="gamma_fourier",
        mu=gamma_fourier["mu"],
        b=gamma_fourier["b"],
        dims=gamma_fourier["dims"],
    )
    
    fourier_contribution = pm.Deterministic(
        name="fourier_contributions",
        var=fourier_data_ * gamma_fourier,
        dims=("date", "fourier_mode"),
    )
    
def get_channel_contributions(
    X_channels: np.ndarray, adstock_max_lag, beta_channel_config, alpha_config, lam_config
) -> pt.TensorVariable: 
    """Also extracted out as well."""
    ...

# Hopefully making the model definition block  
coords = {"date": ..., "fourier_mode": ..., "channel": ...}
with pm.Model(coords=coords) as custom_mmm: 
    periods: npt.NDArray[np.float_] = date_data.dt.dayofyear.to_numpy() / 365.25
    X_fourier = generate_fourier_modes(
        periods=periods,
        n_order=yearly_seasonality,
    )
    seasonal_contribution = get_fourier_contribution(X_fourier, gamma_fourier=...)

    channel_contribution = get_channel_contribution(X_channel, ...)
    
    mu = seasonal_contribution.sum(axis=1) + channel_contribution.sum(axis=1)
    
    # Freedom to change the likelihood, more transformations, customization
    geographical_hierarchical_intercept  = ...

But rely on these in the current class in order keep an easy API for out of the box

[twiecki]

That's great! @ferrine also developed a times-series API recently that allowed for elegant combination of higher-level building blocks. Perhaps there's something we can learn from that too.

[ferrine]

Hi, on a high level, it is beneficial to use ideas from https://arxiv.org/abs/2208.07132 to compose the building blocks that maintain a good prior predictive for the organic component. Organic consists of Seasonal terms, Regression, Trend, and HSGP or HSTP to model outliers and trend violations. In my particular scenario, this is still WIP, but for this simpler case I imagine the model should take form of

log sales_t ~ log(exp(organic(t, X_t)) + marketing(t)) + eps_t

Where organic has multiplicative effect, marketing has additive effect, but the total likelihood is on the log scale.
For discrete sales, Negative Binomial is also an option but should follow the same idea of being semi additive.

The more complicated model is what I'm working on right now (the causal graph can be found in the very bottom of this blogpost). Unfortunately, no manual specification is possible anymore, so the clear separation of the components was essential.

This simple model only has Direct (what comes from the marketing spend on the lowest funnel) and organic effects and can be rewritten as

log sales_t ~ log(exp(organic(t, X_t)) + direct(t)) + eps_t

But the larger model also has more than this and includes transitive effects, e.g., "invest into brand awareness, attract more leads to the lower funnel". It would look like this:

log sales_t ~ log(exp(organic(t, transitive(Z_t), X_t)) + direct(t)) + eps_t

or

log sales_t ~ log(exp(organic(t, X_t)) + direct(t) + transitive(Z_t)) + eps_t

Given the discoveries, I believe pymc-mareting can focus on the simple case with only the lowest funnel observations (sales) and the model is the basic one

log sales_t ~ log(exp(organic(t, X_t)) + direct(t)) + eps_t

The only issue is with the geo effect. One difficulty I found, it felt to me that restricting all geos to have the same number of observations might be too demanding. What do you think of it? So far I'm more towars to model geos separately:

log sales_gt ~ log(exp(organic_g(t, X_gt)) + direct_g(t)) + eps_gt

However, some direct investments into channels might cover all geos, so there are basically two types of marketing effects: "geo specific" and "world". Global effect marketing spend should still somehow have geo specific effects, but the functional form is different.

log sales_gt ~ log(exp(organic_g(t, X_gt)) + direct_g(t) + direct_w_g(t)) + eps_gt

[ricardoV94]

Related to #274