create_variable_indexer#

pymc_marketing.mmm.lift_test.create_variable_indexer(model, indices)[source]#

Create a function to index variables in the model.

Parameters:

modelpm.Model: PyMC model
indicesdict[str, np.ndarray]: Dictionary of indices for the indices in the model.

Returns:

Callable[[str], TensorVariable]: Function to index variables in the model.

Raises:

KeyError: If the variable is not in the model.

Examples

Create a variable indexer:

import numpy as np
import pymc as pm

from pymc_marketing.mmm.lift_test import create_variable_indexer

coords = {"channel": [0, 1, 2], "geo": ["A", "B", "C"]}
with pm.Model(coords=coords) as model:
    pm.Normal("alpha", dims=("channel", "geo"))
    pm.Normal("beta", dims="channel")

# Usually from exact_row_indices
indices = {"channel": [0, 1], "geo": [1, 0]}

variable_indexer = create_variable_indexer(model, indices)

Get alpha at indices:

alpha_at_indices = variable_indexer("alpha")

Get beta at indices:

beta_at_indices = variable_indexer("beta")