Media Mix Model and Experimental Calibration: A Simulation Study

Causal Model

It is very important to view media mix models from a causal perspective. We can not simply add a bunch of variables and hope for the best! See for example Using Data Science for Bad Decision-Making: A Case Study. There is a systematic way to select the variables we should / should not add into the model. To see this how this works in practice, let’s use DoWhy to get the set of variables for our model. First we need to define the DAG (this is where domain knowledge from the marketing mechanism comes into place!).

gml_graph = """
graph [
    directed 1

    node [
        id x1
        label "x1"
    ]
    node [
        id x2
        label "x2"
    ]
    node [
        id z
        label "z"
    ]
    node [
        id seasonality
        label "seasonality"
    ]
    node [
        id trend
        label "trend"
    ]
    node [
        id y
        label "y"
    ]
    edge [
        source seasonality
        target x1
    ]
    edge [
        source z
        target x1
    ]
    edge [
        source seasonality
        target y
    ]
    edge [
        source trend
        target y
    ]
    edge [
        source z
        target y
    ]
    edge [
        source x1
        target y
    ]
    edge [
        source x2
        target y
    ]
]
"""

We want to see if we can find a set of variables which we can put in a regression model to estimate the effect of \(x_1\) and \(x_2\) on \(y\). Hence, we define to causal models:

causal_model_1 = CausalModel(
    data=model_df,
    graph=gml_graph,
    treatment="x1",
    outcome="y",
)

causal_model_2 = CausalModel(
    data=model_df,
    graph=gml_graph,
    treatment="x2",
    outcome="y",
)

causal_model_1.view_model()

We can now query for backdoor paths between the treatment and the outcome:

causal_model_1._graph.get_backdoor_paths(nodes1=["x1"], nodes2=["y"])

[['x1', 'z', 'y'], ['x1', 'seasonality', 'y']]

causal_model_2._graph.get_backdoor_paths(nodes1=["x2"], nodes2=["y"])

[]

Note that for \(x_1\) we need to add \(z\) and seasonality, otherwise we would have open paths:

causal_model_1._graph.check_valid_backdoor_set(nodes1=["x1"], nodes2=["y"], nodes3=[])

{'is_dseparated': False}

Therefore the set of variables we should add to the model are \(x_1\), \(x_2\), \(z\), seasonality and trend (the trend component will allow us to reduce the variance on the estimate).

causal_model_1._graph.check_valid_backdoor_set(
    nodes1=["x1"], nodes2=["y"], nodes3=["z", "seasonality", "trend"]
)

{'is_dseparated': True}

causal_model_1._graph.check_valid_backdoor_set(
    nodes1=["x2"], nodes2=["y"], nodes3=["z", "seasonality", "trend"]
)

{'is_dseparated': True}

Prior Specification

Bayesian models provide great flexibility to add domain knowledge into the model through the prior distributions. PreliZ is a great library to study and set priors.

For example, what should be the prior for the adstock parameter \(\alpha\)? We know it should be between zero and one. We also do not expect the carry over to be very strong. Hence, we can use a Beta distribution with the mass a bit towards zero.

fig, ax = plt.subplots(figsize=(10, 6))
pz.Beta(alpha=2, beta=3).plot_pdf(ax=ax)
ax.set(xlabel=r"$\alpha$")
ax.set_title(label=r"Prior Distribution $\alpha$", fontsize=18, fontweight="bold")

For \(\lambda\) we could also argue similarly: it is a positive number and we expect a mild saturation. Hence, we can use a Gamma distribution with mean close to one. Still, we make sure it is not concentrated to much around one in order to allow for some flexibility.

fig, ax = plt.subplots(figsize=(10, 6))
pz.Gamma(alpha=2, beta=2).plot_pdf(ax=ax)
ax.set(xlabel=r"$\lambda$")
ax.set_title(label=r"Prior Distribution $\lambda$", fontsize=18, fontweight="bold")

A much harder choice for prior are the regression coefficients of the media variables. Without any prior information on channel performance a common choice is to use a weakly informative prior driven by the investment share. The idea motivation behind it is that we expect (without seeing the data!) that channels where we spend the most to contribute more to sales. This can actually not be the case once we see the data. Nevertheless, is a good guidance if you do not have any other information. One could put the same prior for all channels, but in practice I have seen in practice that, given the small sample sizes, smaller channels can get very high contributions just because of the noise (I have also tested via simulations). It is important to emphasize that there is no right answer here. It is a matter of what you believe and what you want to test.

For this example, let’s use such a prior. In addition, note we expect the effect to be always positive. Hence, we can use a half-normal distribution to enforce this.

channel_share = data_df[channels].sum() / data_df[channels].sum().sum()

fig, ax = plt.subplots(figsize=(10, 6))
sns.barplot(
    x=channel_share.index, y=channel_share.values, hue=channel_share.index, ax=ax
)
ax.yaxis.set_major_formatter(mtick.PercentFormatter(xmax=1))
ax.set(xlabel="channel", ylabel="share")
ax.set_title(label="Channel Spend Share", fontsize=18, fontweight="bold")

fig, ax = plt.subplots(figsize=(10, 6))
pz.HalfNormal(sigma=channel_share.loc["x1"]).plot_pdf(ax=ax)
pz.HalfNormal(sigma=channel_share.loc["x2"]).plot_pdf(ax=ax)
ax.set(xlabel=r"$\beta$")
ax.set_title(label=r"Prior Distribution $\beta$", fontsize=18, fontweight="bold")

Model Specification

We are now ready to specify the model. We will not go into details as we are implementing the classical model from “Media Mix Model Calibration With Bayesian Priors”, by Zhang, et al., which is described in detail in the the previous post Media Effect Estimation with PyMC: Adstock, Saturation & Diminishing Returns.

The only difference it we do not add an explicit trend variable but we use a Gaussian process to model this latent variable (therefore we also remove an intercept) as described in the post Estimating the Long-Term Base in Marketing-Mix Models. We are using the same technique as described in the example Time Series Modeling with HSGP: Baby Births Example.

coords = {"date": date, "channel": channels, "fourier_mode": np.arange(2 * n_order)}

