merge_inference_data#

pymc_marketing.mmm.budget_optimizer.merge_inference_data(idatas, prefixes=None, *, merge_on='channel_data', use_every_n_draw=1)[source]#

Merge multiple xarray.DataTree objects with per-model prefixes.

This is the companion to merge_models() for the inference-data side of a multi-model budget optimization. After calling both functions you have a merged pm.Model and a merged DataTree that can be passed directly to BudgetOptimizer.

Parameters:

idataslist[xarray.DataTree]: Posterior samples from each fitted model. All objects must have a posterior group.
prefixeslist[str] or None, optional: Per-model prefix applied to every variable and dimension name that is not in merge_on. If None (default), prefixes are auto-generated as ["model1", "model2", ...].
merge_onstr or None, optional: Variable name (and its associated dimensions) that is shared across all models and therefore not prefixed. Typically "channel_data" so the shared budget variable remains unprefixed. Pass None to prefix every variable.
use_every_n_drawint, optional: Thinning factor, keeps every n-th posterior draw before merging. Useful when merging many models to keep memory usage manageable. Defaults to 1 (no thinning).

Returns:

xarray.DataTree: A single merged DataTree with prefixed variables and dimensions ready for use as the idata argument to BudgetOptimizer.

Raises:

ValueError: If prefixes is provided but its length does not match len(idatas).

Examples

Merge two fitted MMMs for a multi-region optimization:

from pymc_marketing.pytensor_utils import merge_models
from pymc_marketing.mmm.budget_optimizer import (
    merge_inference_data,
    BudgetOptimizer,
)

# Step 1 – build per-model optimization models
m1 = mmm_north.create_optimization_model("2025-01-01", "2025-03-31")
m2 = mmm_south.create_optimization_model("2025-01-01", "2025-03-31")

# Step 2 – merge PyMC models (from pytensor_utils)
merged_model = merge_models(
    [m1, m2], prefixes=["north", "south"], merge_on="channel_data"
)

# Step 3 – merge inference data
merged_idata = merge_inference_data(
    idatas=[mmm_north.idata, mmm_south.idata],
    prefixes=["north", "south"],
    merge_on="channel_data",
    use_every_n_draw=2,
)

# Step 4 – optimize
optimizer = BudgetOptimizer(
    model=merged_model,
    idata=merged_idata,
    num_periods=13,
    # The model's date axis is carry-in + decisions + carry-over, each
    # flank effective_carryover_lags() wide.
    carry_in_periods=mmm_north.effective_carryover_lags(),
    adstock_periods=mmm_north.effective_carryover_lags(),
    response_variable="north_total_media_contribution_original_scale",
)
optimal, result = optimizer.allocate_budget(total_budget=100_000)