Vectorize adstock transformations

[juanitorduz]

In this PR we add a vectorized version of the adstoc transformations to avoid writing for loops when going over channels transformations. For example, for 3 channels and 169 observation we can write:

coords = {"channels": ["ch1", "ch2", "ch3"]}


with pm.Model(coords=coords) as model:
    # --- Priors ---
    intercept = pm.Normal(name="intercept", mu=0, sd=2)
    beta = pm.HalfNormal(name="beta", sigma=2, dims="channels")
    alpha = pm.Beta(name="alpha", alpha=1, beta=1, dims="channels")
    lam = pm.HalfNormal("lam", sigma=2, dims="channels")
    sigma = pm.HalfNormal(name="sigma", sd=2)

    # --- Model Parametrization ---
    channels_adstock = geometric_adstock_vectorized(
        x=scaled_channels_df, alpha=alpha, normalize=True
    )
    channels_adstock_saturated = logistic_saturation(x=channels_adstock, lam=lam)
    mu = intercept + pm.math.dot(channels_adstock_saturated, beta)

    # --- Likelihood ---
    likelihood = pm.Normal("likelihood", mu=mu, sigma=sigma, observed=scaled_y)

pm.model_to_graphviz(model)

[juanitorduz]

In 9257d0b I added an experimental implementation of the geometric adstock function using aesara.scan to try to omit for loops. Note that this one works for both scalars and vectors (see tests). Feedback is welcome!

[juanitorduz]

Well ... I really tried my best to write a geometric adstock function which was able to handle scalar and vector values, could accept the l_max argument and to support normalisation.... the geometric_adstock_scan does all of it BUT is super slow and not really usable in practice.

In my current project I have used the geometric_adstock_vectorized and it works super nice! The only drawback is that does not work for scalars but I guess that is ok as we have the initial implementation.

Hence, we could simply remove geometric_adstock_scan or put it in an experimental phase to see if we can make it work. After this decision I will simple hope to merge this PR so that we can test and use geometric_adstock_vectorized.

I am sorry about the d-tour with the scan approach but I just wanted to try it out.

[juanitorduz]

In 462b69d I removed the implementation using scan.

[juanitorduz]

In 9db1d14 and 434a18d I organised the test of the transformers module better. This is ready for review! 🚀

[ricardoV94]

Looks good. Just left two small comments regarding the tests