with pm.Model(coords=coords) as model_1:
    # --- Data Containers ---

    index_scaled_data = pm.Data(
        name="index_scaled",
        value=index_scaled.to_numpy().flatten(),
        mutable=True,
        dims="date",
    )

    channels_scaled_data = pm.Data(
        name="channels_scaled",
        value=channels_scaled,
        mutable=True,
        dims=("date", "channel"),
    )

    z_scaled_data = pm.Data(
        name="z_scaled",
        value=controls_scaled.to_numpy().flatten(),
        mutable=True,
        dims="date",
    )

    fourier_features_data = pm.Data(
        name="fourier_features",
        value=fourier_features,
        mutable=True,
        dims=("date", "fourier_mode"),
    )

    y_scaled_data = pm.Data(
        name="y_scaled",
        value=target_scaled.to_numpy().flatten(),
        mutable=True,
        dims="date",
    )

    # --- Priors ---

    amplitude_trend = pm.HalfNormal(name="amplitude_trend", sigma=1)
    ls_trend_params = pm.find_constrained_prior(
        distribution=pm.InverseGamma,
        lower=0.1,
        upper=0.9,
        init_guess={"alpha": 3, "beta": 1},
        mass=0.95,
    )
    ls_trend = pm.InverseGamma(name="ls_trend", **ls_trend_params)

    alpha = pm.Beta(name="alpha", alpha=2, beta=3, dims="channel")
    lam = pm.Gamma(name="lam", alpha=2, beta=2, dims="channel")
    beta_channel = pm.HalfNormal(
        name="beta_channel",
        sigma=channel_share.to_numpy(),
        dims="channel",
    )
    beta_z = pm.Normal(name="beta_z", mu=0, sigma=1)
    beta_fourier = pm.Normal(name="beta_fourier", mu=0, sigma=1, dims="fourier_mode")

    sigma = pm.HalfNormal(name="sigma", sigma=1)
    nu = pm.Gamma(name="nu", alpha=25, beta=2)

    # --- Parametrization ---

    cov_trend = amplitude_trend * pm.gp.cov.ExpQuad(input_dim=1, ls=ls_trend)
    gp_trend = pm.gp.HSGP(m=[20], c=1.5, cov_func=cov_trend)
    f_trend = gp_trend.prior(name="f_trend", X=index_scaled_data[:, None], dims="date")

    channel_adstock = pm.Deterministic(
        name="channel_adstock",
        var=geometric_adstock(
            x=channels_scaled_data,
            alpha=alpha,
            l_max=l_max,
            normalize=True,
        ),
        dims=("date", "channel"),
    )
    channel_adstock_saturated = pm.Deterministic(
        name="channel_adstock_saturated",
        var=logistic_saturation(x=channel_adstock, lam=lam),
        dims=("date", "channel"),
    )
    channel_contributions = pm.Deterministic(
        name="channel_contributions",
        var=channel_adstock_saturated * beta_channel,
        dims=("date", "channel"),
    )

    fourier_contribution = pm.Deterministic(
        name="fourier_contribution",
        var=pt.dot(fourier_features_data, beta_fourier),
        dims="date",
    )

    z_contribution = pm.Deterministic(
        name="z_contribution",
        var=z_scaled_data * beta_z,
        dims="date",
    )

    mu = pm.Deterministic(
        name="mu",
        var=f_trend
        + channel_contributions.sum(axis=-1)
        + z_contribution
        + fourier_contribution,
        dims="date",
    )

    # --- Likelihood ---
    y = pm.StudentT(
        name="y",
        nu=nu,
        mu=mu,
        sigma=sigma,
        observed=y_scaled_data,
        dims="date",
    )

pm.model_to_graphviz(model=model_1)

Prior Predictive Checks

Before fitting the model, it is always a good idea to check the prior predictive distribution. This is a way to see if the model is able to generate data that is consistent with our prior beliefs. We can do this by sampling from the prior and plotting the results.

with model_1:
    prior_predictive_1 = pm.sample_prior_predictive(samples=2_000, random_seed=rng)

fig, ax = plt.subplots()
az.plot_ppc(data=prior_predictive_1, group="prior", kind="kde", ax=ax)
ax.set_title(label="Prior Predictive - Model 1", fontsize=18, fontweight="bold")

It looks quite reasonable.

Model Fitting

We now fit the model using the No-U-Turn Sampler (NUTS) algorithm. We will run four chains in parallel and sample \(2000\) draws from each chain. We use NumPyro as the backend for the model fitting (it is super fast).

with model_1:
    idata_1 = pm.sample(
        target_accept=0.9,
        draws=2_000,
        chains=4,
        nuts_sampler="numpyro",
        random_seed=rng,
    )
    posterior_predictive_1 = pm.sample_posterior_predictive(
        trace=idata_1, random_seed=rng
    )

Model Diagnostics

First we verify we do not have divergences:

idata_1["sample_stats"]["diverging"].sum().item()

\(\displaystyle 0\)

Next, we look into the posterior distribution of the parameters.

var_names = [
    "amplitude_trend",
    "ls_trend",
    "alpha",
    "lam",
    "beta_channel",
    "beta_z",
    "beta_fourier",
    "sigma",
    "nu",
]

az.summary(data=idata_1, var_names=var_names, round_to=3)

axes = az.plot_trace(
    data=idata_1,
    var_names=var_names,
    compact=True,
    kind="rank_bars",
    backend_kwargs={"figsize": (15, 17), "layout": "constrained"},
)
plt.gcf().suptitle("Trace - Model 1", fontsize=16)

The trace looks fine!

Posterior Predictive Checks

We now want to see the model fit. We want to do it in the original scale. Therefore we need to scale our posterior samples back.

pp_vars_original_scale_1 = {
    var_name: xr.apply_ufunc(
        target_scaler.inverse_transform,
        idata_1["posterior"][var_name].expand_dims(dim={"_": 1}, axis=-1),
        input_core_dims=[["date", "_"]],
        output_core_dims=[["date", "_"]],
        vectorize=True,
    ).squeeze(dim="_")
    for var_name in ["mu", "f_trend", "fourier_contribution", "z_contribution"]
}

pp_vars_original_scale_1["channel_contributions"] = xr.apply_ufunc(
    target_scaler.inverse_transform,
    idata_1["posterior"]["channel_contributions"],
    input_core_dims=[["date", "channel"]],
    output_core_dims=[["date", "channel"]],
    vectorize=True,
)

# We define the baseline as f_trand + fourier_contribution + z_contribution
pp_vars_original_scale_1["baseline"] = xr.apply_ufunc(
    target_scaler.inverse_transform,
    (
        idata_1["posterior"]["f_trend"]
        + idata_1["posterior"]["fourier_contribution"]
        + idata_1["posterior"]["z_contribution"]
    ).expand_dims(dim={"_": 1}, axis=-1),
    input_core_dims=[["date", "_"]],
    output_core_dims=[["date", "_"]],
    vectorize=True,
).squeeze(dim="_")

We also scale back the likelihood.

pp_likelihood_original_scale_1 = xr.apply_ufunc(
    target_scaler.inverse_transform,
    posterior_predictive_1["posterior_predictive"]["y"].expand_dims(
        dim={"_": 1}, axis=-1
    ),
    input_core_dims=[["date", "_"]],
    output_core_dims=[["date", "_"]],
    vectorize=True,
).squeeze(dim="_")

We can now visualize the results:

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_likelihood_original_scale_1,
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.1, "label": r"likelihood $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["mu"],
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.3, "label": r"$\mu$ $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper left")
ax.set_title(label="Posterior Predictive - Model 1", fontsize=18, fontweight="bold")

The model does explain the data variance quite well. Let’s look into the distribution:

fig, ax = plt.subplots()
az.plot_ppc(
    data=posterior_predictive_1,
    num_pp_samples=1_000,
    observed_rug=True,
    random_seed=seed,
    ax=ax,
)
ax.set_title(label="Posterior Predictive - Model 1", fontsize=18, fontweight="bold")

