r2d2#

R2D2 prior for variance decomposition in regression models.

The R2D2 prior (Zhang et al., 2020) provides automatic shrinkage and variable selection by placing a prior on the R-squared of a regression model, then allocating explained variance across coefficients via a Dirichlet decomposition. Uses a Normal base distribution following Aguilar & Bürkner (2022).

References#

Examples#

Standalone PyMC:

>>> from pymc_marketing.r2d2 import R2D2
>>> from pymc_extras.prior import Prior
>>> import pymc as pm, pymc.dims as pmd, xarray as xr
>>>
>>> ds = xr.Dataset(
...     {
...         "controls": (("obs", "control"), np.random.randn(100, 2)),
...         "y": ("obs", np.random.randn(100)),
...     },
...     coords={"obs": range(100), "control": ["a", "b"]},
... )
>>> r2d2 = R2D2(
...     r2=Prior("Beta", mu=0.8, sigma=0.2),
...     total_sigma=Prior("LogNormal", mu=0, sigma=1),
...     dims={"control": "control"},
... )
>>> with pm.Model(coords={"obs": range(100), "control": ["a", "b"]}) as model:
...     controls = pmd.Data("controls", ds["controls"])
...     beta_control = r2d2.split("control").create_variable("beta_control")
...     intercept = pmd.Normal("intercept", mu=0, sigma=1)
...     mu = intercept + controls @ beta_control
...     sigma = r2d2.error_sigma.create_variable("sigma")
...     pmd.Normal("y_obs", mu=mu, sigma=sigma, observed=ds["y"])

MMM integration:

>>> from pymc_marketing.mmm import MMM
>>> from pymc_marketing.mmm.components.adstock import GeometricAdstock
>>> from pymc_marketing.mmm.components.saturation import LogisticSaturation
>>>
>>> r2d2 = R2D2(
...     r2=Prior("Beta", mu=0.8, sigma=0.2),
...     total_sigma=Prior("LogNormal", mu=0, sigma=1),
...     dims={"control": "control", "fourier": "fourier_mode"},
... )
>>> mmm = MMM(
...     adstock=GeometricAdstock(l_max=8),
...     saturation=LogisticSaturation(),
...     model_config={
...         "likelihood": Prior("Normal", sigma=r2d2.error_sigma),
...         "gamma_control": r2d2.split("control"),
...         "gamma_fourier": r2d2.split("fourier"),
...     },
... )

Gotchas#

  • r2 must be scalar: r2 Prior must not have dims. The R2D2 prior uses a single global R² decomposition.

  • total_sigma must be scalar: No dims allowed. This is per the R2D2 paper. total_sigma represents the global scale of the response variable.

  • Beta prior sigma must be small enough: For r2=Prior("Beta", mu=M, sigma=S), the sigma must satisfy S < sqrt(M * (1-M)). For mu=0.8, this means sigma < 0.4. Values too large produce invalid (alpha=0, beta=0) parameters. Use sigma=0.2.

  • Only components in dims get variance: Covariates not included in dims are not covered by the decomposition. Ensure all relevant covariates are included.

  • Shared decomposition: All split() and error_sigma references share the same underlying Dirichlet. Changing one affects all.

Classes

R2D2(r2, total_sigma, dims)

R2D2 variance decomposition.

R2D2Sigma(decomposition)

Lazy reference to the residual standard deviation variable.

R2D2Split(decomposition, component_name[, _dims])

Lazy reference to a component's split variable.