CustomModelWrapper#

class pymc_marketing.mmm.budget_optimizer.CustomModelWrapper(base_model, idata, channels, adstock_periods=0, channel_data_var='channel_data', adstock=None, **kwargs)[source]#

Wrapper for custom PyMC models to be used with BudgetOptimizer.

This wrapper lets you plug any fitted PyMC model into the budget optimizer without subclassing the built-in MMM. It clones your base_model for each optimization run and resets the channel-data variable to zeros of the appropriate shape.

Parameters:

base_modelpymc.Model: A PyMC model that contains a shared channel_data variable (or the variable named by channel_data_var). The fitted posterior is passed separately via idata.
idataarviz.InferenceData or xarray.DataTree: Posterior samples from the fitted model.
channelsSequence[str]: Names of the channel dimensions, in the same order as the last axis of the channel-data variable.
adstock_periodsint, optional: Number of extra periods to prepend for adstock warm-up. The model is built with num_periods + adstock_periods date steps, and only the last num_periods are used when computing the response. Defaults to 0 (no warm-up).
channel_data_varstr, optional: Name of the shared channel-data variable inside base_model. Defaults to "channel_data".
adstockAny, optional: Deprecated. Pass adstock_periods instead.

Examples

Build a simple custom model and wrap it for budget optimization:

import pymc as pm
import numpy as np
from pymc_marketing.mmm.budget_optimizer import (
    CustomModelWrapper,
    BudgetOptimizer,
)

channels = ["tv", "search", "social"]
n_obs, n_channels = 52, len(channels)

with pm.Model(coords={"date": range(n_obs), "channel": channels}) as base_model:
    channel_data = pm.Data(
        "channel_data", np.zeros((n_obs, n_channels)), dims=("date", "channel")
    )
    beta = pm.Normal("beta", mu=0, sigma=1, dims="channel")
    mu = (channel_data * beta).sum(axis=-1)
    pm.Normal("y", mu=mu, sigma=1, observed=np.zeros(n_obs))

# After sampling, wrap for optimization:
wrapper = CustomModelWrapper(
    base_model=base_model,
    idata=idata,  # your posterior samples
    channels=channels,
)

optimizer = BudgetOptimizer(
    model=wrapper.optimization_model(num_periods=13),
    idata=wrapper.idata,
    adstock_periods=wrapper.adstock_periods,
    channel_scales=wrapper.channel_scales,
    num_periods=13,
)
optimal, result = optimizer.allocate_budget(total_budget=100_000)

Methods

`CustomModelWrapper.__init__`(base_model, ...)	Create a new model by parsing and validating input data from keyword arguments.
`CustomModelWrapper.construct`([_fields_set])
`CustomModelWrapper.copy`(*[, include, ...])	Returns a copy of the model.
`CustomModelWrapper.dict`(*[, include, ...])
`CustomModelWrapper.from_orm`(obj)
`CustomModelWrapper.json`(*[, include, ...])
`CustomModelWrapper.model_parametrized_name`(params)	Compute the class name for parametrizations of generic classes.
`CustomModelWrapper.optimization_model`(...)	Clone `base_model` and resize the channel-data variable for optimization.
`CustomModelWrapper.parse_file`(path, *[, ...])
`CustomModelWrapper.parse_obj`(obj)
`CustomModelWrapper.parse_raw`(b, *[, ...])
`CustomModelWrapper.schema`([by_alias, ...])
`CustomModelWrapper.schema_json`(*[, ...])
`CustomModelWrapper.update_forward_refs`(**localns)
`CustomModelWrapper.validate`(value)

Attributes

`channel_scales`	Per-channel scale factors used by the budget optimizer.
`model_computed_fields`
`model_config`	Configuration for the model, should be a dictionary conforming to [`ConfigDict`][pydantic.config.ConfigDict].
`model_extra`	Get extra fields set during validation.
`model_fields`
`model_fields_set`	Returns the set of fields that have been explicitly set on this model instance.
`base_model`
`idata`
`channel_columns`
`adstock_periods`
`channel_data_var`