Looks good!

Model Components

We now visualize the model components.

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.5, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.5, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["f_trend"],
    hdi_prob=0.94,
    color="C2",
    fill_kwargs={"alpha": 0.3, "label": r"$f_\text{trend}$ $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["fourier_contribution"],
    hdi_prob=0.94,
    color="C3",
    fill_kwargs={"alpha": 0.3, "label": r"seasonality $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["z_contribution"],
    hdi_prob=0.94,
    color="C4",
    fill_kwargs={"alpha": 0.3, "label": r"$z$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set_title(label="Components Contributions - Model 1", fontsize=18, fontweight="bold")

Remarks:

• The Gaussian process is able to capture the trend component quite well.
• Note that the variance of the \(x_1\) estimate is much wider than the one for \(x_2\). This is due to the fact \(x_1\) has a high correlation with the seasonality component. Multicollinearity is a common issue in media mix models.

We now plot the baseline against the media contribution:

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.5, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.5, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["baseline"],
    hdi_prob=0.94,
    color="C7",
    fill_kwargs={"alpha": 0.5, "label": r"baseline $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper left")
ax.set_title(label="Baseline vs Channels - Model 1", fontsize=18, fontweight="bold")

Let’s now deep dive into the media contributions:

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)
sns.lineplot(
    x="date",
    y="x1_effect",
    data=data_df.assign(x1_effect=lambda x: amplitude * x["x1_effect"]),
    color="C0",
    label=r"$x1$ effect",
    ax=ax[0],
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.3, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax[0],
)
ax[0].legend(loc="upper right")

sns.lineplot(
    x="date",
    y="x2_effect",
    data=data_df.assign(x2_effect=lambda x: amplitude * x["x2_effect"]),
    color="C1",
    label=r"$x2$ effect",
    ax=ax[1],
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_1["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.3, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax[1],
)
ax[1].legend(loc="upper right")

fig.suptitle("Channel Contributions - Model 1", fontsize=18, fontweight="bold")

Overall, the model recovered the true effects. Still, the estimate of \(x_2\) is slightly underestimated by the model (the true value is not in the middle of the HDI).

Media Parameters

We now compare the estimated against the true values for the parameters \(\alpha\), \(\lambda\) and the regression coefficients \(\beta\). Recall that since we are working on the scaled space we need to scale \(\lambda\) and \(\beta\) back:

alpha_posterior_1 = idata_1["posterior"]["alpha"]
lam_posterior_1 = idata_1["posterior"]["lam"] * channels_scaler.scale_
beta_channel_posterior_1 = (
    idata_1["posterior"]["beta_channel"]
    * (1 / channels_scaler.scale_)
    * (target_scaler.scale_)
    * (1 / amplitude)
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=alpha_posterior_1, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(
    x=alpha1, color="C0", linestyle="--", linewidth=2, label=r"$\alpha_1$ (true)"
)
ax.axvline(
    x=alpha2, color="C1", linestyle="--", linewidth=2, label=r"$\alpha_2$ (true)"
)
ax.legend(loc="upper left")
ax.set_title(
    label=r"Adstock Parameter $\alpha$ ($94\%$ HDI) - Model 1",
    fontsize=18,
    fontweight="bold",
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=lam_posterior_1, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(x=lam1, color="C0", linestyle="--", linewidth=2, label=r"$\lambda_1$ (true)")
ax.axvline(x=lam2, color="C1", linestyle="--", linewidth=2, label=r"$\lambda_2$ (true)")
ax.legend(loc="upper right")
ax.set_title(
    label=r"Saturation Parameter $\lambda$ ($94\%$ HDI) - Model 1",
    fontsize=18,
    fontweight="bold",
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=beta_channel_posterior_1, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(x=beta1, color="C0", linestyle="--", linewidth=2, label=r"$\beta_1$ (true)")
ax.axvline(x=beta2, color="C1", linestyle="--", linewidth=2, label=r"$\beta_2$ (true)")
ax.legend(loc="upper right")
ax.set_title(
    label=r"Channel $\beta$ ($94\%$ HDI) - Model 1", fontsize=18, fontweight="bold"
)

Remarks:

• We wer able to recover the true values of the carryover parameter \(\alpha\).
• For \(x_1\), the true values of \(\lambda\) and \(\beta\) are in the \(94\%\) HDI.
• The \(94\%\) HDI of the regression coefficient \(\beta\) for \(x_1\) is much wider than the one for \(x_2\). As mentioned before, this is probably related with multicollinearity with the seasonal component.
• For \(x_2\) the model overestimated the saturation parameter \(\lambda\) and underestimated the regression coefficient \(\beta\). Maybe we should have had stricter priors? This is a true challenge in practice as we do not know the true values. This is a motivation for this notebook! Being able to easily pass information through ROAS than beta coefficients in the scaled space is a much more convenient way to pass information to the model.

ROAS Estimation

Once the model is fitted, we can use the posterior samples to compute the ROAS as in above. We simply need to generate predictions when setting the media data to zero. Let’s start with the global ROAS:

predictions_roas_data_all_1 = {"x1": {}, "x2": {}}

for channel in tqdm(["x1", "x2"]):
    with model_1:
        pm.set_data(
            new_data={
                "channels_scaled": channels_scaler.transform(
                    data_df[channels].assign(**{channel: 0})
                )
            }
        )
        predictions_roas_data_all_1[channel] = pm.sample_posterior_predictive(
            trace=idata_1, var_names=["y", "mu"], progressbar=False, random_seed=rng
        )

We now scale the predictions back to the original scale:

predictions_roas_data_scaled_diff_all_1 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_data_scaled_diff_all_1[channel] = pp_vars_original_scale_1[
        "mu"
    ] - xr.apply_ufunc(
        target_scaler.inverse_transform,
        predictions_roas_data_all_1[channel]["posterior_predictive"]["mu"].expand_dims(
            dim={"_": 1}, axis=-1
        ),
        input_core_dims=[["date", "_"]],
        output_core_dims=[["date", "_"]],
        vectorize=True,
    ).squeeze(
        dim="_"
    )

We can now visualize the difference iin the potential outcomes:

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_all_1["x1"],
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.3, "label": r"$E(y - y_{01})$ $94\%$ HDI"},
    ax=ax[0],
)
sns.lineplot(
    data=data_df.assign(x1_effect=lambda x: amplitude * x["x1_effect"]),
    x="date",
    y="x1_effect",
    color="C0",
    label=r"$x1$ effect",
    ax=ax[0],
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_all_1["x2"],
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.3, "label": r"$E(y - y_{02})$ $94\%$ HDI"},
    ax=ax[1],
)
sns.lineplot(
    data=data_df.assign(x2_effect=lambda x: amplitude * x["x2_effect"]),
    x="date",
    y="x2_effect",
    color="C1",
    label=r"$x2$ effect",
    ax=ax[1],
)
ax[0].legend(loc="upper left")
ax[1].legend(loc="upper left")
ax[0].set(ylabel=None)
ax[1].set(ylabel=None)
fig.suptitle("Potential Outcomes Difference - Model 1", fontsize=18, fontweight="bold")

It is not surprise that this is precisely the same plot as the one we saw before. Namely, the media contributions from the model 🙃.

We now compute the ROAS:

predictions_roas_all_1 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_all_1[channel] = (
        predictions_roas_data_scaled_diff_all_1[channel].sum(dim="date")
        / data_df[channel].sum()
    )

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)

az.plot_posterior(predictions_roas_all_1["x1"], ref_val=roas_all_1, ax=ax[0])
ax[0].set(title="x1")
az.plot_posterior(predictions_roas_all_1["x2"], ref_val=roas_all_2, ax=ax[1])
ax[1].set(title="x2", xlabel="ROAS")
fig.suptitle("Global ROAS - Model 1", fontsize=18, fontweight="bold")

We can proceed in an analogous way for the quarterly ROAS. This time we need to loop ever time series on which we set up the quarterly spend to zero for a given channel.

predictions_roas_data_1 = {"x1": {}, "x2": {}}

for channel in tqdm(["x1", "x2"]):
    for q in tqdm(roas_data[channel].keys()):
        with model_1:
            pm.set_data(
                new_data={
                    "channels_scaled": channels_scaler.transform(
                        roas_data[channel][q][channels]
                    )
                }
            )
            predictions_roas_data_1[channel][q] = pm.sample_posterior_predictive(
                trace=idata_1, var_names=["y", "mu"], progressbar=False, random_seed=rng
            )

predictions_roas_data_scaled_diff_1 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_data_scaled_diff_1[channel] = {
        q: pp_vars_original_scale_1["mu"]
        - xr.apply_ufunc(
            target_scaler.inverse_transform,
            idata_q["posterior_predictive"]["mu"].expand_dims(dim={"_": 1}, axis=-1),
            input_core_dims=[["date", "_"]],
            output_core_dims=[["date", "_"]],
            vectorize=True,
        ).squeeze(dim="_")
        for q, idata_q in tqdm(predictions_roas_data_1[channel].items())
    }

In order to have a better understanding on this computation, we can plot one of such time series differences:

fig, ax = plt.subplots()
ax.axvspan(
    xmin=quarter_ranges.loc[quarter_ranges.quarter == "2022Q3"]["date"]["min"].item(),
    xmax=quarter_ranges.loc[quarter_ranges.quarter == "2022Q3"]["date"]["max"].item(),
    color="C1",
    label="2022-Q3",
    alpha=0.2,
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_1["x2"]["2022Q3"],
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.7, "label": r"$x2$ difference against actuals $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper right")
ax.set_title(
    label=r"Potential Outcomes Difference - Model 1 (Channel $x2$ (2023-Q3))",
    fontsize=18,
    fontweight="bold",
)

Observe the difference is non-zero just after the selected quarter. This is due to the carryover effect of the media spend.

We now divide the differences by the spend to get the ROAS:

predictions_roas_1 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_1[channel] = {
        q: idata_q.sum(dim="date") / data_df.query("quarter == @q")[channel].sum()
        for q, idata_q in predictions_roas_data_scaled_diff_1[channel].items()
    }

Here are the results for \(x_1\) and compare them with the true values:

x1_ref_vals = (
    roas.query("channel == 'x1'")
    .drop("channel", axis=1)
    .rename(columns={"roas": "ref_val"})
    .to_dict(orient="records")
)

axes = az.plot_posterior(
    data=xr.concat(
        predictions_roas_1["x1"].values(),
        dim=pd.Index(data=predictions_roas_1["x1"].keys(), name="quarter"),
    ),
    figsize=(18, 15),
    grid=(5, 2),
    ref_val={"mu": x1_ref_vals},
    backend_kwargs={"sharex": True, "layout": "constrained"},
)
fig = axes.ravel()[0].get_figure()
fig.suptitle(
    t="Posterior Predictive - ROAS x1 - Model 1", fontsize=18, fontweight="bold"
)

Here are the results for \(x_2\):

x2_ref_vals = (
    roas.query("channel == 'x2'")
    .drop("channel", axis=1)
    .rename(columns={"roas": "ref_val"})
    .to_dict(orient="records")
)

axes = az.plot_posterior(
    data=xr.concat(
        predictions_roas_1["x2"].values(),
        dim=pd.Index(data=predictions_roas_1["x2"].keys(), name="quarter"),
    ),
    figsize=(18, 15),
    grid=(5, 2),
    ref_val={"mu": x2_ref_vals},
    backend_kwargs={"sharex": True, "layout": "constrained"},
)
fig = axes.ravel()[0].get_figure()
fig.suptitle(
    t="Posterior Predictive - ROAS x2 - Model 1", fontsize=18, fontweight="bold"
)

These results are consistent with the ones from the global ROAS. The model underestimates for \(x_2\) (for \(x_1\) the difference lies within the reasonable uncertainty estimates).

Model Diagnostics

We inspect for divergences the trace of the model:

idata_2["sample_stats"]["diverging"].sum().item()

\(\displaystyle 0\)

var_names = [
    "amplitude_trend",
    "ls_trend",
    "alpha",
    "lam",
    "beta_channel",
    "beta_fourier",
    "sigma",
    "nu",
]

az.summary(data=idata_2, var_names=var_names, round_to=3)

axes = az.plot_trace(
    data=idata_2,
    var_names=var_names,
    compact=True,
    kind="rank_bars",
    backend_kwargs={"figsize": (15, 17), "layout": "constrained"},
)
plt.gcf().suptitle("Trace - Model 2", fontsize=16)

Overall, the trace looks fine. One big difference we immediately see is that the regression coefficient \(\beta_1\) of \(x_1\) increased significantly as compared to the previous model

Posterior Predictive Checks

We run the same analysis to visualize the model fit. As berore, we need to rescale the predictions back to the original scale.

pp_vars_original_scale_2 = {
    var_name: xr.apply_ufunc(
        target_scaler.inverse_transform,
        idata_2["posterior"][var_name].expand_dims(dim={"_": 1}, axis=-1),
        input_core_dims=[["date", "_"]],
        output_core_dims=[["date", "_"]],
        vectorize=True,
    ).squeeze(dim="_")
    for var_name in ["mu", "f_trend", "fourier_contribution"]
}

pp_vars_original_scale_2["channel_contributions"] = xr.apply_ufunc(
    target_scaler.inverse_transform,
    idata_2["posterior"]["channel_contributions"],
    input_core_dims=[["date", "channel"]],
    output_core_dims=[["date", "channel"]],
    vectorize=True,
)

pp_likelihood_original_scale_2 = xr.apply_ufunc(
    target_scaler.inverse_transform,
    posterior_predictive_2["posterior_predictive"]["y"].expand_dims(
        dim={"_": 1}, axis=-1
    ),
    input_core_dims=[["date", "_"]],
    output_core_dims=[["date", "_"]],
    vectorize=True,
).squeeze(dim="_")

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_likelihood_original_scale_2,
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.1, "label": r"likelihood $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["mu"],
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.3, "label": r"$\mu$ $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper left")
ax.set_title(label="Posterior Predictive - Model 2", fontsize=18, fontweight="bold")

fig, ax = plt.subplots()
az.plot_ppc(
    data=posterior_predictive_2,
    num_pp_samples=1_000,
    observed_rug=True,
    random_seed=seed,
    ax=ax,
)
ax.set_title(label="Posterior Predictive - Model 2", fontsize=18, fontweight="bold")

The models seem to have a comparable in-sample fit.

Model Components

Now we inspect the model components.

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.5, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.5, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["f_trend"],
    hdi_prob=0.94,
    color="C2",
    fill_kwargs={"alpha": 0.3, "label": r"$f_\text{trend}$ $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["fourier_contribution"],
    hdi_prob=0.94,
    color="C3",
    fill_kwargs={"alpha": 0.3, "label": r"seasonality $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set_title(label="Components Contributions - Model 2", fontsize=18, fontweight="bold")

Here we clearly see that the model overestimates the effect of \(x_1\) by a lot! 🫠

Let’s compare against the true values for the media contributions:

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)
sns.lineplot(
    x="date",
    y="x1_effect",
    data=data_df.assign(x1_effect=lambda x: amplitude * x["x1_effect"]),
    color="C0",
    label=r"$x1$ effect",
    ax=ax[0],
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.3, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax[0],
)
ax[0].legend(loc="upper right")

sns.lineplot(
    x="date",
    y="x2_effect",
    data=data_df.assign(x2_effect=lambda x: amplitude * x["x2_effect"]),
    color="C1",
    label=r"$x2$ effect",
    ax=ax[1],
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_2["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.3, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax[1],
)
ax[1].legend(loc="upper right")

fig.suptitle("Channel Contributions - Model 2", fontsize=18, fontweight="bold")

The overestimation of \(x_1\) is close to \(100\%\)! Note that the estimate of \(x_2\) is not that bad, but still underestimated.

Media Parameters

We now

alpha_posterior_2 = idata_2["posterior"]["alpha"]
lam_posterior_2 = idata_2["posterior"]["lam"] * channels_scaler.scale_
beta_channel_posterior_2 = idata_2["posterior"]["beta_channel"] * channels_scaler.scale_

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=[alpha_posterior_2], combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(
    x=alpha1, color="C0", linestyle="--", linewidth=2, label=r"$\alpha_1$ (true)"
)
ax.axvline(
    x=alpha2, color="C1", linestyle="--", linewidth=2, label=r"$\alpha_2$ (true)"
)
ax.legend(loc="upper left")
ax.set_title(
    label=r"Adstock Parameter $\alpha$ ($94\%$ HDI) - Model 2",
    fontsize=18,
    fontweight="bold",
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=lam_posterior_2, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(x=lam1, color="C0", linestyle="--", linewidth=2, label=r"$\lambda_1$ (true)")
ax.axvline(x=lam2, color="C1", linestyle="--", linewidth=2, label=r"$\lambda_2$ (true)")
ax.legend(loc="upper right")
ax.set_title(
    label=r"Saturation Parameter $\lambda$ ($94\%$ HDI) - Model 2",
    fontsize=18,
    fontweight="bold",
)

beta_channel_posterior_2 = (
    idata_2["posterior"]["beta_channel"]
    * (1 / channels_scaler.scale_)
    * (target_scaler.scale_)
    * (1 / amplitude)
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=beta_channel_posterior_2, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(x=beta1, color="C0", linestyle="--", linewidth=2, label=r"$\beta_1$ (true)")
ax.axvline(x=beta2, color="C1", linestyle="--", linewidth=2, label=r"$\beta_2$ (true)")
ax.legend(loc="upper right")
ax.set_title(
    label=r"Channel $\beta$ ($94\%$ HDI) - Model 1", fontsize=18, fontweight="bold"
)

The carryover parameter \(\alpha_1\) and the regression coefficient \(\beta_1\) for \(x_1\) are the ones that are most affected by the omitted variable bias. For \(x_2\) the model is able to recover the true values within the \(94\%\) HDI.

ROAS Estimation

We now compute the ROAS as before. We start with the global ROAS:

predictions_roas_data_all_2 = {"x1": {}, "x2": {}}

for channel in tqdm(["x1", "x2"]):
    with model_2:
        pm.set_data(
            new_data={
                "channels_scaled": channels_scaler.transform(
                    data_df[channels].assign(**{channel: 0})
                )
            }
        )
        predictions_roas_data_all_2[channel] = pm.sample_posterior_predictive(
            trace=idata_2, var_names=["y", "mu"], progressbar=False, random_seed=rng
        )

predictions_roas_data_scaled_diff_all_2 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_data_scaled_diff_all_2[channel] = pp_vars_original_scale_2[
        "mu"
    ] - xr.apply_ufunc(
        target_scaler.inverse_transform,
        predictions_roas_data_all_2[channel]["posterior_predictive"]["mu"].expand_dims(
            dim={"_": 1}, axis=-1
        ),
        input_core_dims=[["date", "_"]],
        output_core_dims=[["date", "_"]],
        vectorize=True,
    ).squeeze(
        dim="_"
    )

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_all_2["x1"],
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.3, "label": r"$E(y - y_{01})$ $94\%$ HDI"},
    ax=ax[0],
)
sns.lineplot(
    data=data_df.assign(x1_effect=lambda x: amplitude * x["x1_effect"]),
    x="date",
    y="x1_effect",
    color="C0",
    label=r"$x1$ effect",
    ax=ax[0],
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_all_2["x2"],
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.3, "label": r"$E(y - y_{02})$ $94\%$ HDI"},
    ax=ax[1],
)
sns.lineplot(
    data=data_df.assign(x2_effect=lambda x: amplitude * x["x2_effect"]),
    x="date",
    y="x2_effect",
    color="C1",
    label=r"$x2$ effect",
    ax=ax[1],
)
ax[0].legend(loc="upper left")
ax[1].legend(loc="upper left")
ax[0].set(ylabel=None)
ax[1].set(ylabel=None)
fig.suptitle("Potential Outcomes Difference - Model 2", fontsize=18, fontweight="bold")

predictions_roas_all_2 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_all_2[channel] = (
        predictions_roas_data_scaled_diff_all_2[channel].sum(dim="date")
        / data_df[channel].sum()
    )

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)

az.plot_posterior(predictions_roas_all_2["x1"], ref_val=roas_all_1, ax=ax[0])
ax[0].set(title="x1")
az.plot_posterior(predictions_roas_all_2["x2"], ref_val=roas_all_2, ax=ax[1])
ax[1].set(title="x2", xlabel="ROAS")
fig.suptitle("Global ROAS - Model 2", fontsize=18, fontweight="bold")

The ROAS overestimation of \(x_1\) is massive! It is of the order of \(100\%\). For \(x_2\) the estimated global is less than the true value (but less dramatic than \(x_1\)).

We now compute the quarterly ROAS as before.

predictions_roas_data_2 = {"x1": {}, "x2": {}}

for channel in tqdm(["x1", "x2"]):
    for q in tqdm(roas_data[channel].keys()):
        with model_2:
            pm.set_data(
                new_data={
                    "channels_scaled": channels_scaler.transform(
                        roas_data[channel][q][channels]
                    )
                }
            )
            predictions_roas_data_2[channel][q] = pm.sample_posterior_predictive(
                trace=idata_2, var_names=["y", "mu"], progressbar=False, random_seed=rng
            )

predictions_roas_data_scaled_diff_2 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_data_scaled_diff_2[channel] = {
        q: pp_vars_original_scale_2["mu"]
        - xr.apply_ufunc(
            target_scaler.inverse_transform,
            idata_q["posterior_predictive"]["mu"].expand_dims(dim={"_": 1}, axis=-1),
            input_core_dims=[["date", "_"]],
            output_core_dims=[["date", "_"]],
            vectorize=True,
        ).squeeze(dim="_")
        for q, idata_q in tqdm(predictions_roas_data_2[channel].items())
    }

predictions_roas_2 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_2[channel] = {
        q: idata_q.sum(dim="date") / data_df.query("quarter == @q")[channel].sum()
        for q, idata_q in predictions_roas_data_scaled_diff_2[channel].items()
    }

axes = az.plot_posterior(
    data=xr.concat(
        predictions_roas_2["x1"].values(),
        dim=pd.Index(data=predictions_roas_2["x1"].keys(), name="quarter"),
    ),
    figsize=(18, 15),
    grid=(5, 2),
    ref_val={"mu": x1_ref_vals},
    backend_kwargs={"sharex": True, "layout": "constrained"},
)
fig = axes.ravel()[0].get_figure()
fig.suptitle(
    t="Posterior Predictive - ROAS x1 - Model 2", fontsize=18, fontweight="bold"
)

axes = az.plot_posterior(
    data=xr.concat(
        predictions_roas_2["x2"].values(),
        dim=pd.Index(data=predictions_roas_2["x2"].keys(), name="quarter"),
    ),
    figsize=(18, 15),
    grid=(5, 2),
    ref_val={"mu": x2_ref_vals},
    backend_kwargs={"sharex": True, "layout": "constrained"},
)
fig = axes.ravel()[0].get_figure()
fig.suptitle(
    t="Posterior Predictive - ROAS x2 - Model 2", fontsize=18, fontweight="bold"
)

The quarterly ROAS present the same pattern as the global ROAS. The model overestimates the ROAS for \(x_1\) heavily!

Model Specification

In the following model we set priors on the ROAS instead of the regression coefficients.

eps = np.finfo(float).eps


with pm.Model(coords=coords) as model_3:
    # --- Data Containers ---

    index_scaled_data = pm.Data(
        name="index_scaled",
        value=index_scaled.to_numpy().flatten(),
        mutable=True,
        dims="date",
    )

    channels_scaled_data = pm.Data(
        name="channels_scaled",
        value=channels_scaled,
        mutable=True,
        dims=("date", "channel"),
    )

    channels_cost_data = pm.Data(
        name="channels_cost",
        value=data_df[channels].to_numpy(),
        mutable=True,
        dims=("date", "channel"),
    )

    fourier_features_data = pm.Data(
        name="fourier_features",
        value=fourier_features,
        mutable=True,
        dims=("date", "fourier_mode"),
    )

    y_scaled_data = pm.Data(
        name="y_scaled",
        value=target_scaled.to_numpy().flatten(),
        mutable=True,
        dims="date",
    )

    # --- Priors ---

    amplitude_trend = pm.HalfNormal(name="amplitude_trend", sigma=1)
    ls_trend_params = pm.find_constrained_prior(
        distribution=pm.InverseGamma,
        lower=0.1,
        upper=0.9,
        init_guess={"alpha": 3, "beta": 1},
        mass=0.95,
    )
    ls_trend = pm.InverseGamma(name="ls_trend", **ls_trend_params)

    alpha = pm.Beta(name="alpha", alpha=2, beta=3, dims="channel")
    lam = pm.Gamma(name="lam", alpha=2, beta=2, dims="channel")
    roas = pm.LogNormal(
        name="roas",
        mu=np.log(roas_bar_scaled),
        sigma=error_scaled,
        dims="channel",
    )
    beta_fourier = pm.Normal(name="beta_fourier", mu=0, sigma=1, dims="fourier_mode")

    sigma = pm.HalfNormal(name="sigma", sigma=1)
    nu = pm.Gamma(name="nu", alpha=25, beta=2)

    # --- Parametrization ---

    cov_trend = amplitude_trend * pm.gp.cov.ExpQuad(input_dim=1, ls=ls_trend)
    gp_trend = pm.gp.HSGP(m=[20], c=1.5, cov_func=cov_trend)
    f_trend = gp_trend.prior(name="f_trend", X=index_scaled_data[:, None], dims="date")

    channel_adstock = pm.Deterministic(
        name="channel_adstock",
        var=geometric_adstock(
            x=channels_scaled_data,
            alpha=alpha,
            l_max=l_max,
            normalize=True,
        ),
        dims=("date", "channel"),
    )
    channel_adstock_saturated = pm.Deterministic(
        name="channel_adstock_saturated",
        var=logistic_saturation(x=channel_adstock, lam=lam),
        dims=("date", "channel"),
    )
    beta_channel = pm.Deterministic(
        name="beta_channel",
        var=(channels_cost_data.sum(axis=0) * roas)
        / (channel_adstock_saturated.sum(axis=0) + eps),
        dims="channel",
    )
    channel_contributions = pm.Deterministic(
        name="channel_contributions",
        var=channel_adstock_saturated * beta_channel,
        dims=("date", "channel"),
    )

    fourier_contribution = pm.Deterministic(
        name="fourier_contribution",
        var=pt.dot(fourier_features_data, beta_fourier),
        dims="date",
    )

    mu = pm.Deterministic(
        name="mu",
        var=f_trend + channel_contributions.sum(axis=-1) + fourier_contribution,
        dims="date",
    )

    # --- Likelihood ---
    y = pm.StudentT(
        name="y",
        nu=nu,
        mu=mu,
        sigma=sigma,
        observed=y_scaled_data,
        dims="date",
    )

pm.model_to_graphviz(model=model_3)

Remark: The channels_cost_data has the cost of each channel without any scaling. Note that you could use impressions to model the contribution of each channel (channels_scaled_data). This method is agnostic to the type of data you use.

Prior Predictive Checks

Let’s see how this parametrization changes the prior predictive distribution:

with model_3:
    prior_predictive_3 = pm.sample_prior_predictive(samples=2_000, random_seed=rng)

fig, ax = plt.subplots()
az.plot_ppc(data=prior_predictive_3, group="prior", kind="kde", ax=ax)
ax.set_title(label="Prior Predictive - Model 3", fontsize=18, fontweight="bold")

It actually looks very similar to the first model.

Model Fitting

We now fit the model.

with model_3:
    idata_3 = pm.sample(
        target_accept=0.9,
        draws=2_000,
        chains=4,
        nuts_sampler="numpyro",
        random_seed=rng,
    )
    posterior_predictive_3 = pm.sample_posterior_predictive(
        trace=idata_3, random_seed=rng
    )

The sampling time is comparable with the models above.

Model Diagnostics

We now take a look into the results.

idata_3["sample_stats"]["diverging"].sum().item()

\(\displaystyle 0\)

var_names = [
    "amplitude_trend",
    "ls_trend",
    "alpha",
    "lam",
    "roas",
    "beta_channel",
    "beta_fourier",
    "sigma",
    "nu",
]

az.summary(data=idata_3, var_names=var_names, round_to=3)

axes = az.plot_trace(
    data=idata_3,
    var_names=var_names,
    compact=True,
    backend_kwargs={"figsize": (15, 17), "layout": "constrained"},
)
plt.gcf().suptitle("Trace - Model 3", fontsize=16)

Everything looks in order. Note that the ROAS posterior distribution is very close to the prior.

Posterior Predictive Checks

We now evaluate the model fit.

pp_vars_original_scale_3 = {
    var_name: xr.apply_ufunc(
        target_scaler.inverse_transform,
        idata_3["posterior"][var_name].expand_dims(dim={"_": 1}, axis=-1),
        input_core_dims=[["date", "_"]],
        output_core_dims=[["date", "_"]],
        vectorize=True,
    ).squeeze(dim="_")
    for var_name in ["mu", "f_trend", "fourier_contribution"]
}

pp_vars_original_scale_3["channel_contributions"] = xr.apply_ufunc(
    target_scaler.inverse_transform,
    idata_3["posterior"]["channel_contributions"],
    input_core_dims=[["date", "channel"]],
    output_core_dims=[["date", "channel"]],
    vectorize=True,
)

pp_vars_original_scale_3["baseline"] = xr.apply_ufunc(
    target_scaler.inverse_transform,
    (
        idata_3["posterior"]["f_trend"] + idata_3["posterior"]["fourier_contribution"]
    ).expand_dims(dim={"_": 1}, axis=-1),
    input_core_dims=[["date", "_"]],
    output_core_dims=[["date", "_"]],
    vectorize=True,
).squeeze(dim="_")

pp_likelihood_original_scale_3 = xr.apply_ufunc(
    target_scaler.inverse_transform,
    posterior_predictive_3["posterior_predictive"]["y"].expand_dims(
        dim={"_": 1}, axis=-1
    ),
    input_core_dims=[["date", "_"]],
    output_core_dims=[["date", "_"]],
    vectorize=True,
).squeeze(dim="_")

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_likelihood_original_scale_3,
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.1, "label": r"likelihood $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["mu"],
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.3, "label": r"$\mu$ $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper left")
ax.set_title(label="Posterior Predictive - Model 3", fontsize=18, fontweight="bold")

