"""
Multivariate MMM implementation.
Models multiple correlated outcomes with cross-effects
(cannibalization, halo effects) between products/brands.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
from .base import BaseExtendedMMM
from ..results import CrossEffectSummary
if TYPE_CHECKING:
from ..config import MultivariateModelConfig
[docs]
class MultivariateMMM(BaseExtendedMMM):
"""
MMM with multiple correlated outcomes and cross-effects.
This model handles:
- Multiple outcome variables (e.g., sales by product/brand)
- Correlated residuals using LKJ prior
- Cross-effects between outcomes (cannibalization, halo)
- Promotion-modulated cross-effects
Parameters
----------
X_media : np.ndarray
Media variable matrix (n_obs, n_channels)
outcome_data : dict[str, np.ndarray]
Dictionary mapping outcome names to their data
channel_names : list[str]
Names of media channels
config : MultivariateModelConfig
Configuration for outcomes and cross-effects
promotion_data : dict[str, np.ndarray] | None
Optional promotion indicators for modulated effects
index : pd.Index | None
Optional time index
"""
[docs]
def __init__(
self,
X_media: np.ndarray,
outcome_data: dict[str, np.ndarray],
channel_names: list[str],
config: "MultivariateModelConfig",
promotion_data: dict[str, np.ndarray] | None = None,
index: pd.Index | None = None,
model_config=None,
trend_config=None,
):
# Use first outcome for y (placeholder)
first_outcome = list(outcome_data.values())[0]
super().__init__(
X_media,
first_outcome,
channel_names,
index,
model_config=model_config,
trend_config=trend_config,
)
self.outcome_data = outcome_data
self.config = config
self.promotion_data = promotion_data or {}
self.outcome_names = [o.name for o in config.outcomes]
self.n_outcomes = len(config.outcomes)
# Per-outcome standardization (the fixed priors assume z-scale
# outcomes). Data attributes stay on the caller's raw scale; the
# likelihood fits standardized outcomes and report-consumed
# deterministics are registered back in original units.
self.outcome_means = {
name: float(np.asarray(outcome_data[name], dtype=float).mean())
for name in self.outcome_names
}
self.outcome_stds = {
name: float(np.asarray(outcome_data[name], dtype=float).std()) + 1e-8
for name in self.outcome_names
}
# Build cross-effect structure
self._cross_effect_specs = self._build_cross_effect_specs()
def _build_coords(self) -> dict:
coords = super()._build_coords()
coords["outcome"] = self.outcome_names
return coords
def _build_cross_effect_specs(self):
"""Convert config cross-effects to internal specs."""
from ..components.cross_effects import CrossEffectSpec
from ..config import CrossEffectType
specs = []
for ce in self.config.cross_effects:
source_idx = self.outcome_names.index(ce.source_outcome)
target_idx = self.outcome_names.index(ce.target_outcome)
specs.append(
CrossEffectSpec(
source_idx=source_idx,
target_idx=target_idx,
effect_type=ce.effect_type.value,
prior_sigma=ce.prior_sigma,
)
)
# Add reverse direction for symmetric effects
if ce.effect_type == CrossEffectType.SYMMETRIC:
specs.append(
CrossEffectSpec(
source_idx=target_idx,
target_idx=source_idx,
effect_type=ce.effect_type.value,
prior_sigma=ce.prior_sigma,
)
)
return specs
def _baseline_dynamics_terms(self) -> tuple[dict, dict]:
"""Per-outcome trend + seasonality terms from the config's share flags."""
return self._build_baseline_dynamics(
self.config.outcomes,
self.config.share_trend,
self.config.share_seasonality,
)
def _build_model(self) -> pm.Model:
"""Build the multivariate model."""
from ..components.cross_effects import (
build_cross_effect_matrix,
compute_cross_effect_contribution,
)
from ..components.observation import build_multivariate_likelihood
coords = self._build_coords()
with pm.Model(coords=coords) as model:
# Data. ``Y`` stays raw: cross-effects are defined on the raw
# source outcome (``psi[s, t] * Y_raw[:, s]``), preserving psi's
# original-unit interpretation. The likelihood fits standardized
# outcomes so the fixed priors are well-calibrated.
X_media = pm.Data("X_media", self.X_media, dims=("obs", "channel"))
Y_matrix = np.column_stack(
[self.outcome_data[name] for name in self.outcome_names]
)
means = np.array([self.outcome_means[n] for n in self.outcome_names])
stds = np.array([self.outcome_stds[n] for n in self.outcome_names])
Y_standardized = (Y_matrix - means) / stds
Y = pm.Data("Y", Y_matrix, dims=("obs", "outcome"))
# Media transformations: normalize -> geometric adstock -> logistic
# saturation (the previously-unused ``alpha`` is now the carryover).
media_transformed = []
channel_tx = {} # channel -> (x_sat, apply, x_input) for experiments
for i, channel in enumerate(self.channel_names):
x = X_media[:, i]
if self.config.share_media_adstock:
if i == 0:
alpha_shared = pm.Beta("alpha_shared", alpha=2, beta=2)
alpha = alpha_shared
else:
alpha = pm.Beta(f"alpha_{channel}", alpha=2, beta=2)
lam = pm.Gamma(f"lambda_{channel}", alpha=3, beta=1)
apply = self._media_transform_apply(i, alpha, lam)
x_sat = apply(x)
media_transformed.append(x_sat)
channel_tx[channel] = (x_sat, apply, x)
media_transformed = pt.stack(media_transformed, axis=1)
# Outcome-level parameters. Per-outcome prior scales are read from
# each OutcomeConfig (intercept_prior_sigma / media_effect.sigma) so
# a spec/DAG outcome-prior override actually takes effect; the
# defaults (2.0 / 0.5) reproduce the historical shared-scalar priors.
intercept_sigma = np.array(
[oc.intercept_prior_sigma for oc in self.config.outcomes], dtype=float
)
media_sigma = np.array(
[oc.media_effect.sigma for oc in self.config.outcomes], dtype=float
)
alpha = pm.Normal("alpha", mu=0, sigma=intercept_sigma, dims="outcome")
beta_media = pm.Normal(
"beta_media",
mu=0,
sigma=media_sigma[:, None],
dims=("outcome", "channel"),
)
# Baseline dynamics per outcome: a trend and/or Fourier seasonality
# (shared across outcomes or per-outcome per the config flags), gated
# on each outcome's include_trend / include_seasonality. None when
# unconfigured, so a model built without trend/seasonality is
# byte-identical.
trend_terms, seasonality_terms = self._baseline_dynamics_terms()
# Cross-effects
if self._cross_effect_specs:
psi_matrix, psi_params = build_cross_effect_matrix(
self._cross_effect_specs,
self.n_outcomes,
name_prefix="psi",
)
pm.Deterministic("psi_matrix", psi_matrix)
else:
psi_matrix = None
# Build expected values on the standardized scale. Cross-effects
# act on the raw source outcome and are divided by the target's
# std, so ``psi`` keeps its raw source->target interpretation.
mu_list = []
for k, outcome_config in enumerate(self.config.outcomes):
mu_k = alpha[k]
# Media effects
mu_k = mu_k + pt.dot(media_transformed, beta_media[k, :])
# Baseline dynamics (standardized scale)
if trend_terms.get(k) is not None:
mu_k = mu_k + trend_terms[k]
if seasonality_terms.get(k) is not None:
mu_k = mu_k + seasonality_terms[k]
# Cross-effects
if psi_matrix is not None:
# Get promotion modulation
modulation = {}
for ce in self.config.cross_effects:
source_idx = self.outcome_names.index(ce.source_outcome)
if ce.promotion_modulated and ce.promotion_column:
if ce.promotion_column in self.promotion_data:
modulation[source_idx] = pm.Data(
f"promo_{ce.source_outcome}",
self.promotion_data[ce.promotion_column],
dims="obs",
)
cross_contrib = compute_cross_effect_contribution(
Y, psi_matrix, k, self.n_outcomes, modulation
)
mu_k = mu_k + cross_contrib / stds[k]
mu_list.append(mu_k)
mu_standardized = pt.stack(mu_list, axis=1)
# Original-unit predictions (reports and oracles consume this)
pm.Deterministic(
"mu",
mu_standardized * stds[None, :] + means[None, :],
dims=("obs", "outcome"),
)
# Per-channel contribution to the PRIMARY outcome, original units
# (obs, channel) — the single-KPI report surface. The joint model
# is additive in the media per outcome (mu_k gains
# ``Σ_c beta_media[k,c]·sat_c``), so the primary outcome's per-channel
# contribution is ``beta_media[k0,c]·sat_c·std_k0``. Cross-effects
# are outcome→outcome (not media) and are excluded.
k0 = self._primary_outcome_index()
pm.Deterministic(
"channel_contributions",
media_transformed * beta_media[k0, :] * stds[k0],
dims=("obs", "channel"),
)
# Multivariate likelihood (standardized scale)
build_multivariate_likelihood(
"Y_obs",
mu_standardized,
Y_standardized,
self.n_outcomes,
self.config.lkj_eta,
dims=("obs", "outcome"),
)
# Experiment calibration: per (channel, outcome) the model-implied
# contribution is beta_media[outcome, channel] * saturated spend on
# the standardized scale; the outcome's std converts the estimand
# back to original units.
if self.experiments:
handles = {}
for k, outcome in enumerate(self.outcome_names):
for c, channel in enumerate(self.channel_names):
x_sat, apply, x_input = channel_tx[channel]
handles[(channel, outcome)] = {
"coef": beta_media[k, c],
"x_sat": x_sat,
"apply": apply,
"x_input": x_input,
"spend_obs": self.X_media[:, c],
"scale": self.outcome_stds[outcome],
}
self._add_experiment_likelihoods(handles)
return model
[docs]
def get_cross_effects_summary(self) -> pd.DataFrame:
"""Extract cross-effect estimates."""
self._check_fitted()
if "psi_matrix" not in self._trace.posterior:
return pd.DataFrame()
results = []
psi = self._trace.posterior["psi_matrix"]
for spec in self._cross_effect_specs:
vals = psi[:, :, spec.source_idx, spec.target_idx].values.flatten()
from mmm_framework.utils.arviz_compat import hdi_bounds
hdi = hdi_bounds(vals, 0.94)
results.append(
CrossEffectSummary(
source=self.outcome_names[spec.source_idx],
target=self.outcome_names[spec.target_idx],
effect_type=spec.effect_type,
mean=float(np.mean(vals)),
sd=float(np.std(vals)),
hdi_low=float(hdi[0]),
hdi_high=float(hdi[1]),
)
)
return pd.DataFrame([r.to_dict() for r in results])
[docs]
def get_correlation_matrix(self) -> pd.DataFrame:
"""Extract posterior mean correlation matrix."""
self._check_fitted()
corr = (
self._trace.posterior["Y_obs_correlation"]
.mean(dim=["chain", "draw"])
.values
)
return pd.DataFrame(
corr,
index=self.outcome_names,
columns=self.outcome_names,
)
__all__ = ["MultivariateMMM"]