Seasonality in time varying media coefficients

[tobia-cassandra]

Hello all!

I was digging a little bit deeper inside the time_varying_media functionality lately and I was thinking about changing the covariance kernel to one that accounts for seasonal patterns (like the Periodic one) and combining it with the default Matern52, so I wrote this very simple code to test it out:

class TrendSeasonHoliday(Stationary):
    def __init__(self, period, holidays_idx, ls=2, ls_inv=None, **kwargs):
        super().__init__(input_dim=1, ls=ls, ls_inv=ls_inv, **kwargs)
        self.k_trend = Matern52(input_dim=1, ls=ls)
        self.k_seas  = Periodic(input_dim=1, period=period, ls=ls)
        self.holidays_idx = holidays_idx

    def full(self, X, Xs=None):
        Xs = X if Xs is None else Xs
        K  = self.k_trend.full(X, Xs) + self.k_seas.full(X, Xs)
        # add diagonal spikes for holidays
        for idx in self.holidays_idx:
            K = pm.math.set_subtensor(K[idx, idx], K[idx, idx] + 1)
        return K
    
    def to_dict(self):
        return {
            "class_name": "TrendSeasonHoliday",
            "period": self.period,
            "holidays_idx": self.holidays_idx,
            "ls": float(self.ls) if hasattr(self.ls, '__float__') else self.ls,
            "input_dim": self.input_dim
        }
    
    def __dict__(self):
        return self.to_dict()


hsgp_cfg = HSGPKwargs(cov_func=TrendSeasonHoliday(period=52, holidays_idx=holiday_locs))
model_cfg = {"media_tvp_config": hsgp_cfg}


mmm = PyMCModel(
    date_column="date_week",
    channel_columns=media_channels,
    control_columns=controls,
    calibrate_prior_from_contribution=False,
    calibration_data=data,
    calibration_output_var="revenue",
    time_varying_media=True, 
    model_config=model_cfg
)
_ = mmm.fit(X=X, y=y, chains=4, target_accept=0.90, nuts_sampler="numpyro")

Here, PyMCModel is just a wrapper that I built, but the fit method is absolutely the same.

Basically, the fit works fine but then, after it, I got a PydanticSerializationError and I basically discovered that we can only use kernels included in the enum CovFunc.

What I would like to ask is: is there a way to bypass this behavior and use a custom kernel for TVP?

Thank you all!