fig, ax = plt.subplots()
az.plot_ppc(
    data=posterior_predictive_3,
    num_pp_samples=1_000,
    observed_rug=True,
    random_seed=seed,
    ax=ax,
)
ax.set_title(label="Posterior Predictive - Model 3", fontsize=18, fontweight="bold")

The posterior predictive distribution looks worse than the previous models.

Model Components

We now plot the model components.

fig, ax = plt.subplots()
sns.lineplot(data=model_df, x="date", y="y", color="black", label=target, ax=ax)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.5, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.5, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["f_trend"],
    hdi_prob=0.94,
    color="C2",
    fill_kwargs={"alpha": 0.3, "label": r"$f_\text{trend}$ $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["fourier_contribution"],
    hdi_prob=0.94,
    color="C3",
    fill_kwargs={"alpha": 0.3, "label": r"seasonality $94\%$ HDI"},
    smooth=False,
    ax=ax,
)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set_title(label="Components Contributions - Model 3", fontsize=18, fontweight="bold")

Observe that the \(94\%\) HDI for both channels is much narrower than the previous models. This is because we are constraining the values through the ROAS tight priors.

Let’s compare now the media contributions against the true values:

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)
sns.lineplot(
    x="date",
    y="x1_effect",
    data=data_df.assign(x1_effect=lambda x: amplitude * x["x1_effect"]),
    color="C0",
    label=r"$x1$ effect",
    ax=ax[0],
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["channel_contributions"].sel(channel="x1"),
    hdi_prob=0.94,
    color="C0",
    fill_kwargs={"alpha": 0.3, "label": r"$x1$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax[0],
)
ax[0].legend(loc="upper right")

sns.lineplot(
    x="date",
    y="x2_effect",
    data=data_df.assign(x2_effect=lambda x: amplitude * x["x2_effect"]),
    color="C1",
    label=r"$x2$ effect",
    ax=ax[1],
)
az.plot_hdi(
    x=date,
    y=pp_vars_original_scale_3["channel_contributions"].sel(channel="x2"),
    hdi_prob=0.94,
    color="C1",
    fill_kwargs={"alpha": 0.3, "label": r"$x2$ contribution $94\%$ HDI"},
    smooth=False,
    ax=ax[1],
)
ax[1].legend(loc="upper right")

fig.suptitle("Channel Contributions - Model 3", fontsize=18, fontweight="bold")

The result for both channels is much better than the previous models *remember we do not have \(z\) in the model).

