MMM Extensions

The mmm_extensions module provides advanced modeling capabilities including nested/mediated models, multivariate outcomes, and variable selection.

Note

This module uses lazy loading for heavy PyMC dependencies. Import the specific classes you need rather than using from mmm_extensions import *.

Configuration

Configuration Classes for MMM Extensions

Immutable configuration objects for nested and multivariate models.

Note: Shared enums like SaturationType are imported from the main config module to avoid duplication.

class mmm_framework.mmm_extensions.config.MediatorType(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Type of mediating variable based on observability.

FULLY_OBSERVED = 'fully_observed'

PARTIALLY_OBSERVED = 'partially_observed'

AGGREGATED_SURVEY = 'aggregated_survey'

FULLY_LATENT = 'fully_latent'

class mmm_framework.mmm_extensions.config.CrossEffectType(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Type of cross-product effect.

CANNIBALIZATION (psi = -HalfNormal) and HALO (psi = +HalfNormal) impose the sign a priori. UNCONSTRAINED (psi ~ Normal(0, sigma)) lets the data choose the sign – preferable when you do not want to assume the direction, and the honest default given that a one-sided prior makes “the posterior is below zero” near-automatic. Note that on observed sibling outcomes the directional cross-effect psi is confounded with the residual correlation (only their sum is identified), so an unconstrained psi measures a cross-outcome association, not causal cannibalization – see mmm_framework.mmm_extensions.builders.cross_effect().

CANNIBALIZATION = 'cannibalization'

HALO = 'halo'

SYMMETRIC = 'symmetric'

ASYMMETRIC = 'asymmetric'

UNCONSTRAINED = 'unconstrained'

class mmm_framework.mmm_extensions.config.EffectConstraint(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Constraint on effect direction.

NONE = 'none'

POSITIVE = 'positive'

NEGATIVE = 'negative'

class mmm_framework.mmm_extensions.config.VariableSelectionMethod(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Available variable selection methods for control variables.

NONE = 'none'

REGULARIZED_HORSESHOE = 'regularized_horseshoe'

FINNISH_HORSESHOE = 'finnish_horseshoe'

SPIKE_SLAB = 'spike_slab'

BAYESIAN_LASSO = 'bayesian_lasso'

class mmm_framework.mmm_extensions.config.AdstockConfig(l_max=8, prior_type='beta', prior_alpha=2.0, prior_beta=2.0, normalize=True)

Bases: object

Configuration for adstock transformation.

l_max: int = 8

prior_type: str = 'beta'

prior_alpha: float = 2.0

prior_beta: float = 2.0

normalize: bool = True

__init__(l_max=8, prior_type='beta', prior_alpha=2.0, prior_beta=2.0, normalize=True)

class mmm_framework.mmm_extensions.config.SaturationConfig(type=SaturationType.LOGISTIC, lam_prior_alpha=3.0, lam_prior_beta=1.0, kappa_prior_alpha=2.0, kappa_prior_beta=2.0, slope_prior_alpha=3.0, slope_prior_beta=1.0)

Bases: object

Configuration for saturation transformation.

type: SaturationType = 'logistic'

lam_prior_alpha: float = 3.0

lam_prior_beta: float = 1.0

kappa_prior_alpha: float = 2.0

kappa_prior_beta: float = 2.0

slope_prior_alpha: float = 3.0

slope_prior_beta: float = 1.0

__init__(type=SaturationType.LOGISTIC, lam_prior_alpha=3.0, lam_prior_beta=1.0, kappa_prior_alpha=2.0, kappa_prior_beta=2.0, slope_prior_alpha=3.0, slope_prior_beta=1.0)

class mmm_framework.mmm_extensions.config.EffectPriorConfig(constraint=EffectConstraint.NONE, mu=0.0, sigma=1.0)

Bases: object

Configuration for effect coefficient prior.

constraint: EffectConstraint = 'none'

mu: float = 0.0

sigma: float = 1.0

__init__(constraint=EffectConstraint.NONE, mu=0.0, sigma=1.0)

class mmm_framework.mmm_extensions.config.MediatorConfig(name, mediator_type=MediatorType.PARTIALLY_OBSERVED, media_effect=<factory>, outcome_effect=<factory>, observation_noise_sigma=0.1, allow_direct_effect=True, direct_effect=<factory>, apply_adstock=True, apply_saturation=True, adstock=<factory>, saturation=<factory>)

Bases: object

Configuration for a mediating variable.

name: str

mediator_type: MediatorType = 'partially_observed'

media_effect: EffectPriorConfig

outcome_effect: EffectPriorConfig

observation_noise_sigma: float = 0.1

allow_direct_effect: bool = True

direct_effect: EffectPriorConfig

apply_adstock: bool = True

apply_saturation: bool = True

adstock: AdstockConfig

saturation: SaturationConfig

__init__(name, mediator_type=MediatorType.PARTIALLY_OBSERVED, media_effect=<factory>, outcome_effect=<factory>, observation_noise_sigma=0.1, allow_direct_effect=True, direct_effect=<factory>, apply_adstock=True, apply_saturation=True, adstock=<factory>, saturation=<factory>)

class mmm_framework.mmm_extensions.config.OutcomeConfig(name, column, intercept_prior_sigma=2.0, media_effect=<factory>, include_trend=True, include_seasonality=True)

Bases: object

Configuration for an outcome variable.

name: str

column: str

intercept_prior_sigma: float = 2.0

media_effect: EffectPriorConfig

include_trend: bool = True

include_seasonality: bool = True

__init__(name, column, intercept_prior_sigma=2.0, media_effect=<factory>, include_trend=True, include_seasonality=True)

class mmm_framework.mmm_extensions.config.CrossEffectConfig(source_outcome, target_outcome, effect_type=CrossEffectType.CANNIBALIZATION, prior_sigma=0.3, promotion_modulated=True, promotion_column=None, lag=0)

Bases: object

Configuration for cross-product effects.

source_outcome: str

target_outcome: str

effect_type: CrossEffectType = 'cannibalization'

prior_sigma: float = 0.3

promotion_modulated: bool = True

promotion_column: str | None = None

lag: int = 0

__init__(source_outcome, target_outcome, effect_type=CrossEffectType.CANNIBALIZATION, prior_sigma=0.3, promotion_modulated=True, promotion_column=None, lag=0)

class mmm_framework.mmm_extensions.config.NestedModelConfig(mediators=<factory>, media_to_mediator_map=<factory>, share_adstock_across_mediators=True, share_saturation_across_mediators=False)

Bases: object

Configuration for nested/mediated model.

mediators: tuple[MediatorConfig, ...]

media_to_mediator_map: dict[str, tuple[str, ...]]

share_adstock_across_mediators: bool = True

share_saturation_across_mediators: bool = False

__init__(mediators=<factory>, media_to_mediator_map=<factory>, share_adstock_across_mediators=True, share_saturation_across_mediators=False)

class mmm_framework.mmm_extensions.config.MultivariateModelConfig(outcomes=<factory>, cross_effects=<factory>, lkj_eta=2.0, share_media_adstock=True, share_media_saturation=False, share_trend=False, share_seasonality=True)

Bases: object

Configuration for multivariate outcome model.

outcomes: tuple[OutcomeConfig, ...]

cross_effects: tuple[CrossEffectConfig, ...]

lkj_eta: float = 2.0

share_media_adstock: bool = True

share_media_saturation: bool = False

share_trend: bool = False

share_seasonality: bool = True

__init__(outcomes=<factory>, cross_effects=<factory>, lkj_eta=2.0, share_media_adstock=True, share_media_saturation=False, share_trend=False, share_seasonality=True)

class mmm_framework.mmm_extensions.config.CombinedModelConfig(nested, multivariate, mediator_to_outcome_map=<factory>)

Bases: object

Configuration for combined nested + multivariate model.

nested: NestedModelConfig

multivariate: MultivariateModelConfig

mediator_to_outcome_map: dict[str, tuple[str, ...]]

__init__(nested, multivariate, mediator_to_outcome_map=<factory>)

class mmm_framework.mmm_extensions.config.HorseshoeConfig(expected_nonzero=3, slab_scale=2.0, slab_df=4.0, local_df=5.0, global_df=1.0)

Bases: object

Configuration for horseshoe-family priors.

The regularized horseshoe (Piironen & Vehtari, 2017) provides: - Strong shrinkage of small effects toward zero - Minimal shrinkage of large effects (signal preservation) - Regularized slab to prevent unrealistic effect sizes

Parameters

expected_nonzeroint

Prior expectation of the number of nonzero coefficients (D0). Used to calibrate the global shrinkage parameter tau.

slab_scalefloat

Scale parameter for the slab (c in the formulation). Controls maximum expected coefficient magnitude in std units.

slab_dffloat

Degrees of freedom for the slab’s distribution. Lower = heavier tails = allow larger effects.

local_dffloat

Degrees of freedom for local shrinkage parameters (lambda). Default 5.0; use 1.0 for half-Cauchy (original horseshoe).

global_dffloat

Degrees of freedom for global shrinkage parameter (tau). Default 1.0 gives half-Cauchy (standard horseshoe).

expected_nonzero: int = 3

slab_scale: float = 2.0

slab_df: float = 4.0

local_df: float = 5.0

global_df: float = 1.0

__init__(expected_nonzero=3, slab_scale=2.0, slab_df=4.0, local_df=5.0, global_df=1.0)

class mmm_framework.mmm_extensions.config.SpikeSlabConfig(prior_inclusion_prob=0.5, spike_scale=0.01, slab_scale=1.0, use_continuous_relaxation=True, temperature=0.1)

Bases: object

Configuration for spike-and-slab priors.

The spike-and-slab uses a mixture of two distributions: - Spike: concentrated near zero (for excluded variables) - Slab: diffuse prior (for included variables)

Parameters

prior_inclusion_probfloat

Prior probability that each coefficient is nonzero. 0.5 represents maximum uncertainty about inclusion.

spike_scalefloat

Standard deviation of the spike (near-zero distribution). Should be small (0.01-0.05) to effectively zero coefficients.

slab_scalefloat

Standard deviation of the slab (nonzero distribution). Should reflect expected magnitude of true effects.

use_continuous_relaxationbool

If True, use continuous relaxation for gradient-based sampling. Required for NUTS; set False only for Gibbs samplers.

temperaturefloat

Temperature for continuous relaxation (lower = sharper selection).

prior_inclusion_prob: float = 0.5

spike_scale: float = 0.01

slab_scale: float = 1.0

use_continuous_relaxation: bool = True

temperature: float = 0.1

__init__(prior_inclusion_prob=0.5, spike_scale=0.01, slab_scale=1.0, use_continuous_relaxation=True, temperature=0.1)

class mmm_framework.mmm_extensions.config.LassoConfig(regularization=1.0, adaptive=False)

Bases: object

Configuration for Bayesian LASSO prior.

The Bayesian LASSO (Park & Casella, 2008) places Laplace priors on coefficients, providing L1-like shrinkage in a Bayesian context.

Parameters

regularizationfloat

Regularization strength (lambda). Higher = more shrinkage.

adaptivebool

If True, use adaptive LASSO with coefficient-specific penalties.

regularization: float = 1.0

adaptive: bool = False

__init__(regularization=1.0, adaptive=False)

class mmm_framework.mmm_extensions.config.VariableSelectionConfig(method=VariableSelectionMethod.NONE, horseshoe=<factory>, spike_slab=<factory>, lasso=<factory>, exclude_variables=(), include_only_variables=None)

Bases: object

Complete configuration for control variable selection.

CAUSAL WARNING: Variable selection should ONLY be applied to precision control variables—variables that affect the outcome but do NOT affect treatment assignment (media spending). Applying selection to confounders can introduce severe bias in causal effect estimates.

Parameters

methodVariableSelectionMethod

Which selection method to use.

horseshoeHorseshoeConfig

Configuration for horseshoe methods.

spike_slabSpikeSlabConfig

Configuration for spike-and-slab.

lassoLassoConfig

Configuration for Bayesian LASSO.

exclude_variablestuple[str, …]

Variables to EXCLUDE from selection (always include with standard priors). Use for known confounders that must remain in the model.

include_only_variablestuple[str, …] | None

If specified, only apply selection to these variables. All others use standard priors.

Examples

>>> # Sparse selection with excluded confounders
>>> config = VariableSelectionConfig(
...     method=VariableSelectionMethod.REGULARIZED_HORSESHOE,
...     horseshoe=HorseshoeConfig(expected_nonzero=3),
...     exclude_variables=("distribution", "price", "competitor_media"),
... )

method: VariableSelectionMethod = 'none'

horseshoe: HorseshoeConfig

spike_slab: SpikeSlabConfig

lasso: LassoConfig

exclude_variables: tuple[str, ...] = ()

include_only_variables: tuple[str, ...] | None = None

get_selectable_variables(all_control_names)

Partition control variables into selectable and non-selectable.

Return type:

tuple[list[str], list[str]]

Parameters

all_control_nameslist[str]

All control variable names in the model.

Returns

tuple[list[str], list[str]]

(variables_with_selection, variables_without_selection)

__init__(method=VariableSelectionMethod.NONE, horseshoe=<factory>, spike_slab=<factory>, lasso=<factory>, exclude_variables=(), include_only_variables=None)

mmm_framework.mmm_extensions.config.sparse_selection_config(expected_relevant=3, confounders=())

Create configuration for sparse control selection.

Use when you expect only a few controls are truly relevant.

Return type:

VariableSelectionConfig

Parameters

expected_relevantint

Prior expectation of relevant controls.

confounderstuple[str, …]

Confounder variables to exclude from selection.

mmm_framework.mmm_extensions.config.dense_selection_config(regularization=1.0, confounders=())

Create configuration for dense control selection.

Use when you expect many controls have small effects.

Return type:

VariableSelectionConfig

Parameters

regularizationfloat

Regularization strength.

confounderstuple[str, …]

Confounder variables to exclude from selection.

mmm_framework.mmm_extensions.config.inclusion_prob_selection_config(prior_inclusion=0.5, confounders=())

Create configuration with explicit inclusion probabilities.

Use when you want interpretable posterior inclusion probabilities.

Return type:

VariableSelectionConfig

Parameters

prior_inclusionfloat

Prior probability of inclusion for each variable.

confounderstuple[str, …]

Confounder variables to exclude from selection.

class mmm_framework.mmm_extensions.config.MediatorObservationType(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

How the mediator is observed.

Extends the original MediatorType to add aggregated survey support.

FULLY_OBSERVED = 'fully_observed'

PARTIALLY_OBSERVED = 'partially_observed'

AGGREGATED_SURVEY = 'aggregated_survey'

FULLY_LATENT = 'fully_latent'

class mmm_framework.mmm_extensions.config.AggregatedSurveyLikelihood(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Likelihood for aggregated survey observations.

BINOMIAL = 'binomial'

NORMAL = 'normal'

BETA_BINOMIAL = 'beta_binomial'

class mmm_framework.mmm_extensions.config.AggregatedSurveyConfig(aggregation_map, sample_sizes, likelihood=AggregatedSurveyLikelihood.BINOMIAL, design_effect=1.0, aggregation_function='mean', overdispersion_prior_sigma=0.1)

Bases: object

Configuration for temporally aggregated survey observations.

Used when surveys are fielded continuously over a period (e.g., monthly) and results are aggregated, rather than point-in-time snapshots.

Attributes

aggregation_mapdict[int, tuple[int, …]]

Maps observation index to constituent time indices. E.g., {0: (0, 1, 2, 3), 1: (4, 5, 6, 7)} for monthly surveys in weekly model.

sample_sizestuple[int, …]

Number of respondents per survey wave. Length must match aggregation_map.

likelihoodAggregatedSurveyLikelihood

Which likelihood to use for the observation model.

design_effectfloat

Survey design effect multiplier on variance (default 1.0). Use >1 for clustered samples, complex weighting, etc.

aggregation_functionLiteral[“mean”, “sum”, “last”]

How to aggregate latent values within each period. “mean” is typical for awareness (average state during fielding).

overdispersion_prior_sigmafloat

Prior sigma for overdispersion parameter (beta-binomial only).

aggregation_map: dict[int, tuple[int, ...]]

sample_sizes: tuple[int, ...]

likelihood: AggregatedSurveyLikelihood = 'binomial'

design_effect: float = 1.0

aggregation_function: Literal['mean', 'sum', 'last'] = 'mean'

overdispersion_prior_sigma: float = 0.1

__init__(aggregation_map, sample_sizes, likelihood=AggregatedSurveyLikelihood.BINOMIAL, design_effect=1.0, aggregation_function='mean', overdispersion_prior_sigma=0.1)

class mmm_framework.mmm_extensions.config.MediatorDynamics(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Latent-state dynamics of a mediator (or latent factor).

STATIC z_t = level + drivers_t (no state carryover) AR1 z_t = level + sum_{s<=t} rho^(t-s) * (drivers_s + sigma*eps_s) RANDOM_WALK AR1 with rho fixed at 1 (accumulating stock)

STATIC = 'static'

AR1 = 'ar1'

RANDOM_WALK = 'random_walk'

class mmm_framework.mmm_extensions.config.MediatorLikelihood(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Measurement family for a mediator’s observed data.

GAUSSIAN continuous index, z-scored, masked Normal observation BINOMIAL per-period success counts + per-period trials, logit link ORDERED per-period Likert category counts, cumulative-logit Multinomial LATENT never observed (state identified by in-graph standardization)

GAUSSIAN = 'gaussian'

BINOMIAL = 'binomial'

ORDERED = 'ordered'

LATENT = 'latent'

class mmm_framework.mmm_extensions.config.MediatorMeasurement(likelihood=MediatorLikelihood.GAUSSIAN, noise_sigma=0.3, design_effect=1.0, n_categories=None, cutpoint_prior_sigma=2.0)

Bases: object

How a mediator’s latent state is observed.

Each measured family pins the latent’s scale through its own geometry: GAUSSIAN defines the state on the standardized survey scale (loading fixed at 1), BINOMIAL pins it absolutely on the logit/probability scale, ORDERED anchors location in the cutpoints and scale against the unit-logistic response noise. design_effect deflates survey information for clustered or weighted samples (n_eff = n / design_effect).

likelihood: MediatorLikelihood = 'gaussian'

noise_sigma: float = 0.3

design_effect: float = 1.0

n_categories: int | None = None

cutpoint_prior_sigma: float = 2.0

__init__(likelihood=MediatorLikelihood.GAUSSIAN, noise_sigma=0.3, design_effect=1.0, n_categories=None, cutpoint_prior_sigma=2.0)

class mmm_framework.mmm_extensions.config.MediatorSpec(name, channels=(), parents=(), controls=(), latent_factors=(), dynamics=MediatorDynamics.STATIC, rho_prior_alpha=6.0, rho_prior_beta=2.0, innovation_sigma=0.3, state_parameterization='auto', measurement=<factory>, media_effect=<factory>, parent_effect=<factory>, control_effect=<factory>, outcome_effect=<factory>, affects_outcome=True, allow_direct_effect=True, direct_effect=<factory>, apply_adstock=None)

Bases: object

One structural mediator equation in a StructuralNestedMMM.

The latent state is driven by media channels (saturated, optionally adstocked), upstream mediators (parents – must form a DAG), control columns (e.g. price), and shared latent factors; its dynamics are STATIC, AR1, or RANDOM_WALK; and it is observed through measurement.

apply_adstock=None (default) resolves by dynamics: adstock ON for STATIC equations, OFF for AR1/RANDOM_WALK – the state itself carries the media effect forward there, and adstock + AR would be two nearly-interchangeable geometric carryovers (an alpha/rho ridge). Set it explicitly to override (overriding to True on an AR equation warns at build).

direct_effect defaults tight (Normal(0, 0.3)): an over-wide direct path steals the mediated signal (see technical-docs/nested-recovery-search.md).

name: str

channels: tuple[str, ...] = ()

parents: tuple[str, ...] = ()

controls: tuple[str, ...] = ()

latent_factors: tuple[str, ...] = ()

dynamics: MediatorDynamics = 'static'

rho_prior_alpha: float = 6.0

rho_prior_beta: float = 2.0

innovation_sigma: float = 0.3

state_parameterization: str = 'auto'

measurement: MediatorMeasurement

media_effect: EffectPriorConfig

parent_effect: EffectPriorConfig

control_effect: EffectPriorConfig

outcome_effect: EffectPriorConfig

affects_outcome: bool = True

allow_direct_effect: bool = True

direct_effect: EffectPriorConfig

apply_adstock: bool | None = None

property adstock_enabled: bool

The resolved adstock setting (dynamics-dependent when unset).

__init__(name, channels=(), parents=(), controls=(), latent_factors=(), dynamics=MediatorDynamics.STATIC, rho_prior_alpha=6.0, rho_prior_beta=2.0, innovation_sigma=0.3, state_parameterization='auto', measurement=<factory>, media_effect=<factory>, parent_effect=<factory>, control_effect=<factory>, outcome_effect=<factory>, affects_outcome=True, allow_direct_effect=True, direct_effect=<factory>, apply_adstock=None)

class mmm_framework.mmm_extensions.config.LatentFactorSpec(name, dynamics=MediatorDynamics.AR1, rho_prior_alpha=8.0, rho_prior_beta=2.0, affects_outcome=True, outcome_effect_sigma=1.0, mediator_effect_sigma=1.0, anchor='auto')

Bases: object

A shared latent factor (e.g. a demand trend) entering one or more mediator equations and/or the outcome.

The realized series is standardized in-graph to unit variance (the scale would otherwise trade off against the loadings), so loadings carry the factor’s units. Sign is anchored at the outcome loading (HalfNormal) when affects_outcome is True, else at the first consuming mediator’s loading; all other loadings are free-sign Normal.

name: str

dynamics: MediatorDynamics = 'ar1'

rho_prior_alpha: float = 8.0

rho_prior_beta: float = 2.0

affects_outcome: bool = True

outcome_effect_sigma: float = 1.0

mediator_effect_sigma: float = 1.0

anchor: str = 'auto'

__init__(name, dynamics=MediatorDynamics.AR1, rho_prior_alpha=8.0, rho_prior_beta=2.0, affects_outcome=True, outcome_effect_sigma=1.0, mediator_effect_sigma=1.0, anchor='auto')

class mmm_framework.mmm_extensions.config.StructuralNestedConfig(mediators=<factory>, latent_factors=<factory>, outcome_controls=None, nonmediated_effect=<factory>)

Bases: object

Configuration for StructuralNestedMMM – a DAG of mediator equations with per-mediator dynamics + measurement, shared latent factors, and an outcome equation.

outcome_controls=None means every control column provided to the model also enters the outcome equation. Channels not routed to any mediator get a plain direct effect with the nonmediated_effect prior.

mediators: tuple[MediatorSpec, ...]

latent_factors: tuple[LatentFactorSpec, ...]

outcome_controls: tuple[str, ...] | None = None

nonmediated_effect: EffectPriorConfig

topological_order()

Kahn topological sort of the mediator DAG (parents before children).

Return type:

list[str]

__init__(mediators=<factory>, latent_factors=<factory>, outcome_controls=None, nonmediated_effect=<factory>)

class mmm_framework.mmm_extensions.config.MediatorConfigExtended(name, observation_type=MediatorObservationType.PARTIALLY_OBSERVED, media_effect_constraint='positive', media_effect_sigma=1.0, outcome_effect_sigma=1.0, observation_noise_sigma=0.1, allow_direct_effect=True, direct_effect_sigma=0.5, apply_adstock=True, apply_saturation=True, aggregated_survey_config=None)

Bases: object

Extended MediatorConfig with aggregated survey support.

This replaces the original MediatorConfig when aggregated surveys are needed. All original fields are preserved for backward compatibility.

name: str

observation_type: MediatorObservationType = 'partially_observed'

media_effect_constraint: str = 'positive'

media_effect_sigma: float = 1.0

outcome_effect_sigma: float = 1.0

observation_noise_sigma: float = 0.1

allow_direct_effect: bool = True

direct_effect_sigma: float = 0.5

apply_adstock: bool = True

apply_saturation: bool = True

aggregated_survey_config: AggregatedSurveyConfig | None = None

__init__(name, observation_type=MediatorObservationType.PARTIALLY_OBSERVED, media_effect_constraint='positive', media_effect_sigma=1.0, outcome_effect_sigma=1.0, observation_noise_sigma=0.1, allow_direct_effect=True, direct_effect_sigma=0.5, apply_adstock=True, apply_saturation=True, aggregated_survey_config=None)

Builders

Builder Patterns for MMM Extension Configurations.

Fluent API for constructing complex configurations with sensible defaults. Each builder follows the pattern:

config = Builder().with_x(...).with_y(...).build()

class mmm_framework.mmm_extensions.builders.AdstockConfigBuilder

Bases: object

Builder for AdstockConfig.

__init__()

with_max_lag(l_max)

Set maximum lag length.

Return type:

Self

with_beta_prior(alpha=2.0, beta=2.0)

Use Beta prior for decay rate.

Return type:

Self

with_uniform_prior()

Use Uniform(0,1) prior for decay rate.

Return type:

Self

with_slow_decay()

Configure for slow decay (long memory).

Return type:

Self

with_fast_decay()

Configure for fast decay (short memory).

Return type:

Self

without_normalization()

Disable weight normalization.

Return type:

Self

build()

Build the configuration.

Return type:

AdstockConfig

class mmm_framework.mmm_extensions.builders.SaturationConfigBuilder

Bases: object

Builder for SaturationConfig.

__init__()

logistic(lam_alpha=3.0, lam_beta=1.0)

Use logistic saturation.

Return type:

Self

hill(kappa_alpha=2.0, kappa_beta=2.0, slope_alpha=3.0, slope_beta=1.0)

Use Hill saturation.

Return type:

Self

with_strong_saturation()

Configure for strong diminishing returns.

Return type:

Self

with_weak_saturation()

Configure for weak diminishing returns.

Return type:

Self

build()

Build the configuration.

Return type:

SaturationConfig

class mmm_framework.mmm_extensions.builders.EffectPriorConfigBuilder

Bases: object

Builder for EffectPriorConfig.

__init__()

unconstrained(mu=0.0, sigma=1.0)

Unconstrained effect (Normal prior).

Return type:

Self

positive(sigma=1.0)

Positive-constrained effect (HalfNormal prior).

Return type:

Self

negative(sigma=1.0)

Negative-constrained effect (-HalfNormal prior).

Return type:

Self

with_tight_prior()

Tighter prior (more regularization).

Return type:

Self

with_wide_prior()

Wider prior (less regularization).

Return type:

Self

build()

Build the configuration.

Return type:

EffectPriorConfig

class mmm_framework.mmm_extensions.builders.MediatorConfigBuilder(name)

Bases: object

Builder for MediatorConfig.

__init__(name)

fully_latent()

Mediator is never observed (pure latent variable).

Return type:

Self

partially_observed(observation_noise=0.1)

Mediator observed in some periods (e.g., surveys).

Return type:

Self

fully_observed(observation_noise=0.05)

Mediator observed in all periods (e.g., traffic counters).

Return type:

Self

with_media_effect(effect)

Set media → mediator effect prior.

Return type:

Self

with_positive_media_effect(sigma=1.0)

Media should increase mediator (e.g., awareness).

Return type:

Self

with_outcome_effect(effect)

Set mediator → outcome effect prior.

Return type:

Self

with_direct_effect(sigma=0.5)

Allow direct media → outcome effect.

Return type:

Self

without_direct_effect()

No direct effect (all media effect flows through mediator).

Return type:

Self

with_adstock(config)

Set adstock configuration.

Return type:

Self

with_slow_adstock(l_max=12)

Configure slow-decaying adstock (awareness builds slowly).

Return type:

Self

without_adstock()

Disable adstock transformation.

Return type:

Self

with_saturation(config)

Set saturation configuration.

Return type:

Self

without_saturation()

Disable saturation transformation.

Return type:

Self

build()

Build the configuration.

Return type:

MediatorConfig

aggregated_survey()

Mediator observed via temporally aggregated surveys.

Return type:

Self

with_survey_config(config)

Set aggregated survey configuration.

Return type:

Self

class mmm_framework.mmm_extensions.builders.OutcomeConfigBuilder(name, column=None)

Bases: object

Builder for OutcomeConfig.

__init__(name, column=None)

with_column(column)

Set data column name.

Return type:

Self

with_intercept_prior(sigma)

Set intercept prior scale.

Return type:

Self

with_media_effect(effect)

Set media effect prior.

Return type:

Self

with_positive_media_effects(sigma=0.5)

Constrain media effects to be positive.

Return type:

Self

with_trend()

Include trend component.

Return type:

Self

without_trend()

Exclude trend component.

Return type:

Self

with_seasonality()

Include seasonality component.

Return type:

Self

without_seasonality()

Exclude seasonality component.

Return type:

Self

build()

Build the configuration.

Return type:

OutcomeConfig

class mmm_framework.mmm_extensions.builders.CrossEffectConfigBuilder(source, target)

Bases: object

Builder for CrossEffectConfig.

__init__(source, target)

cannibalization()

Source reduces target sales (negative effect).

Return type:

Self

halo()

Source increases target sales (positive effect).

Return type:

Self

unconstrained()

No predefined sign: psi ~ Normal(0, sigma), the data picks the direction.

Prefer this when you do not want to assume cannibalization vs. halo. On observed outcomes the directional effect is confounded with the residual correlation (only their sum is identified), so the estimate is a cross-outcome association, not a causal cannibalization – see cross_effect().

Return type:

Self

symmetric()

Bidirectional effect.

Return type:

Self

asymmetric()

One-way effect (default).

Return type:

Self

with_prior_sigma(sigma)

Set effect prior scale.

Return type:

Self

modulated_by_promotion(column=None)

Effect only active when source is promoted.

Return type:

Self

always_active()

Effect always present (not promotion-modulated).

Return type:

Self

with_lag(lag)

Set temporal lag (0 = contemporaneous).

Return type:

Self

lagged()

Use lagged effect (lag=1) for identification.

Return type:

Self

build()

Build the configuration.

Return type:

CrossEffectConfig

class mmm_framework.mmm_extensions.builders.NestedModelConfigBuilder

Bases: object

Builder for NestedModelConfig.

__init__()

add_mediator(mediator)

Add a mediator to the model.

Return type:

Self

with_awareness_mediator(name='awareness', observation_noise=0.15)

Add awareness mediator with typical configuration.

Return type:

Self

with_traffic_mediator(name='foot_traffic', observation_noise=0.05)

Add foot traffic mediator with typical configuration.

Return type:

Self

map_channels_to_mediator(mediator_name, channel_names)

Specify which channels affect a mediator.

Return type:

Self

share_adstock(share=True)

Share adstock parameters across mediators.

Return type:

Self

share_saturation(share=True)

Share saturation parameters across mediators.

Return type:

Self

build()

Build the configuration.

Return type:

NestedModelConfig

class mmm_framework.mmm_extensions.builders.MultivariateModelConfigBuilder

Bases: object

Builder for MultivariateModelConfig.

__init__()

add_outcome(outcome)

Add an outcome to the model.

Return type:

Self

with_outcomes(*names)

Add multiple outcomes with default configuration.

Return type:

Self

add_cross_effect(effect)

Add a cross-effect between outcomes.

Return type:

Self

with_cannibalization(source, target, promotion_column=None)

Add cannibalization effect.

Return type:

Self

with_halo_effect(source, target)

Add halo effect.

Return type:

Self

with_lkj_eta(eta)

Set LKJ correlation prior parameter.

Return type:

Self

with_weak_correlations()

Prior favoring weak correlations (eta > 1).

Return type:

Self

with_strong_correlations()

Prior allowing strong correlations (eta < 1).

Return type:

Self

share_media_adstock(share=True)

Share adstock parameters across outcomes.

Return type:

Self

share_media_saturation(share=True)

Share saturation parameters across outcomes.

Return type:

Self

share_trend(share=True)

Share trend parameters across outcomes.

Return type:

Self

share_seasonality(share=True)

Share seasonality parameters across outcomes.

Return type:

Self

build()

Build the configuration.

Return type:

MultivariateModelConfig

class mmm_framework.mmm_extensions.builders.CombinedModelConfigBuilder

Bases: object

Builder for combined nested + multivariate model.

__init__()

add_mediator(mediator)

Return type:

Self

with_awareness_mediator(name='awareness', **kwargs)

Return type:

Self

with_traffic_mediator(name='foot_traffic', **kwargs)

Return type:

Self

map_channels_to_mediator(mediator, channels)

Return type:

Self

add_outcome(outcome)

Return type:

Self

with_outcomes(*names)

Return type:

Self

add_cross_effect(effect)

Add a cross-effect between outcomes.

Return type:

Self

with_cannibalization(source, target, **kwargs)

Return type:

Self

with_halo_effect(source, target)

Return type:

Self

with_lkj_eta(eta)

Return type:

Self

map_mediator_to_outcomes(mediator_name, outcome_names)

Specify which outcomes are affected by a mediator.

Return type:

Self

build()

Build the configuration.

Return type:

CombinedModelConfig

class mmm_framework.mmm_extensions.builders.HorseshoeConfigBuilder

Bases: object

Builder for HorseshoeConfig with sensible defaults.

Examples

>>> config = (HorseshoeConfigBuilder()
...     .with_expected_nonzero(5)
...     .with_slab_scale(2.5)
...     .with_heavy_tails()
...     .build())

__init__()

with_expected_nonzero(n)

Set expected number of nonzero coefficients.

Return type:

Self

with_slab_scale(scale)

Set slab scale (max expected effect in std units).

Return type:

Self

with_slab_df(df)

Set slab degrees of freedom.

Return type:

Self

with_local_df(df)

Set local shrinkage degrees of freedom.

Return type:

Self

with_global_df(df)

Set global shrinkage degrees of freedom.

Return type:

Self

with_heavy_tails()

Configure for heavier-tailed slab (allow larger effects).

Return type:

Self

with_light_tails()

Configure for lighter-tailed slab (more regularization).

Return type:

Self

with_half_cauchy_local()

Use half-Cauchy for local shrinkage (original horseshoe).

Return type:

Self

with_aggressive_shrinkage()

Configure for more aggressive shrinkage of small effects.

Return type:

Self

build()

Build the HorseshoeConfig object.

Return type:

HorseshoeConfig

class mmm_framework.mmm_extensions.builders.SpikeSlabConfigBuilder

Bases: object

Builder for SpikeSlabConfig with sensible defaults.

Examples

>>> config = (SpikeSlabConfigBuilder()
...     .with_prior_inclusion(0.3)
...     .with_sharp_selection()
...     .build())

__init__()

with_prior_inclusion(prob)

Set prior inclusion probability.

Return type:

Self

with_spike_scale(scale)

Set spike scale (should be small, e.g., 0.01).

Return type:

Self

with_slab_scale(scale)

Set slab scale (expected magnitude of nonzero effects).

Return type:

Self

with_temperature(temp)

Set temperature for continuous relaxation.

Return type:

Self

continuous()

Use continuous relaxation (required for NUTS).

Return type:

Self

discrete()

Use discrete selection (requires Gibbs sampler).

Return type:

Self

with_sharp_selection()

Configure for sharper variable selection.

Return type:

Self

with_soft_selection()

Configure for softer variable selection.

Return type:

Self

build()

Build the SpikeSlabConfig object.

Return type:

SpikeSlabConfig

class mmm_framework.mmm_extensions.builders.LassoConfigBuilder

Bases: object

Builder for LassoConfig.

Examples

>>> config = (LassoConfigBuilder()
...     .with_regularization(2.0)
...     .adaptive()
...     .build())

__init__()

with_regularization(strength)

Set regularization strength.

Return type:

Self

adaptive()

Use adaptive LASSO with coefficient-specific penalties.

Return type:

Self

non_adaptive()

Use standard (non-adaptive) LASSO.

Return type:

Self

with_strong_regularization()

Configure for strong shrinkage.

Return type:

Self

with_weak_regularization()

Configure for weak shrinkage.

Return type:

Self

build()

Build the LassoConfig object.

Return type:

LassoConfig

class mmm_framework.mmm_extensions.builders.VariableSelectionConfigBuilder

Bases: object

Builder for VariableSelectionConfig with fluent API.

This is the main builder for configuring variable selection in MMM.

CAUSAL WARNING: Variable selection should only be applied to precision controls, not confounders. Use exclude_confounders() to ensure confounders are always included with standard priors.

Examples

>>> # Regularized horseshoe with excluded confounders
>>> config = (VariableSelectionConfigBuilder()
...     .regularized_horseshoe(expected_nonzero=5)
...     .with_slab_scale(2.0)
...     .exclude_confounders("distribution", "price", "competitor_media")
...     .build())

>>> # Spike-and-slab for explicit inclusion probabilities
>>> config = (VariableSelectionConfigBuilder()
...     .spike_slab(prior_inclusion=0.3)
...     .with_sharp_selection()
...     .apply_only_to("weather", "gas_price", "minor_holiday")
...     .build())

>>> # Using sub-builders for full control
>>> config = (VariableSelectionConfigBuilder()
...     .regularized_horseshoe()
...     .with_horseshoe_config(
...         HorseshoeConfigBuilder()
...         .with_expected_nonzero(3)
...         .with_heavy_tails()
...         .build()
...     )
...     .build())

__init__()

none()

No variable selection (standard priors for all controls).

Return type:

Self

regularized_horseshoe(expected_nonzero=3)

Use regularized horseshoe prior (recommended default).

The regularized horseshoe provides excellent shrinkage of small effects while preserving large effects, with a slab to prevent unrealistic coefficient magnitudes.

Return type:

Self

Parameters

expected_nonzeroint

Prior expectation of relevant control variables.

finnish_horseshoe(expected_nonzero=3)

Use Finnish horseshoe prior.

Mathematically equivalent to regularized horseshoe; name emphasizes the slab regularization from Piironen & Vehtari.

Return type:

Self

Parameters

expected_nonzeroint

Prior expectation of relevant control variables.

spike_slab(prior_inclusion=0.5, continuous=True)

Use spike-and-slab prior.

Provides explicit posterior inclusion probabilities for each variable.

Return type:

Self

Parameters

prior_inclusionfloat

Prior probability that each variable is included.

continuousbool

Use continuous relaxation for NUTS sampling.

bayesian_lasso(regularization=1.0)

Use Bayesian LASSO prior.

Better when expecting many small effects rather than sparse signals.

Return type:

Self

Parameters

regularizationfloat

Regularization strength (higher = more shrinkage).

with_horseshoe_config(config)

Set horseshoe configuration from pre-built config.

Return type:

Self

with_expected_nonzero(n)

Set expected number of nonzero coefficients (horseshoe).

Return type:

Self

with_slab_scale(scale)

Set slab scale for horseshoe (max effect in std units).

Return type:

Self

with_slab_df(df)

Set slab degrees of freedom (horseshoe).

Return type:

Self

with_spike_slab_config(config)

Set spike-slab configuration from pre-built config.

Return type:

Self

with_prior_inclusion(prob)

Set prior inclusion probability (spike-slab).

Return type:

Self

with_temperature(temp)

Set temperature for continuous spike-slab.

Return type:

Self

with_sharp_selection()

Configure spike-slab for sharper selection.

Return type:

Self

with_lasso_config(config)

Set LASSO configuration from pre-built config.

Return type:

Self

with_regularization(strength)

Set LASSO regularization strength.

Return type:

Self

exclude_confounders(*variables)

Exclude variables from selection (always include with standard priors).

IMPORTANT: Use this for known confounders that affect both media spending and sales. These must be controlled for to identify causal effects and should NOT be subject to variable selection.

Parameters:

*variables (str) – Variable names to exclude from selection.

Return type:

Self

exclude(*variables)

Alias for exclude_confounders.

Return type:

Self

apply_only_to(*variables)

Apply selection only to specified variables.

All other variables will use standard priors.

Parameters:

*variables (str) – Variable names to apply selection to.

Return type:

Self

clear_exclusions()

Clear all exclusions.

Return type:

Self

build()

Build the VariableSelectionConfig object.

Return type:

VariableSelectionConfig

mmm_framework.mmm_extensions.builders.awareness_mediator(name='awareness', observation_noise=0.15)

Create typical awareness mediator configuration.

Return type:

MediatorConfig

mmm_framework.mmm_extensions.builders.foot_traffic_mediator(name='foot_traffic', observation_noise=0.05)

Create typical foot traffic mediator configuration.

Return type:

MediatorConfig

mmm_framework.mmm_extensions.builders.cannibalization_effect(source, target, promotion_column=None, lagged=False)

Create cannibalization cross-effect configuration.

Return type:

CrossEffectConfig

mmm_framework.mmm_extensions.builders.halo_effect(source, target)

Create halo cross-effect configuration.

Return type:

CrossEffectConfig

mmm_framework.mmm_extensions.builders.cross_effect(source, target, prior_sigma=0.5, promotion_column=None)

Create a signless (unconstrained) cross-outcome effect source -> target.

Unlike cannibalization_effect() (which forces psi <= 0) or halo_effect() (psi >= 0), this places a symmetric Normal(0, prior_sigma) prior on psi and lets the data choose the direction. Use it when you do not want to assume cannibalization vs. halo – and note that a one-sided prior makes “P(psi < 0) ~ 1” near-automatic, so an unconstrained prior is the honest default for measuring a cross-outcome association.

Return type:

CrossEffectConfig

Identification caveat

The cross-effect enters as psi * Y_source on the observed sibling outcome, while the residual covariance also links the outcomes. Only their sum (the total residual covariance) is identified – the split between the directional psi and the symmetric residual correlation is set by the prior, not the data. So an unconstrained psi is a cross-outcome association, not a causal cannibalization estimate, and it is confounded with the shared-demand correlation already captured by the residual covariance (the actual “joint shock”). To estimate causal substitution, model it on exogenous drivers (cross-product media/price), use a share/compositional model, or run an experiment.

mmm_framework.mmm_extensions.builders.sparse_controls(expected_nonzero=3, *confounders)

Create sparse control selection configuration.

Convenience factory for the most common use case: expecting only a few control variables are truly relevant.

Parameters:

• expected_nonzero (int) – Prior expectation of relevant controls.

• *confounders (str) – Confounder variable names to exclude from selection.

Return type:

VariableSelectionConfig

Returns:

Configuration for regularized horseshoe selection.

Example:

config = sparse_controls(3, "distribution", "price")

mmm_framework.mmm_extensions.builders.selection_with_inclusion_probs(prior_inclusion=0.5, *confounders)

Create selection configuration with explicit inclusion probabilities.

Parameters:

• prior_inclusion (float) – Prior probability of inclusion for each variable.

• *confounders (str) – Confounder variable names to exclude from selection.

Return type:

VariableSelectionConfig

Returns:

Configuration for spike-slab selection.

mmm_framework.mmm_extensions.builders.dense_controls(regularization=1.0, *confounders)

Create dense control selection configuration.

Use when expecting many controls have small effects.

Parameters:

• regularization (float) – Regularization strength.

• *confounders (str) – Confounder variable names to exclude from selection.

Return type:

VariableSelectionConfig

Returns:

Configuration for Bayesian LASSO selection.

class mmm_framework.mmm_extensions.builders.AggregatedSurveyConfigBuilder

Bases: object

Builder for AggregatedSurveyConfig.

__init__()

with_aggregation_map(aggregation_map)

Set the aggregation map directly.

Return type:

AggregatedSurveyConfigBuilder

from_frequencies(model_frequency, survey_frequency, n_periods, start_date=None)

Compute aggregation map from frequency specifications.

Return type:

AggregatedSurveyConfigBuilder

monthly_in_weekly_model(n_weeks, start_date=None)

Convenience for monthly surveys in weekly model.

Return type:

AggregatedSurveyConfigBuilder

quarterly_in_weekly_model(n_weeks, start_date=None)

Convenience for quarterly surveys in weekly model.

Return type:

AggregatedSurveyConfigBuilder

with_sample_sizes(sample_sizes)

Set sample sizes per survey wave.

Return type:

AggregatedSurveyConfigBuilder

with_constant_sample_size(n)

Set constant sample size across all waves.

Return type:

AggregatedSurveyConfigBuilder

binomial()

Use exact binomial likelihood (default, preferred).

Return type:

AggregatedSurveyConfigBuilder

normal_approximation()

Use normal approximation with derived SEs.

Return type:

AggregatedSurveyConfigBuilder

beta_binomial(overdispersion_prior_sigma=0.1)

Use beta-binomial for overdispersed data.

Return type:

AggregatedSurveyConfigBuilder

with_design_effect(deff)

Set survey design effect.

Use >1.0 for clustered samples, complex weighting, etc. Common values: 1.5-2.5 for area probability samples.

Return type:

AggregatedSurveyConfigBuilder

with_effective_sample_size(nominal_n, effective_n)

Compute design effect from nominal and effective sample sizes.

Return type:

AggregatedSurveyConfigBuilder

aggregate_by_mean()

Average latent values within survey period (default).

Return type:

AggregatedSurveyConfigBuilder

aggregate_by_last()

Use last latent value in survey period.

Return type:

AggregatedSurveyConfigBuilder

build()

Build the configuration.

Return type:

AggregatedSurveyConfig

class mmm_framework.mmm_extensions.builders.MediatorConfigBuilderExtended(name)

Bases: object

Extended MediatorConfigBuilder with aggregated survey support.

Drop-in replacement for MediatorConfigBuilder that adds methods for configuring aggregated survey observations.

Example

>>> config = (
...     MediatorConfigBuilderExtended("brand_awareness")
...     .aggregated_survey()
...     .with_survey_config(
...         AggregatedSurveyConfigBuilder()
...         .monthly_in_weekly_model(104)
...         .with_sample_sizes([500, 450, 520, ...])
...         .binomial()
...         .build()
...     )
...     .with_positive_media_effect()
...     .build()
... )

__init__(name)

fully_latent()

Mediator is never observed.

Return type:

MediatorConfigBuilderExtended

partially_observed(observation_noise=0.1)

Point-in-time observations (sparse).

Return type:

MediatorConfigBuilderExtended

fully_observed(observation_noise=0.05)

Observations every period.

Return type:

MediatorConfigBuilderExtended

aggregated_survey()

Mediator observed via temporally aggregated surveys.

Must call with_survey_config() after this.

Return type:

MediatorConfigBuilderExtended

with_survey_config(config)

Set aggregated survey configuration.

Return type:

MediatorConfigBuilderExtended

with_positive_media_effect(sigma=1.0)

Media increases mediator.

Return type:

MediatorConfigBuilderExtended

with_unconstrained_media_effect(sigma=1.0)

No sign constraint on media effect.

Return type:

MediatorConfigBuilderExtended

with_direct_effect(sigma=0.5)

Allow direct media → outcome effect.

Return type:

MediatorConfigBuilderExtended

without_direct_effect()

All media effect flows through mediator.

Return type:

MediatorConfigBuilderExtended

with_adstock()

Apply adstock to media → mediator pathway.

Return type:

MediatorConfigBuilderExtended

without_adstock()

No adstock on media → mediator pathway.

Return type:

MediatorConfigBuilderExtended

with_saturation()

Apply saturation to media → mediator pathway.

Return type:

MediatorConfigBuilderExtended

without_saturation()

No saturation on media → mediator pathway.

Return type:

MediatorConfigBuilderExtended

build()

Build the extended mediator configuration.

Return type:

MediatorConfigExtended

mmm_framework.mmm_extensions.builders.survey_awareness_mediator(name='awareness', n_model_periods=104, sample_sizes=500, model_frequency='weekly', survey_frequency='monthly', start_date=None, design_effect=1.0, use_beta_binomial=False)

Create awareness mediator with aggregated survey observation model.

Convenience factory for the common case of brand tracking surveys fielded continuously over monthly periods.

Return type:

MediatorConfigExtended

Parameters

namestr

Mediator name.

n_model_periodsint

Number of periods in the model (e.g., weeks).

sample_sizeslist[int] or int

Sample sizes per survey wave. If int, uses constant size.

model_frequencystr

Model time frequency (“weekly” or “daily”).

survey_frequencystr

Survey aggregation frequency (“monthly” or “quarterly”).

start_datestr, optional

Start date for calendar-based aggregation.

design_effectfloat

Survey design effect (default 1.0).

use_beta_binomialbool

If True, use beta-binomial for overdispersion.

Returns

MediatorConfigExtended

Configured mediator with aggregated survey observation.

Example

>>> mediator = survey_awareness_mediator(
...     name="brand_awareness",
...     n_model_periods=104,  # 2 years weekly
...     sample_sizes=[500, 480, 520, 490, ...],  # per month
...     design_effect=1.5,  # clustered sample
... )

mmm_framework.mmm_extensions.builders.binary_survey_mediator(name, channels, *, parents=(), controls=(), latent_factors=(), persistence='high', design_effect=1.0, affects_outcome=True, allow_direct_effect=True, direct_effect_sigma=0.3)

A binary-tracker mediator (“have you seen this brand?”): AR(1) latent state on the logit scale, binomial measurement with per-week trials.

persistence sets the Beta prior on the AR(1) rho: “high” (Beta(9, 1.5), awareness-like carryover), “medium” (Beta(6, 2)), “low” (Beta(2, 4)). The AR(1) state carries the media effect forward, so adstock is OFF for this equation (adstock + AR(1) would double-count carryover).

Return type:

MediatorSpec

mmm_framework.mmm_extensions.builders.likert_mediator(name, channels, *, n_categories=5, parents=(), controls=(), latent_factors=(), dynamics=MediatorDynamics.STATIC, design_effect=1.0, affects_outcome=True, allow_direct_effect=True, direct_effect_sigma=0.3)

A Likert-measured mediator (e.g. consideration): cumulative-logit Multinomial measurement over per-week category counts. Location lives in the cutpoints; feed upstream mediators via parents (e.g. awareness) and drivers like price via controls.

Return type:

MediatorSpec

mmm_framework.mmm_extensions.builders.latent_demand_factor(name='demand', *, affects_outcome=True, outcome_effect_sigma=1.0, mediator_effect_sigma=1.0)

A smooth latent demand trend: AR(1) with a high-persistence prior (Beta(9, 1.5)), unit-standardized in-graph, sign-anchored at the outcome loading. Route it into mediator equations via MediatorSpec.latent_factors=(name,).

Return type:

LatentFactorSpec

Base Builders

Base classes and utilities for extension builders.

This module provides shared functionality for extension builders that operate on extension-specific config classes (frozen dataclasses).

Note: The extension builders use different config classes than the base builders (frozen dataclasses vs Pydantic models), so they cannot directly inherit. This module provides common patterns and utilities.

class mmm_framework.mmm_extensions.builders_base.ExtensionBuilderProtocol(*args, **kwargs)

Bases: Protocol[T]

Protocol for extension configuration builders.

All extension builders should implement a build() method that returns the final configuration object.

build()

Build and return the configuration object.

Return type:

TypeVar(T)

__init__(*args, **kwargs)

class mmm_framework.mmm_extensions.builders_base.AdstockBuilderMixin

Bases: object

Mixin providing common adstock configuration methods.

This mixin provides shared convenience methods for configuring adstock parameters. It can be used by any builder that configures adstock.

Attributes

_l_maxint

Maximum lag for adstock transformation.

_prior_alphafloat

Alpha parameter for Beta prior on decay rate.

_prior_betafloat

Beta parameter for Beta prior on decay rate.

_normalizebool

Whether to normalize adstock weights.

with_max_lag(l_max)

Set maximum lag length.

Return type:

Self

Parameters

l_maxint

Maximum number of periods for carryover effects.

Returns

Self

Builder instance for method chaining.

with_slow_decay()

Configure for slow decay (long memory).

Sets Beta(3, 1) prior which favors higher decay rates, meaning effects persist longer.

Return type:

Self

Returns

Self

Builder instance for method chaining.

with_fast_decay()

Configure for fast decay (short memory).

Sets Beta(1, 3) prior which favors lower decay rates, meaning effects dissipate quickly.

Return type:

Self

Returns

Self

Builder instance for method chaining.

without_normalization()

Disable weight normalization.

Return type:

Self

Returns

Self

Builder instance for method chaining.

class mmm_framework.mmm_extensions.builders_base.SaturationBuilderMixin

Bases: object

Mixin providing common saturation configuration methods.

This mixin provides shared convenience methods for configuring saturation parameters. It can be used by any builder that configures saturation.

Attributes

_lam_alphafloat

Alpha parameter for lambda prior (logistic saturation).

_lam_betafloat

Beta parameter for lambda prior.

_kappa_alphafloat

Alpha parameter for kappa prior (Hill saturation).

_kappa_betafloat

Beta parameter for kappa prior.

_slope_alphafloat

Alpha parameter for slope prior (Hill saturation).

_slope_betafloat

Beta parameter for slope prior.

with_strong_saturation()

Configure for strong diminishing returns.

Return type:

Self

Returns

Self

Builder instance for method chaining.

with_weak_saturation()

Configure for weak diminishing returns.

Return type:

Self

Returns

Self

Builder instance for method chaining.

class mmm_framework.mmm_extensions.builders_base.EffectPriorMixin

Bases: object

Mixin for configuring effect prior parameters.

This mixin provides shared methods for configuring effect priors (constraint type and scale).

Components

Components subpackage for MMM Extensions.

This package provides modular, composable building blocks for nested and multivariate models using strategy patterns.

Contains both NumPy (for preprocessing) and PyTensor (for model building) versions of transforms.

mmm_framework.mmm_extensions.components.geometric_adstock(x, alpha, l_max=8, normalize=True)

Apply geometric adstock using matrix multiplication (no scan).

This is often more efficient and avoids scan complexity. Requires knowing n_obs at graph construction time.

Return type:

TensorVariable

Parameters

xTensorVariable

Input media variable (n_obs,)

alphaTensorVariable

Decay rate [0, 1]

l_maxint

Maximum lag length

normalizebool

Whether to normalize weights to sum to 1

Returns

TensorVariable

Adstocked media variable

mmm_framework.mmm_extensions.components.geometric_adstock_2d(X, alpha)

Apply geometric adstock to a 2D array (multiple channels).

Applies the geometric adstock transformation independently to each column (channel) of the input matrix.

Parameters

XNDArray[np.floating]

Input matrix of shape (n_periods, n_channels).

alphafloat

Decay rate parameter in [0, 1).

Returns

NDArray[np.floating]

Adstocked matrix, same shape as input.

Examples

>>> import numpy as np
>>> from mmm_framework.transforms import geometric_adstock_2d
>>>
>>> # Two channels with different spend patterns
>>> X = np.array([
...     [100.0, 50.0],
...     [0.0, 50.0],
...     [0.0, 0.0],
... ])
>>> adstocked = geometric_adstock_2d(X, alpha=0.5)

See Also

geometric_adstock : 1D version of this function.

mmm_framework.mmm_extensions.components.logistic_saturation(x, lam)

Apply logistic saturation transformation (PyTensor version).

Return type:

TensorVariable

Parameters

xTensorVariable

Input (already adstocked)

lamTensorVariable

Saturation rate (higher = faster saturation)

Returns

TensorVariable

Saturated output in [0, 1]

mmm_framework.mmm_extensions.components.hill_saturation(x, kappa, slope)

Apply Hill saturation transformation (PyTensor version).

Return type:

TensorVariable

Parameters

xTensorVariable

Input (already adstocked)

kappaTensorVariable

Half-saturation point (EC50)

slopeTensorVariable

Steepness of curve

Returns

TensorVariable

Saturated output in [0, 1]

mmm_framework.mmm_extensions.components.create_fourier_features(t, period, order)

Create Fourier features for capturing seasonality.

Generates sine and cosine features at multiple harmonics of the specified period. This is the standard approach for modeling periodic patterns in time series (e.g., weekly, yearly seasonality).

Parameters

tNDArray[np.floating]

Time index values. Can be any numeric scale (e.g., week numbers, day of year, etc.).

periodfloat

The fundamental period length in the same units as t. For example, if t is in weeks, period=52 captures yearly seasonality.

orderint

Number of Fourier terms (harmonics) to include. Higher order captures more complex seasonal patterns but may overfit. order=0 returns an empty array.

Returns

NDArray[np.floating]

Feature matrix of shape (len(t), 2 * order). Columns are [sin_1, cos_1, sin_2, cos_2, …, sin_order, cos_order]. Returns shape (len(t), 0) if order=0.

Examples

>>> import numpy as np
>>> from mmm_framework.transforms import create_fourier_features
>>>
>>> # Weekly data with yearly seasonality
>>> weeks = np.arange(104)  # 2 years of data
>>> features = create_fourier_features(weeks, period=52.0, order=3)
>>> print(features.shape)
(104, 6)
>>>
>>> # Values repeat after one period
>>> np.allclose(features[0], features[52])
True

Notes

The Fourier features at order k are:

sin(2 * pi * k * t / period) cos(2 * pi * k * t / period)

Using both sine and cosine allows the model to capture phase shifts in the seasonal pattern.

For most applications: - order=3-4 is sufficient for smooth seasonal patterns - order=6-10 can capture more complex patterns - Very high order risks overfitting and should be used with

regularization

See Also

Prophet (Facebook) uses this same approach for seasonality modeling.

mmm_framework.mmm_extensions.components.create_bspline_basis(t, n_knots, degree=3)

Create B-spline basis matrix for flexible trend modeling.

B-splines are piecewise polynomial functions that provide a smooth, flexible way to model trends. The basis functions are localized, which helps with interpretability and numerical stability.

Parameters

tNDArray[np.floating]

Time values, should be scaled to [0, 1] for best results.

n_knotsint

Number of interior knots. More knots allow more flexibility but risk overfitting.

degreeint, default=3

Spline degree. degree=3 gives cubic splines, which are smooth through the second derivative.

Returns

NDArray[np.floating]

Basis matrix of shape (len(t), n_knots + degree + 1). Each column is a B-spline basis function evaluated at t.

Examples

>>> import numpy as np
>>> from mmm_framework.transforms import create_bspline_basis
>>>
>>> t = np.linspace(0, 1, 100)
>>> basis = create_bspline_basis(t, n_knots=5, degree=3)
>>> print(basis.shape)
(100, 9)
>>>
>>> # Basis sums to 1 (partition of unity)
>>> np.allclose(basis.sum(axis=1), 1.0)
True

Notes

The basis is “clamped” at the boundaries, meaning the first and last basis functions are 1 at t=0 and t=1 respectively. This ensures the fitted curve passes through the endpoint predictions.

In a Bayesian model, coefficients for each basis function are given priors, and the posterior captures uncertainty in the trend.

Raises

ImportError

If scipy is not installed.

See Also

create_piecewise_trend_matrix : Alternative trend representation.

mmm_framework.mmm_extensions.components.geometric_adstock_np(x, alpha)

Apply geometric adstock transformation to a 1D array.

Implements the recurrence relation:

y[t] = x[t] + alpha * y[t-1]

where y[0] = x[0].

This models the carryover effect of marketing activities, where past spending continues to have an effect in future periods, decaying exponentially with rate alpha.

Parameters

xNDArray[np.floating]

Input time series (e.g., media spend), shape (n_periods,).

alphafloat

Decay rate parameter in [0, 1). Higher values mean slower decay (longer-lasting effects). alpha=0 means no carryover.

Returns

NDArray[np.floating]

Adstocked time series, same shape as input.

Examples

>>> import numpy as np
>>> from mmm_framework.transforms import geometric_adstock
>>>
>>> # Single pulse of spend
>>> spend = np.array([100.0, 0.0, 0.0, 0.0, 0.0])
>>> adstocked = geometric_adstock(spend, alpha=0.5)
>>> print(adstocked)
[100.  50.  25.  12.5  6.25]

Notes

The sum of the adstock weights is 1/(1-alpha), so the total effect of a unit spend is scaled by this factor. For alpha=0.5, total effect is 2x the immediate effect.

mmm_framework.mmm_extensions.components.geometric_adstock_pt(x, alpha, l_max=8, normalize=True)

Apply geometric adstock transformation using PyTensor scan.

This version uses scan for proper gradient flow in complex models. For most use cases, geometric_adstock_convolution is preferred.

Return type:

TensorVariable

Parameters

xTensorVariable

Input media variable (n_obs,)

alphaTensorVariable

Decay rate [0, 1]

l_maxint

Maximum lag length

normalizebool

Whether to normalize weights to sum to 1

Returns

TensorVariable

Adstocked media variable

mmm_framework.mmm_extensions.components.geometric_adstock_convolution(x, alpha, l_max=8, normalize=True)

Apply geometric adstock using matrix multiplication (no scan).

This is often more efficient and avoids scan complexity. Requires knowing n_obs at graph construction time.

Return type:

TensorVariable

Parameters

xTensorVariable

Input media variable (n_obs,)

alphaTensorVariable

Decay rate [0, 1]

l_maxint

Maximum lag length

normalizebool

Whether to normalize weights to sum to 1

Returns

TensorVariable

Adstocked media variable

mmm_framework.mmm_extensions.components.geometric_adstock_matrix(X, alphas, l_max=8)

Apply geometric adstock to multiple channels.

Return type:

TensorVariable

Parameters

XTensorVariable

Media matrix (n_obs, n_channels)

alphasTensorVariable

Decay rates per channel (n_channels,)

l_maxint

Maximum lag length

Returns

TensorVariable

Adstocked media matrix (n_obs, n_channels)

mmm_framework.mmm_extensions.components.apply_transformation_pipeline(x, transforms)

Apply a sequence of transformations.

Return type:

TensorVariable

Parameters

xTensorVariable

Input variable

transformslist[tuple[Callable, dict]]

List of (transform_fn, params_dict) tuples

Returns

TensorVariable

Transformed output

mmm_framework.mmm_extensions.components.create_adstock_prior(name, prior_type='beta', **kwargs)

Create adstock decay prior.

Parameters:

• name (str) – Parameter name.

• prior_type (str) – Prior type, either “beta” or “uniform”.

• **kwargs – Additional prior parameters.

Return type:

TensorVariable

Returns:

Prior random variable.

mmm_framework.mmm_extensions.components.create_saturation_prior(name, saturation_type='logistic', **kwargs)

Create saturation parameter priors.

Parameters:

• name (str) – Base parameter name.

• saturation_type (str) – Saturation type, either “logistic” or “hill”.

• **kwargs – Prior hyperparameters.

Return type:

dict[str, TensorVariable]

Returns:

Dictionary of prior random variables.

mmm_framework.mmm_extensions.components.create_effect_prior(name, constrained='none', mu=0.0, sigma=1.0, dims=None)

Create effect coefficient prior.

Return type:

TensorVariable

Parameters

namestr

Parameter name

constrainedstr

“none”, “positive”, “negative”

mufloat

Prior mean (for unconstrained)

sigmafloat

Prior scale

dimsstr | tuple | None

PyMC dimensions

Returns

TensorVariable

Prior random variable

class mmm_framework.mmm_extensions.components.MediaTransformResult(transformed, adstock_params, saturation_params)

Bases: object

Result of media transformation.

transformed: TensorVariable

adstock_params: dict[str, TensorVariable]

saturation_params: dict[str, TensorVariable]

__init__(transformed, adstock_params, saturation_params)

class mmm_framework.mmm_extensions.components.EffectResult(contribution, coefficients, components=None)

Bases: object

Result of effect computation.

contribution: pt.TensorVariable

coefficients: pt.TensorVariable

components: pt.TensorVariable | None = None

__init__(contribution, coefficients, components=None)

mmm_framework.mmm_extensions.components.build_media_transforms(X_media, channel_names, adstock_config, saturation_config, share_params=False, name_prefix='')

Build media transformation block.

Return type:

MediaTransformResult

Parameters

X_mediaTensorVariable

Raw media matrix (n_obs, n_channels)

channel_nameslist[str]

Channel names

adstock_configdict

Adstock configuration

saturation_configdict

Saturation configuration

share_paramsbool

Whether to share parameters across channels

name_prefixstr

Prefix for parameter names

Returns

MediaTransformResult

Transformed media and parameters

mmm_framework.mmm_extensions.components.build_linear_effect(X, var_names, name_prefix, constrained='none', prior_sigma=0.5, dims=None)

Build linear effect block.

Return type:

EffectResult

Parameters

XTensorVariable

Design matrix (n_obs, n_vars)

var_nameslist[str]

Variable names

name_prefixstr

Prefix for parameter names

constrainedstr

Constraint type

prior_sigmafloat

Prior scale

dimsstr | None

Dimension name for coefficients

Returns

EffectResult

Effect contribution and coefficients

mmm_framework.mmm_extensions.components.build_gaussian_likelihood(name, mu, observed, sigma_prior_sigma=0.5, dims=None)

Build Gaussian likelihood.

Return type:

tuple[TensorVariable, TensorVariable]

Parameters

namestr

Variable name

muTensorVariable

Expected value

observednp.ndarray

Observed data

sigma_prior_sigmafloat

Prior on observation noise

dimsstr | None

PyMC dimension

Returns

tuple

(likelihood_rv, sigma_rv)

mmm_framework.mmm_extensions.components.build_partial_observation_model(name, latent, observed, mask, sigma_prior_sigma=0.1)

Build observation model for partially observed variable.

Return type:

tuple[pt.TensorVariable | None, pt.TensorVariable]

Parameters

namestr

Variable name

latentTensorVariable

Latent true values

observednp.ndarray

Observed values (with NaN for missing)

masknp.ndarray

Boolean mask (True = observed)

sigma_prior_sigmafloat

Prior on measurement noise

Returns

tuple

(likelihood_rv or None, sigma_rv)

mmm_framework.mmm_extensions.components.build_multivariate_likelihood(name, mu, observed, n_outcomes, lkj_eta=2.0, sigma_prior_sigma=0.5, dims=None)

Build multivariate normal likelihood with LKJ correlation prior.

Return type:

tuple[TensorVariable, TensorVariable, TensorVariable]

Parameters

namestr

Variable name

muTensorVariable

Expected values (n_obs, n_outcomes)

observednp.ndarray

Observed data

n_outcomesint

Number of outcomes

lkj_etafloat

LKJ prior concentration parameter

sigma_prior_sigmafloat

Prior on outcome standard deviations

dimstuple | None

PyMC dimensions

Returns

tuple

(likelihood_rv, cholesky_factor, correlation_matrix)

mmm_framework.mmm_extensions.components.build_aggregated_survey_observation(name, latent, observed_data, config, is_proportion=True)

Build observation model for temporally aggregated survey data.

This handles the case where: 1. Surveys are fielded continuously over a period (e.g., month) 2. Results are aggregated across that period 3. Sample size varies by wave (heteroskedastic observation noise)

Return type:

None

Parameters

namestr

Name prefix for PyMC variables.

latentTensorVariable

Latent mediator values at model frequency (e.g., weekly).

observed_datanp.ndarray

Observed survey results.

configAggregatedSurveyConfig

Configuration for the aggregated observation model.

is_proportionbool

If True and using binomial, observed_data is proportions.

mmm_framework.mmm_extensions.components.compute_survey_observation_indices(model_frequency, survey_frequency, n_periods, start_date=None)

Compute aggregation map from model and survey frequencies.

Return type:

dict[int, tuple[int, ...]]

Parameters

model_frequencystr

Model time frequency: “daily”, “weekly”

survey_frequencystr

Survey aggregation frequency: “weekly”, “monthly”, “quarterly”

n_periodsint

Number of model periods.

start_datestr, optional

Start date for calendar-based aggregation (ISO format).

Returns

dict[int, tuple[int, …]]

Aggregation map suitable for AggregatedSurveyConfig.

mmm_framework.mmm_extensions.components.build_mediator_observation_dispatch(med_config, mediator_latent, mediator_data, mediator_masks)

Dispatch to appropriate observation model based on mediator type.

Return type:

None

Parameters

med_configMediatorConfigExtended

Mediator configuration.

mediator_latentTensorVariable

Latent mediator values.

mediator_datadict[str, np.ndarray]

Observed mediator data.

mediator_masksdict[str, np.ndarray]

Observation masks for partial observation.

mmm_framework.mmm_extensions.components.ar1_decay_matrix(rho, n_obs)

Lower-triangular Toeplitz matrix D[t, s] = rho^(t-s) for s <= t.

D @ u realizes the AR(1) accumulation of the impulse series u. rho may be an RV or a constant tensor (RANDOM_WALK passes 1.0).

Return type:

TensorVariable

mmm_framework.mmm_extensions.components.build_latent_state(name, drivers, dynamics, n_obs, *, level=0.0, rho_prior_alpha=6.0, rho_prior_beta=2.0, innovation_sigma=None, centered=False)

Realize a latent state series from its driver sum.

STATIC -> level + drivers AR1 -> level + D(rho) @ drivers + AR-noise RANDOM_WALK -> AR1 with rho fixed at 1 (no rho RV)

The level intercept sits OUTSIDE the recursion (state deviations follow z_t - level = rho * (z_{t-1} - level) + u_t), so the baseline does not ramp up from zero over a burn-in window while media impulses correctly accumulate through the state (sustained unit driver -> steady-state lift 1/(1-rho)).

innovation_sigma=None builds a deterministic state (no process noise): required for unmeasured mediators, where n_obs free innovations with no measurement would just absorb outcome residual.

centered picks the AR-noise parameterization (dynamic states only; the two are the SAME model, different sampling geometry): :rtype: pt.TensorVariable

• non-centered (default): unit innovations {name}_innovation ~ N(0, 1) scaled by sigma and accumulated through the decay matrix. Right for WEAK measurements (sparse/noisy surveys), where the state is mostly prior.

• centered: the AR noise {name}_state_noise is sampled directly (pm.AR / pm.GaussianRandomWalk, zero-history init N(0, sigma)). Right for STRONG measurements (a dense high-n_t tracker pins z_t, and the non-centered form then puts a funnel between sigma and every innovation).

Must be called inside a pm.Model context.

mmm_framework.mmm_extensions.components.standardize_in_graph(x)

Center + scale a realized series to unit empirical variance in-graph.

Pins the scale of an unmeasured latent FACTOR so its loadings carry the identified units – without this the AR(1) marginal variance 1/(1-rho^2) trades off against the loadings and the sampler can collapse them (the latent_factor_mmm lesson).

ONLY safe for media-independent series (factors built purely from innovation RVs): a media-driven latent standardized this way would have its constants recomputed under a counterfactual set_data swap, contaminating every contrast – which is why LATENT mediators are not standardized (their beta*gamma products are identified through the priors).

The variance guard sits INSIDE the sqrt: at the model’s initial point the innovations are all zero, so x.std() has a NaN gradient at var == 0 (0.5/sqrt(0) * 0) and MAP/ADVI would silently stall at the start point.

Return type:

TensorVariable

mmm_framework.mmm_extensions.components.build_gaussian_state_measurement(name, latent, observed_std, mask, noise_sigma=0.3)

Masked Normal observation of a latent state on the STANDARDIZED scale.

observed_std must already be z-scored over its observed entries (the caller owns data prep); the latent is thereby defined on the standardized survey scale with loading fixed at 1 – the survey pins media -> mediator.

Return type:

None

mmm_framework.mmm_extensions.components.build_binomial_state_measurement(name, latent, counts, trials, mask, design_effect=1.0)

Binomial survey observation of a latent state through a logit link.

p_t = sigmoid(z_t); counts_t ~ Binomial(n_t, p_t) on observed weeks. Weekly-varying trials gives each week exactly its finite-sample precision (variance n_t * p * (1-p)); the logit link replaces the old pt.clip-as-probability construction (dead gradients outside [0, 1]) and pins the latent’s location AND scale absolutely.

Returns the full-length probability tensor sigmoid(latent) so the caller can register it as a Deterministic / feed it downstream.

Return type:

TensorVariable

mmm_framework.mmm_extensions.components.build_ordered_state_measurement(name, latent, counts, mask, n_categories, cutpoint_prior_sigma=2.0, design_effect=1.0)

Cumulative-logit Multinomial observation of Likert category counts.

P(Y <= k) = sigmoid(c_k - z_t) with ordered cutpoints c; per-week category counts (rows of counts, shape (n_obs, K)) are Multinomial with n_t = the week’s row sum – so weekly-varying Likert sample sizes carry their own precision. Location lives in the cutpoints (the mediator equation must NOT have a free level – exactly confounded); scale is identified against the implicit unit-logistic response noise.

Return type:

None

class mmm_framework.mmm_extensions.components.CrossEffectSpec(source_idx, target_idx, effect_type, prior_sigma=0.3)

Bases: object

Specification for a single cross-effect.

source_idx: int

target_idx: int

effect_type: str

prior_sigma: float = 0.3

__init__(source_idx, target_idx, effect_type, prior_sigma=0.3)

mmm_framework.mmm_extensions.components.build_cross_effect_matrix(specs, n_outcomes, name_prefix='psi')

Build cross-effect coefficient matrix.

Return type:

tuple[TensorVariable, dict[tuple[int, int], TensorVariable]]

Parameters

specslist[CrossEffectSpec]

Cross-effect specifications

n_outcomesint

Number of outcomes

name_prefixstr

Prefix for parameter names

Returns

tuple

(cross_effect_matrix, individual_params_dict)

mmm_framework.mmm_extensions.components.compute_cross_effect_contribution(Y, psi_matrix, target_idx, n_outcomes, modulation=None)

Compute cross-effect contribution for a single target outcome.

Return type:

TensorVariable

Parameters

YTensorVariable

Outcome matrix (n_obs, n_outcomes)

psi_matrixTensorVariable

Cross-effect coefficients (n_outcomes, n_outcomes)

target_idxint

Index of target outcome

n_outcomesint

Total number of outcomes

modulationdict | None

Optional modulation by source index (e.g., promotion indicators)

Returns

TensorVariable

Cross-effect contribution (n_obs,)

class mmm_framework.mmm_extensions.components.VariableSelectionResult(beta, inclusion_indicators=None, local_shrinkage=None, global_shrinkage=None, effective_nonzero=None, kappa=None)

Bases: object

Container for variable selection prior outputs.

Attributes

betapt.TensorVariable

The coefficient vector with shrinkage/selection applied.

inclusion_indicatorspt.TensorVariable | None

For spike-slab: soft inclusion indicators (gamma).

local_shrinkagept.TensorVariable | None

For horseshoe: local shrinkage parameters (lambda).

global_shrinkagept.TensorVariable | None

For horseshoe: global shrinkage parameter (tau).

effective_nonzeropt.TensorVariable | None

Estimated number of effectively nonzero coefficients.

kappapt.TensorVariable | None

Shrinkage factors for each coefficient (horseshoe).

beta: pt.TensorVariable

inclusion_indicators: pt.TensorVariable | None = None

local_shrinkage: pt.TensorVariable | None = None

global_shrinkage: pt.TensorVariable | None = None

effective_nonzero: pt.TensorVariable | None = None

kappa: pt.TensorVariable | None = None

__init__(beta, inclusion_indicators=None, local_shrinkage=None, global_shrinkage=None, effective_nonzero=None, kappa=None)

class mmm_framework.mmm_extensions.components.ControlEffectResult(contribution, beta_selected=None, beta_fixed=None, selection_result=None, components=<factory>)

Bases: object

Container for control variable effects with optional selection.

Attributes

contributionpt.TensorVariable

Total control contribution (n_obs,).

beta_selectedpt.TensorVariable | None

Coefficients for selected (shrinkage) variables.

beta_fixedpt.TensorVariable | None

Coefficients for fixed (non-shrinkage) variables.

selection_resultVariableSelectionResult | None

Full selection result for diagnostics.

componentsdict[str, pt.TensorVariable]

Individual variable contributions.

contribution: pt.TensorVariable

beta_selected: pt.TensorVariable | None = None

beta_fixed: pt.TensorVariable | None = None

selection_result: VariableSelectionResult | None = None

components: dict[str, pt.TensorVariable]

__init__(contribution, beta_selected=None, beta_fixed=None, selection_result=None, components=<factory>)

mmm_framework.mmm_extensions.components.create_regularized_horseshoe_prior(name, n_variables, n_obs, sigma, config, dims=None)

Create regularized horseshoe prior (Piironen & Vehtari, 2017).

The regularized horseshoe provides: - Strong shrinkage of small effects toward zero - Minimal shrinkage of large effects (they “escape” the horseshoe) - Slab regularization to prevent unrealistically large effects

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the parameters.

n_variablesint

Number of variables (D).

n_obsint

Number of observations (N).

sigmapt.TensorVariable

Observation noise standard deviation.

configHorseshoeConfig

Horseshoe configuration.

dimsstr | None

PyMC dimension name for coefficients.

Returns

VariableSelectionResult

Container with beta and diagnostic quantities.

mmm_framework.mmm_extensions.components.create_finnish_horseshoe_prior(name, n_variables, n_obs, sigma, config, dims=None)

Create Finnish horseshoe prior (Piironen & Vehtari, 2017).

Mathematically identical to the regularized horseshoe.

Return type:

VariableSelectionResult

mmm_framework.mmm_extensions.components.create_spike_slab_prior(name, n_variables, config, dims=None)

Create spike-and-slab prior for variable selection.

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the parameters.

n_variablesint

Number of variables.

configSpikeSlabConfig

Spike-slab configuration.

dimsstr | None

PyMC dimension name.

Returns

VariableSelectionResult

Container with beta and inclusion indicators.

mmm_framework.mmm_extensions.components.create_bayesian_lasso_prior(name, n_variables, config, dims=None)

Create Bayesian LASSO prior (Park & Casella, 2008).

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the parameters.

n_variablesint

Number of variables.

configLassoConfig

LASSO configuration.

dimsstr | None

PyMC dimension name.

Returns

VariableSelectionResult

Container with beta and scale parameters.

mmm_framework.mmm_extensions.components.create_variable_selection_prior(name, n_variables, n_obs, sigma, config, dims=None)

Factory function to create variable selection priors.

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the coefficient parameters.

n_variablesint

Number of control variables subject to selection.

n_obsint

Number of observations.

sigmapt.TensorVariable

Observation noise standard deviation.

configVariableSelectionConfig

Complete configuration specifying method and hyperparameters.

dimsstr | None

PyMC dimension name for the coefficient vector.

Returns

VariableSelectionResult

Container with coefficient vector and diagnostic quantities.

mmm_framework.mmm_extensions.components.build_control_effects_with_selection(X_controls, control_names, n_obs, sigma, selection_config, name_prefix='control')

Build control variable effects with optional variable selection.

Return type:

ControlEffectResult

Parameters

X_controlsarray-like

Control variable matrix (n_obs, n_controls).

control_nameslist[str]

Names of control variables.

n_obsint

Number of observations.

sigmapt.TensorVariable

Observation noise (for horseshoe calibration).

selection_configVariableSelectionConfig

Configuration for variable selection.

name_prefixstr

Prefix for parameter names.

Returns

ControlEffectResult

Container with contributions and coefficients.

mmm_framework.mmm_extensions.components.compute_inclusion_probabilities(trace, config, name='beta_controls', threshold=0.1)

Compute posterior inclusion probabilities from fitted model.

Return type:

dict[str, ndarray]

Parameters

traceaz.InferenceData

Posterior samples from fitted model.

configVariableSelectionConfig

Configuration used for fitting.

namestr

Base name of the coefficient parameters.

thresholdfloat

For horseshoe: signal-to-noise threshold for “inclusion”.

Returns

dict

Dictionary with ‘inclusion_prob’ array and ‘effective_nonzero’.

mmm_framework.mmm_extensions.components.summarize_variable_selection(trace, control_names, config, name='beta_controls')

Create summary table of variable selection results.

Return type:

DataFrame

Parameters

traceaz.InferenceData

Posterior samples.

control_nameslist[str]

Names of control variables.

configVariableSelectionConfig

Configuration used.

namestr

Base parameter name.

Returns

pd.DataFrame

Summary with columns: variable, mean, std, hdi_3%, hdi_97%, inclusion_prob, selected.

Cross Effects

Cross-effect builders for multivariate MMM.

These functions build cross-effect structures (cannibalization, halo effects) between outcomes in multivariate models.

class mmm_framework.mmm_extensions.components.cross_effects.CrossEffectSpec(source_idx, target_idx, effect_type, prior_sigma=0.3)

Bases: object

Specification for a single cross-effect.

source_idx: int

target_idx: int

effect_type: str

prior_sigma: float = 0.3

__init__(source_idx, target_idx, effect_type, prior_sigma=0.3)

mmm_framework.mmm_extensions.components.cross_effects.build_cross_effect_matrix(specs, n_outcomes, name_prefix='psi')

Build cross-effect coefficient matrix.

Return type:

tuple[TensorVariable, dict[tuple[int, int], TensorVariable]]

Parameters

specslist[CrossEffectSpec]

Cross-effect specifications

n_outcomesint

Number of outcomes

name_prefixstr

Prefix for parameter names

Returns

tuple

(cross_effect_matrix, individual_params_dict)

mmm_framework.mmm_extensions.components.cross_effects.compute_cross_effect_contribution(Y, psi_matrix, target_idx, n_outcomes, modulation=None)

Compute cross-effect contribution for a single target outcome.

Return type:

TensorVariable

Parameters

YTensorVariable

Outcome matrix (n_obs, n_outcomes)

psi_matrixTensorVariable

Cross-effect coefficients (n_outcomes, n_outcomes)

target_idxint

Index of target outcome

n_outcomesint

Total number of outcomes

modulationdict | None

Optional modulation by source index (e.g., promotion indicators)

Returns

TensorVariable

Cross-effect contribution (n_obs,)

Variable Selection

Variable selection priors for MMM Extensions.

Provides regularized horseshoe, spike-and-slab, and Bayesian LASSO priors for control variable selection.

class mmm_framework.mmm_extensions.components.variable_selection.VariableSelectionResult(beta, inclusion_indicators=None, local_shrinkage=None, global_shrinkage=None, effective_nonzero=None, kappa=None)

Bases: object

Container for variable selection prior outputs.

Attributes

betapt.TensorVariable

The coefficient vector with shrinkage/selection applied.

inclusion_indicatorspt.TensorVariable | None

For spike-slab: soft inclusion indicators (gamma).

local_shrinkagept.TensorVariable | None

For horseshoe: local shrinkage parameters (lambda).

global_shrinkagept.TensorVariable | None

For horseshoe: global shrinkage parameter (tau).

effective_nonzeropt.TensorVariable | None

Estimated number of effectively nonzero coefficients.

kappapt.TensorVariable | None

Shrinkage factors for each coefficient (horseshoe).

beta: pt.TensorVariable

inclusion_indicators: pt.TensorVariable | None = None

local_shrinkage: pt.TensorVariable | None = None

global_shrinkage: pt.TensorVariable | None = None

effective_nonzero: pt.TensorVariable | None = None

kappa: pt.TensorVariable | None = None

__init__(beta, inclusion_indicators=None, local_shrinkage=None, global_shrinkage=None, effective_nonzero=None, kappa=None)

class mmm_framework.mmm_extensions.components.variable_selection.ControlEffectResult(contribution, beta_selected=None, beta_fixed=None, selection_result=None, components=<factory>)

Bases: object

Container for control variable effects with optional selection.

Attributes

contributionpt.TensorVariable

Total control contribution (n_obs,).

beta_selectedpt.TensorVariable | None

Coefficients for selected (shrinkage) variables.

beta_fixedpt.TensorVariable | None

Coefficients for fixed (non-shrinkage) variables.

selection_resultVariableSelectionResult | None

Full selection result for diagnostics.

componentsdict[str, pt.TensorVariable]

Individual variable contributions.

contribution: pt.TensorVariable

beta_selected: pt.TensorVariable | None = None

beta_fixed: pt.TensorVariable | None = None

selection_result: VariableSelectionResult | None = None

components: dict[str, pt.TensorVariable]

__init__(contribution, beta_selected=None, beta_fixed=None, selection_result=None, components=<factory>)

mmm_framework.mmm_extensions.components.variable_selection.create_regularized_horseshoe_prior(name, n_variables, n_obs, sigma, config, dims=None)

Create regularized horseshoe prior (Piironen & Vehtari, 2017).

The regularized horseshoe provides: - Strong shrinkage of small effects toward zero - Minimal shrinkage of large effects (they “escape” the horseshoe) - Slab regularization to prevent unrealistically large effects

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the parameters.

n_variablesint

Number of variables (D).

n_obsint

Number of observations (N).

sigmapt.TensorVariable

Observation noise standard deviation.

configHorseshoeConfig

Horseshoe configuration.

dimsstr | None

PyMC dimension name for coefficients.

Returns

VariableSelectionResult

Container with beta and diagnostic quantities.

mmm_framework.mmm_extensions.components.variable_selection.create_finnish_horseshoe_prior(name, n_variables, n_obs, sigma, config, dims=None)

Create Finnish horseshoe prior (Piironen & Vehtari, 2017).

Mathematically identical to the regularized horseshoe.

Return type:

VariableSelectionResult

mmm_framework.mmm_extensions.components.variable_selection.create_spike_slab_prior(name, n_variables, config, dims=None)

Create spike-and-slab prior for variable selection.

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the parameters.

n_variablesint

Number of variables.

configSpikeSlabConfig

Spike-slab configuration.

dimsstr | None

PyMC dimension name.

Returns

VariableSelectionResult

Container with beta and inclusion indicators.

mmm_framework.mmm_extensions.components.variable_selection.create_bayesian_lasso_prior(name, n_variables, config, dims=None)

Create Bayesian LASSO prior (Park & Casella, 2008).

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the parameters.

n_variablesint

Number of variables.

configLassoConfig

LASSO configuration.

dimsstr | None

PyMC dimension name.

Returns

VariableSelectionResult

Container with beta and scale parameters.

mmm_framework.mmm_extensions.components.variable_selection.create_variable_selection_prior(name, n_variables, n_obs, sigma, config, dims=None)

Factory function to create variable selection priors.

Return type:

VariableSelectionResult

Parameters

namestr

Base name for the coefficient parameters.

n_variablesint

Number of control variables subject to selection.

n_obsint

Number of observations.

sigmapt.TensorVariable

Observation noise standard deviation.

configVariableSelectionConfig

Complete configuration specifying method and hyperparameters.

dimsstr | None

PyMC dimension name for the coefficient vector.

Returns

VariableSelectionResult

Container with coefficient vector and diagnostic quantities.

mmm_framework.mmm_extensions.components.variable_selection.build_control_effects_with_selection(X_controls, control_names, n_obs, sigma, selection_config, name_prefix='control')

Build control variable effects with optional variable selection.

Return type:

ControlEffectResult

Parameters

X_controlsarray-like

Control variable matrix (n_obs, n_controls).

control_nameslist[str]

Names of control variables.

n_obsint

Number of observations.

sigmapt.TensorVariable

Observation noise (for horseshoe calibration).

selection_configVariableSelectionConfig

Configuration for variable selection.

name_prefixstr

Prefix for parameter names.

Returns

ControlEffectResult

Container with contributions and coefficients.

mmm_framework.mmm_extensions.components.variable_selection.compute_inclusion_probabilities(trace, config, name='beta_controls', threshold=0.1)

Compute posterior inclusion probabilities from fitted model.

Return type:

dict[str, ndarray]

Parameters

traceaz.InferenceData

Posterior samples from fitted model.

configVariableSelectionConfig

Configuration used for fitting.

namestr

Base name of the coefficient parameters.

thresholdfloat

For horseshoe: signal-to-noise threshold for “inclusion”.

Returns

dict

Dictionary with ‘inclusion_prob’ array and ‘effective_nonzero’.

mmm_framework.mmm_extensions.components.variable_selection.summarize_variable_selection(trace, control_names, config, name='beta_controls')

Create summary table of variable selection results.

Return type:

DataFrame

Parameters

traceaz.InferenceData

Posterior samples.

control_nameslist[str]

Names of control variables.

configVariableSelectionConfig

Configuration used.

namestr

Base parameter name.

Returns

pd.DataFrame

Summary with columns: variable, mean, std, hdi_3%, hdi_97%, inclusion_prob, selected.

Transforms

Extension-specific transformation functions.

Most transforms are re-exported from mmm_framework.transforms. This module contains only extensions-specific implementations that don’t exist in the base module.

mmm_framework.mmm_extensions.components.transforms.geometric_adstock_pt(x, alpha, l_max=8, normalize=True)

Apply geometric adstock transformation using PyTensor scan.

This version uses scan for proper gradient flow in complex models. For most use cases, geometric_adstock_convolution is preferred.

Return type:

TensorVariable

Parameters

xTensorVariable

Input media variable (n_obs,)

alphaTensorVariable

Decay rate [0, 1]

l_maxint

Maximum lag length

normalizebool

Whether to normalize weights to sum to 1

Returns

TensorVariable

Adstocked media variable

mmm_framework.mmm_extensions.components.transforms.geometric_adstock_convolution(x, alpha, l_max=8, normalize=True)

Apply geometric adstock using matrix multiplication (no scan).

This is often more efficient and avoids scan complexity. Requires knowing n_obs at graph construction time.

Return type:

TensorVariable

Parameters

xTensorVariable

Input media variable (n_obs,)

alphaTensorVariable

Decay rate [0, 1]

l_maxint

Maximum lag length

normalizebool

Whether to normalize weights to sum to 1

Returns

TensorVariable

Adstocked media variable

mmm_framework.mmm_extensions.components.transforms.geometric_adstock_matrix(X, alphas, l_max=8)

Apply geometric adstock to multiple channels.

Return type:

TensorVariable

Parameters

XTensorVariable

Media matrix (n_obs, n_channels)

alphasTensorVariable

Decay rates per channel (n_channels,)

l_maxint

Maximum lag length

Returns

TensorVariable

Adstocked media matrix (n_obs, n_channels)

mmm_framework.mmm_extensions.components.transforms.adstock_weights_pt(kind, l_max, *, alpha=0.5, theta=0.0, shape=2.0, scale=2.0, normalize=True)

Build a lag-indexed FIR adstock weight vector (PyTensor).

Return type:

pt.TensorVariable

Parameters

kind{“geometric”, “delayed”, “weibull”, “none”}

Kernel shape.

l_maxint

Kernel length (number of lags).

alphaTensorVariable or float

Decay rate for "geometric" / "delayed".

thetaTensorVariable or float

Peak/delay lag for "delayed" (0 reproduces geometric).

shape, scaleTensorVariable or float

Weibull shape k and scale lambda.

normalizebool

If True, weights are scaled to sum to 1.

mmm_framework.mmm_extensions.components.transforms.apply_adstock_pt(x, weights, l_max)

Convolve a series with an FIR adstock kernel (causal, no scan).

Computes y[t] = sum_k weights[k] * x[t - k] with causal zero padding, matching mmm_framework.transforms.adstock.apply_adstock().

Return type:

TensorVariable

mmm_framework.mmm_extensions.components.transforms.parametric_adstock_pt(x, kind, l_max=8, *, alpha=0.5, theta=0.0, shape=2.0, scale=2.0, normalize=True)

Apply geometric/delayed/Weibull FIR adstock to a 1D series (PyTensor).

Dispatches on kind to build the kernel, then convolves. This is the in-graph counterpart used when adstock parameters are estimated.

Return type:

pt.TensorVariable

mmm_framework.mmm_extensions.components.transforms.apply_transformation_pipeline(x, transforms)

Apply a sequence of transformations.

Return type:

TensorVariable

Parameters

xTensorVariable

Input variable

transformslist[tuple[Callable, dict]]

List of (transform_fn, params_dict) tuples

Returns

TensorVariable

Transformed output

mmm_framework.mmm_extensions.components.transforms.logistic_saturation_pt(x, lam)

Apply logistic saturation transformation (PyTensor version).

Return type:

TensorVariable

Parameters

xTensorVariable

Input (already adstocked)

lamTensorVariable

Saturation rate (higher = faster saturation)

Returns

TensorVariable

Saturated output in [0, 1]

mmm_framework.mmm_extensions.components.transforms.hill_saturation(x, kappa, slope)

Apply Hill saturation transformation (PyTensor version).

Return type:

TensorVariable

Parameters

xTensorVariable

Input (already adstocked)

kappaTensorVariable

Half-saturation point (EC50)

slopeTensorVariable

Steepness of curve

Returns

TensorVariable

Saturated output in [0, 1]

Models

Base Models

Base class for extended MMM models.

Provides common functionality for NestedMMM, MultivariateMMM, and CombinedMMM.

class mmm_framework.mmm_extensions.models.base.BaseExtendedMMM(X_media, y, channel_names, index=None, model_config=None, trend_config=None)

Bases: object

Base class for extended MMM models.

Provides common functionality for all extended model types: - Data storage and validation - Model building lifecycle - MCMC fitting - Result extraction

Subclasses must implement: - _build_coords(): Return PyMC coordinate dict - _build_model(): Return built PyMC model

__init__(X_media, y, channel_names, index=None, model_config=None, trend_config=None)

Initialize the base model.

Parameters

X_medianp.ndarray

Media variable matrix (n_obs, n_channels)

ynp.ndarray

Target variable (n_obs,)

channel_nameslist[str]

Names of media channels

indexpd.Index | None

Optional time index for the data

model_configModelConfig | None

Optional core model configuration carrying the seasonality settings and the outcome likelihood family. When None the extension keeps its historical hard-coded baseline (no seasonality, Normal outcome) — so a directly-constructed model is byte-identical to before.

trend_configTrendConfig | None

Optional trend configuration. None or type none → no trend term (historical behavior); linear adds a standardized-scale trend so a real drift does not contaminate the media coefficients.

property model: pymc.Model

Get or build the PyMC model.

save(path)

Persist the model (config, data, experiments) and its fitted trace.

The PyMC graph is dropped (rebuilt on demand) and the trace is written separately as NetCDF; everything else – including any registered ExperimentMeasurement calibrations – is pickled, so a reloaded model can be inspected, predicted from, or re-fit with the same experiment anchoring.

Return type:

None

classmethod load(path)

Load a model saved with save(), reattaching its trace.

Also reads MMMSerializer’s extended-flavor saves (which gzip the trace to trace.nc.gz by default) — the trace loader handles both layouts, so a cross-flavor load never silently drops the posterior.

Return type:

BaseExtendedMMM

add_experiment_calibration(experiments)

Register experiment likelihood terms and invalidate the built graph.

The experiments are folded into the model as likelihood terms on the model-implied estimand (contribution / ROAS / marginal ROAS) the next time the graph is built, so call this before fit(). Returns self for chaining.

Multi-outcome models (MultivariateMMM, CombinedMMM) require each measurement’s outcome to be set.

Return type:

BaseExtendedMMM

fit(draws=1000, tune=1000, chains=4, target_accept=0.9, random_seed=None, nuts_sampler='pymc', method='nuts', **kwargs)

Fit the model with NUTS (default), SMC, or a fast approximate method.

method ∈ {nuts, smc, map, laplace, advi, fullrank_advi, pathfinder}. NUTS is full MCMC for real inference. smc is tempered Sequential Monte Carlo — also exact (ModelResults.approximate stays False): slower than NUTS but robust to the multimodal geometries extension models are prone to (reflected factor modes, label switching, adstock↔AR ridges) and it estimates the log marginal likelihood for model comparison. The approximate methods fit in seconds for quick model checks — the returned ModelResults.approximate is True, R-hat/ESS are None and the uncertainty is not calibrated (re-fit with NUTS before trusting intervals/decisions). The approximate posterior is a drop-in for the NUTS trace (deterministics included), so the extended reports and pathway analysis work off it. Extension models share the base model’s engines (run_approximate_fit(), run_smc_fit()).

Parameters:

• draws (int) – Posterior draws per chain (NUTS), particles per SMC run, or number of approximate draws (ADVI/Laplace/Pathfinder; MAP is a single point and ignores it).

• nuts_sampler (tune / target_accept /) – NUTS controls (unused by SMC and the approximate methods). chains also sets the number of independent SMC runs (R-hat is computed across them).

• random_seed (int | None) – Random seed for reproducibility.

• method (FitMethod | str) – Inference method (see above).

• **kwargs – forwarded to pm.sample (NUTS), pm.sample_smc (SMC), or the approximate fitter.

Return type:

ModelResults

Returns:

Container with trace, model and diagnostics.

property trace: arviz.InferenceData

Get the fitted trace.

predict(X_media=None, X_controls=None, return_original_scale=True, hdi_prob=0.94, random_seed=None)

Posterior-predictive outcome draws, mirroring mmm_framework.model.base.BayesianMMM.predict().

Resamples the outcome likelihood (y_obs) under the posterior — optionally with counterfactual raw-scale X_media / X_controls swapped into the graph — and returns a PredictionResults whose y_pred_samples has shape (n_draws, n_obs). Latent states (mediators, factors) keep their posterior draws, so counterfactuals are structure-preserving.

Single-outcome models (NestedMMM, StructuralNestedMMM) predict their one outcome; the multi-outcome models (MultivariateMMM, CombinedMMM) predict their primary outcome (see _primary_outcome_index()) from the joint Y_obs likelihood, so the single-KPI reporting and goodness-of-fit tooling work on them too.

property X_media_raw: ndarray

Raw (pre-transform) media matrix (n_obs, n_channels).

property y_raw: ndarray

Observed outcome on the caller’s original scale (n_obs,).

Multi-outcome models (MultivariateMMM / CombinedMMM) report the primary outcome (see _primary_outcome_index()); a 2-D self.y is sliced to that column.

property time_idx: ndarray

Period index per observation. Extension models are a single national series (one obs per period), so this is a plain 0..n_obs-1 range.

sample_channel_contributions(X_media=None, max_draws=None, random_seed=None)

Posterior draws of per-channel contributions to the (primary) outcome.

Evaluates the graph’s channel_contributions deterministic — which each extension registers in ORIGINAL KPI units — with X_media (raw scale) optionally swapped in for a counterfactual scenario, returning (n_draws, n_obs, n_channels). Mirrors mmm_framework.model.base.BayesianMMM.sample_channel_contributions() so the interactive report and response-curve tooling work unchanged.

For the mediation / structural models the contribution is the channel’s total (direct + mediated) linear effect on the outcome (linearized for the nonlinear structural paths, consistent with the pathway table); for the multi-outcome models it is the contribution to the primary outcome. Unlike the base model this returns ORIGINAL-scale contributions directly (the deterministic already carries y_std), so callers must NOT rescale.

Return type:

ndarray

summary(var_names=None)

Get posterior summary statistics.

Return type:

DataFrame

sample_prior_predictive(samples=500, random_seed=None)

Sample from the prior predictive distribution.

Return type:

InferenceData

compute_parameter_learning(var_names=None, *, prior_samples=1000, random_seed=None, **kwargs)

Quantify how much the data updated each parameter relative to its prior.

Draws prior samples and compares them to the fitted posterior, returning the prior-to-posterior contraction, overlap, and location shift for every parameter. Particularly useful for sign-constrained effects such as a cannibalization cross-effect (psi = -HalfNormal), where a posterior “entirely below zero” can simply restate the prior: contraction/overlap reveal whether the data actually learned it. See mmm_framework.diagnostics.parameter_learning().

Return type:

DataFrame

Parameters

var_names:

Parameters to diagnose. None (default) uses the model’s free random variables (which include the psi_*_raw cross-effect magnitudes).

prior_samples:

Number of prior draws used to estimate the prior moments/overlap.

random_seed:

Seed for the prior draw (reproducibility).

**kwargs:

Forwarded to parameter_learning().

Returns

pandas.DataFrame

One row per parameter, sorted by contraction ascending (most prior-dominated first).

Nested Models

Nested/Mediated MMM implementation.

Models media effects that flow through intermediate mediators (e.g., awareness, foot traffic) to the final outcome.

class mmm_framework.mmm_extensions.models.nested.NestedMMM(X_media, y, channel_names, config, mediator_data=None, mediator_masks=None, index=None, model_config=None, trend_config=None)

Bases: BaseExtendedMMM

MMM with nested/mediated causal pathways.

Models: Media → Mediator(s) → Outcome With optional direct effects: Media → Outcome

This model is useful when: - You have intermediate metrics (awareness, consideration, traffic) - You want to decompose effects into direct vs mediated - You have partial observations of mediator variables

Parameters

X_medianp.ndarray

Media variable matrix (n_obs, n_channels)

ynp.ndarray

Target outcome variable (n_obs,)

channel_nameslist[str]

Names of media channels

configNestedModelConfig

Configuration specifying mediators and their properties

mediator_datadict[str, np.ndarray] | None

Observed mediator data (where available)

mediator_masksdict[str, np.ndarray] | None

Boolean masks indicating which observations are available

indexpd.Index | None

Optional time index

__init__(X_media, y, channel_names, config, mediator_data=None, mediator_masks=None, index=None, model_config=None, trend_config=None)

Initialize the base model.

Parameters

X_medianp.ndarray

Media variable matrix (n_obs, n_channels)

ynp.ndarray

Target variable (n_obs,)

channel_nameslist[str]

Names of media channels

indexpd.Index | None

Optional time index for the data

model_configModelConfig | None

Optional core model configuration carrying the seasonality settings and the outcome likelihood family. When None the extension keeps its historical hard-coded baseline (no seasonality, Normal outcome) — so a directly-constructed model is byte-identical to before.

trend_configTrendConfig | None

Optional trend configuration. None or type none → no trend term (historical behavior); linear adds a standardized-scale trend so a real drift does not contaminate the media coefficients.

get_mediation_effects()

Extract mediation effect estimates.

Return type:

DataFrame

Multivariate Models

Multivariate MMM implementation.

Models multiple correlated outcomes with cross-effects (cannibalization, halo effects) between products/brands.

class mmm_framework.mmm_extensions.models.multivariate.MultivariateMMM(X_media, outcome_data, channel_names, config, promotion_data=None, index=None, model_config=None, trend_config=None)

Bases: 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_medianp.ndarray

Media variable matrix (n_obs, n_channels)

outcome_datadict[str, np.ndarray]

Dictionary mapping outcome names to their data

channel_nameslist[str]

Names of media channels

configMultivariateModelConfig

Configuration for outcomes and cross-effects

promotion_datadict[str, np.ndarray] | None

Optional promotion indicators for modulated effects

indexpd.Index | None

Optional time index

__init__(X_media, outcome_data, channel_names, config, promotion_data=None, index=None, model_config=None, trend_config=None)

Initialize the base model.

Parameters

X_medianp.ndarray

Media variable matrix (n_obs, n_channels)

ynp.ndarray

Target variable (n_obs,)

channel_nameslist[str]

Names of media channels

indexpd.Index | None

Optional time index for the data

model_configModelConfig | None

Optional core model configuration carrying the seasonality settings and the outcome likelihood family. When None the extension keeps its historical hard-coded baseline (no seasonality, Normal outcome) — so a directly-constructed model is byte-identical to before.

trend_configTrendConfig | None

Optional trend configuration. None or type none → no trend term (historical behavior); linear adds a standardized-scale trend so a real drift does not contaminate the media coefficients.

get_cross_effects_summary()

Extract cross-effect estimates.

Return type:

DataFrame

get_correlation_matrix()

Extract posterior mean correlation matrix.

Return type:

DataFrame

Combined Models

Combined Nested + Multivariate MMM implementation.

Combines mediated pathways with correlated multi-outcome modeling.

class mmm_framework.mmm_extensions.models.combined.CombinedMMM(X_media, outcome_data, channel_names, config, mediator_data=None, mediator_masks=None, promotion_data=None, index=None, model_config=None, trend_config=None)

Bases: BaseExtendedMMM

Combined nested + multivariate model.

Supports both mediated pathways AND correlated outcomes with cross-effects.

This model is useful when: - You have multiple outcomes (products/brands) - Effects flow through intermediate mediators - Outcomes are correlated and may affect each other

Parameters

X_medianp.ndarray

Media variable matrix (n_obs, n_channels)

outcome_datadict[str, np.ndarray]

Dictionary mapping outcome names to their data

channel_nameslist[str]

Names of media channels

configCombinedModelConfig

Configuration for nested + multivariate structure

mediator_datadict[str, np.ndarray] | None

Observed mediator data (where available)

mediator_masksdict[str, np.ndarray] | None

Boolean masks indicating which observations are available

promotion_datadict[str, np.ndarray] | None

Optional promotion indicators for modulated effects

indexpd.Index | None

Optional time index

__init__(X_media, outcome_data, channel_names, config, mediator_data=None, mediator_masks=None, promotion_data=None, index=None, model_config=None, trend_config=None)

Initialize the base model.

Parameters

X_medianp.ndarray

Media variable matrix (n_obs, n_channels)

ynp.ndarray

Target variable (n_obs,)

channel_nameslist[str]

Names of media channels

indexpd.Index | None

Optional time index for the data

model_configModelConfig | None

Optional core model configuration carrying the seasonality settings and the outcome likelihood family. When None the extension keeps its historical hard-coded baseline (no seasonality, Normal outcome) — so a directly-constructed model is byte-identical to before.

trend_configTrendConfig | None

Optional trend configuration. None or type none → no trend term (historical behavior); linear adds a standardized-scale trend so a real drift does not contaminate the media coefficients.

get_cross_effects_summary()

Extract cross-effect estimates.

Same convention as MultivariateMMM.get_cross_effects_summary(): each row reports the configured source -> target effect read from psi_matrix[source_idx, target_idx] (cannibalization rows carry the imposed negative sign).

Return type:

DataFrame

get_effect_decomposition()

Get full effect decomposition by channel and outcome.

Return type:

DataFrame

Results

Result containers for MMM Extensions.

This module contains dataclasses used to return results from extended model fitting, prediction, and analysis.

class mmm_framework.mmm_extensions.results.MediationEffects(channel, direct_effect, direct_effect_sd, indirect_effects, total_indirect, total_effect, proportion_mediated)

Bases: object

Container for mediation analysis results.

channel: str

direct_effect: float

direct_effect_sd: float

indirect_effects: dict[str, float]

total_indirect: float

total_effect: float

proportion_mediated: float

to_dict()

Convert to dictionary representation.

Return type:

dict

__init__(channel, direct_effect, direct_effect_sd, indirect_effects, total_indirect, total_effect, proportion_mediated)

class mmm_framework.mmm_extensions.results.CrossEffectSummary(source, target, effect_type, mean, sd, hdi_low, hdi_high)

Bases: object

Container for cross-effect analysis results.

source: str

target: str

effect_type: str

mean: float

sd: float

hdi_low: float

hdi_high: float

to_dict()

Convert to dictionary representation.

Return type:

dict

__init__(source, target, effect_type, mean, sd, hdi_low, hdi_high)

class mmm_framework.mmm_extensions.results.ModelResults(trace, model, config, diagnostics=<factory>)

Bases: object

Container for fitted extended model results.

Analogous to MMMResults in the main model module.

trace: InferenceData

model: Model

config: Any

diagnostics: dict

property approximate: bool

True when this came from an approximate fit (MAP / ADVI / Pathfinder) rather than NUTS — its uncertainty is NOT calibrated. Mirrors MMMResults.approximate so callers can branch uniformly.

property converged: bool | None

MCMC convergence verdict (R-hat / ESS / divergences).

True/False for NUTS fits; None when not assessable. None is NOT “converged”. Do not act on a fit where this is False.

property convergence_flags: list[str]

subset of {divergences, rhat, ess}.

Type:

Which convergence checks failed

summary(var_names=None)

Get posterior summary statistics.

Return type:

DataFrame

plot_trace(var_names=None, **kwargs)

Plot trace diagnostics.

plot_posterior(var_names=None, **kwargs)

Plot posterior distributions (az.plot_posterior is gone on arviz >=1.x — the compat shim routes to arviz_plots.plot_dist).

__init__(trace, model, config, diagnostics=<factory>)

class mmm_framework.mmm_extensions.results.EffectDecomposition(channel, outcome, direct_mean, direct_sd, indirect_mean, indirect_sd, total_mean, total_sd)

Bases: object

Container for effect decomposition results.

Holds direct, indirect, and total effects for each channel-outcome pair.

channel: str

outcome: str

direct_mean: float

direct_sd: float

indirect_mean: float

indirect_sd: float

total_mean: float

total_sd: float

to_dict()

Convert to dictionary representation.

Return type:

dict

__init__(channel, outcome, direct_mean, direct_sd, indirect_mean, indirect_sd, total_mean, total_sd)