Media Parameters

alpha_posterior_3 = idata_3["posterior"]["alpha"]
lam_posterior_3 = idata_3["posterior"]["lam"] * channels_scaler.scale_
beta_channel_posterior_3 = idata_3["posterior"]["beta_channel"] * channels_scaler.scale_

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=alpha_posterior_3, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(
    x=alpha1, color="C0", linestyle="--", linewidth=2, label=r"$\alpha_1$ (true)"
)
ax.axvline(
    x=alpha2, color="C1", linestyle="--", linewidth=2, label=r"$\alpha_2$ (true)"
)
ax.legend(loc="upper left")
ax.set_title(
    label=r"Adstock Parameter $\alpha$ ($94\%$ HDI) - Model 2",
    fontsize=18,
    fontweight="bold",
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=lam_posterior_3, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(x=lam1, color="C0", linestyle="--", linewidth=2, label=r"$\lambda_1$ (true)")
ax.axvline(x=lam2, color="C1", linestyle="--", linewidth=2, label=r"$\lambda_2$ (true)")
ax.legend(loc="upper right")
ax.set_title(
    label=r"Saturation Parameter $\lambda$ ($94\%$ HDI) - Model 2",
    fontsize=18,
    fontweight="bold",
)

beta_channel_posterior_3 = (
    idata_3["posterior"]["beta_channel"]
    * (1 / channels_scaler.scale_)
    * (target_scaler.scale_)
    * (1 / amplitude)
)

fig, ax = plt.subplots(figsize=(10, 6))
az.plot_forest(
    data=beta_channel_posterior_3, combined=True, hdi_prob=0.94, colors="black", ax=ax
)
ax.axvline(x=beta1, color="C0", linestyle="--", linewidth=2, label=r"$\beta_1$ (true)")
ax.axvline(x=beta2, color="C1", linestyle="--", linewidth=2, label=r"$\beta_2$ (true)")
ax.legend(loc="upper right")
ax.set_title(
    label=r"Channel $\beta$ ($94\%$ HDI) - Model 1", fontsize=18, fontweight="bold"
)

• For \(x_2\) we recovered all the media parameters within the \(94\%\) HDI.
• For \(x_1\) the model heavily underestimated the carryover parameter \(\alpha_1\) and on the other hand overestimated the regression coefficient \(\beta_1\). The saturation parameter \(\lambda_1\) was closer to the true value, but still outside of the \(94\%\) HDI.

ROAS Estimation

Finally, we compute the ROAS. We should expect them to be much closer to the true values than model 2 since we have explicit tight priors over them. Let’s start with the global ROAS:

predictions_roas_data_all_3 = {"x1": {}, "x2": {}}

for channel in tqdm(["x1", "x2"]):
    with model_3:
        pm.set_data(
            new_data={
                "channels_scaled": channels_scaler.transform(
                    data_df[channels].assign(**{channel: 0}),
                ),
                "channels_cost": data_df[channels].assign(**{channel: 0}),
            }
        )
        predictions_roas_data_all_3[channel] = pm.sample_posterior_predictive(
            trace=idata_3, var_names=["y", "mu"], progressbar=False, random_seed=rng
        )

Note we hat to replace the media model data (channels_scaled) and the media cost data (channels_cost). Here you can see that we could use impressions instead of cost to model the contribution of each channel.

predictions_roas_data_scaled_diff_all_3 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_data_scaled_diff_all_3[channel] = pp_vars_original_scale_3[
        "mu"
    ] - xr.apply_ufunc(
        target_scaler.inverse_transform,
        predictions_roas_data_all_3[channel]["posterior_predictive"]["mu"].expand_dims(
            dim={"_": 1}, axis=-1
        ),
        input_core_dims=[["date", "_"]],
        output_core_dims=[["date", "_"]],
        vectorize=True,
    ).squeeze(dim="_")

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_all_3["x1"],
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.3, "label": r"$E(y - y_{01})$ $94\%$ HDI"},
    ax=ax[0],
)
sns.lineplot(
    data=data_df.assign(x1_effect=lambda x: amplitude * x["x1_effect"]),
    x="date",
    y="x1_effect",
    color="C0",
    label=r"$x1$ effect",
    ax=ax[0],
)
az.plot_hdi(
    x=date,
    y=predictions_roas_data_scaled_diff_all_3["x2"],
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.3, "label": r"$E(y - y_{02})$ $94\%$ HDI"},
    ax=ax[1],
)
sns.lineplot(
    data=data_df.assign(x2_effect=lambda x: amplitude * x["x2_effect"]),
    x="date",
    y="x2_effect",
    color="C1",
    label=r"$x2$ effect",
    ax=ax[1],
)
ax[0].legend(loc="upper left")
ax[1].legend(loc="upper left")
ax[0].set(ylabel=None)
ax[1].set(ylabel=None)
fig.suptitle("Potential Outcomes Difference - Model 3", fontsize=18, fontweight="bold")

predictions_roas_all_3 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_all_3[channel] = (
        predictions_roas_data_scaled_diff_all_3[channel].sum(dim="date")
        / data_df[channel].sum()
    )

fig, ax = plt.subplots(
    nrows=2, ncols=1, figsize=(12, 7), sharex=True, sharey=False, layout="constrained"
)

az.plot_posterior(predictions_roas_all_3["x1"], ref_val=roas_all_1, ax=ax[0])
ax[0].set(title="x1")
az.plot_posterior(predictions_roas_all_3["x2"], ref_val=roas_all_2, ax=ax[1])
ax[1].set(title="x2", xlabel="ROAS")
fig.suptitle("Global ROAS - Model 3", fontsize=18, fontweight="bold")

The results look as expected 🚀! The ROAS estimate for \(x_1\) is much closer to the true value than in the second model.

We wrap up the analysis by computing the quarterly ROAS.

predictions_roas_data_3 = {"x1": {}, "x2": {}}

for channel in tqdm(["x1", "x2"]):
    for q in tqdm(roas_data[channel].keys()):
        with model_3:
            pm.set_data(
                new_data={
                    "channels_scaled": channels_scaler.transform(
                        roas_data[channel][q][channels]
                    ),
                    "channels_cost": roas_data[channel][q][channels],
                }
            )
            predictions_roas_data_3[channel][q] = pm.sample_posterior_predictive(
                trace=idata_3, var_names=["y", "mu"], progressbar=False, random_seed=rng
            )

predictions_roas_data_scaled_diff_3 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_data_scaled_diff_3[channel] = {
        q: pp_vars_original_scale_3["mu"]
        - xr.apply_ufunc(
            target_scaler.inverse_transform,
            idata_q["posterior_predictive"]["mu"].expand_dims(dim={"_": 1}, axis=-1),
            input_core_dims=[["date", "_"]],
            output_core_dims=[["date", "_"]],
            vectorize=True,
        ).squeeze(dim="_")
        for q, idata_q in tqdm(predictions_roas_data_3[channel].items())
    }

predictions_roas_3 = {"x1": {}, "x2": {}}

for channel in ["x1", "x2"]:
    predictions_roas_3[channel] = {
        q: idata_q.sum(dim="date") / data_df.query("quarter == @q")[channel].sum()
        for q, idata_q in predictions_roas_data_scaled_diff_3[channel].items()
    }

axes = az.plot_posterior(
    data=xr.concat(
        predictions_roas_3["x1"].values(),
        dim=pd.Index(data=predictions_roas_3["x1"].keys(), name="quarter"),
    ),
    figsize=(18, 15),
    grid=(5, 2),
    ref_val={"mu": x1_ref_vals},
    backend_kwargs={"sharex": True, "layout": "constrained"},
)
fig = axes.ravel()[0].get_figure()
fig.suptitle(
    t="Posterior Predictive - ROAS x1 - Model 3", fontsize=18, fontweight="bold"
)

axes = az.plot_posterior(
    data=xr.concat(
        predictions_roas_3["x2"].values(),
        dim=pd.Index(data=predictions_roas_3["x2"].keys(), name="quarter"),
    ),
    figsize=(18, 15),
    grid=(5, 2),
    ref_val={"mu": x2_ref_vals},
    backend_kwargs={"sharex": True, "layout": "constrained"},
)
fig = axes.ravel()[0].get_figure()
fig.suptitle(
    t="Posterior Predictive - ROAS x2 - Model 3", fontsize=18, fontweight="bold"
)

The results at quarterly level are consistent with the global ROAS.