How can I parameterize lags parameters for each independent variable in my model in PyMC (5.10.0)?

Question

I am new to PyMC and I am using the 5.10.0 version. I am running a simple Media Mix Model and I need to parameterize the individual lags between the independent(media activity) and the dependent(sales volume) variables. Parameterizing these cause the following error

Cell In[15], line 75
     72 # Creating a tensor to store transformed values
     73 adstocked = tt.zeros_like(x)
---> 75 for i in range(max_lag, x.shape[0]):
     76     weights = tt.power(rate, tt.arange(max_lag + 1))
     77     adstocked = tt.set_subtensor(adstocked[i], tt.dot(x[i-max_lag:i+1][::-1], weights))

TypeError: 'TensorVariable' object cannot be interpreted as an integer

A reproducible example is as follows.

## Create a simple MMM data 
import pandas as pd
from random import randint
import numpy as np
import pytensor.tensor as tt
import pytensor as pt
import pymc as pm
import pymc.sampling.jax as pmjax
import arviz as az


# Generate date range
dates = pd.date_range(start="2021-01-01", end="2022-01-01")

data = {
    "date": dates,
    "gcm_direct_Impressions": [randint(10000, 20000) for _ in dates],
    "display_direct_Impressions" :[randint(100000,150000) for _ in dates],
    "tv_grps": [randint(30, 50) for _ in dates],
    "tiktok_direct_Impressions": [randint(10000, 15000) for _ in dates],
    "sell_out_quantity": [randint(150, 250) for _ in dates]
}
df = pd.DataFrame(data)
m = max(df['sell_out_quantity'].values)

print(f"Max sales Volume {m}")

channel_columns = [col for col in df.columns if 'Impressions' in col or 'grps' in col]

transform_variables = channel_columns


delay_channels = channel_columns

media_channels = channel_columns

target = 'sell_out_quantity'

### Transform each channel variable

data_transformed = df.copy()

numerical_encoder_dict = {}


for feature in transform_variables:
    # Extracting the original values of the feature.
    original = df[feature].values

    # Calculating the maximum value of the feature.
    max_value = original.max()

    # Dividing each value in the feature by the maximum value.
    transformed = original / max_value

    # Storing the transformed data back into the 'data_transformed' DataFrame.
    data_transformed[feature] = transformed

    # Storing the maximum value used for scaling in the dictionary.
    # This will be used for reversing the transformation if needed.
    numerical_encoder_dict[feature] = max_value



def adstock_transform(x, rate,max_lag):
    """ Apply adstock transformation with PyTensor.
    :param x: PyTensor tensor, original data for the channel
    :param rate: PyTensor tensor, decay rate of the adstock transformation
    :param max_lag: int, maximum lag to consider for the adstock effect
    :return: PyTensor tensor, transformed data
    """
    # Creating a tensor to store transformed values
    adstocked = tt.zeros_like(x)
    
    for i in range(max_lag, x.shape[0]):
        weights = tt.power(rate, tt.arange(max_lag + 1))
        adstocked = tt.set_subtensor(adstocked[i], tt.dot(x[i-max_lag:i+1][::-1], weights))
    
    return adstocked

### Create a model
response_mean = []

with pm.Model() as model_2:
    # Looping through each channel in the list of delay channels.
    for channel_name in delay_channels:
        print(f"Delay Channels: Adding {channel_name}")

        # Extracting the transformed data for the current channel.
        x = data_transformed[channel_name].values

        # Defining Bayesian priors for the adstock, gamma, and alpha parameters for the current channel.
        adstock_param = pm.Beta(f"{channel_name}_adstock", 2, 2)
        saturation_gamma = pm.Beta(f"{channel_name}_gamma", 2, 2)
        saturation_alpha = pm.Gamma(f"{channel_name}_alpha", 3, 1)
        rate = pm.Beta(f'{channel_name}_rate', alpha=1, beta=1)
        lag = pm.DiscreteUniform(f"{channel_name}_lag",lower=0,upper=17)
        
        transformed_X1 = adstock_transform(x,rate,max_lag=lag)
        transformed_X2 = tt.zeros_like(x)
        for i in range(1,len(x)):
            transformed_X2 = tt.set_subtensor(transformed_X2[i],(transformed_X1[i]**saturation_alpha)/(transformed_X1[i]**saturation_alpha+saturation_gamma**saturation_alpha))
        channel_b = pm.HalfNormal(f"{channel_name}_media_coef", sigma = m)
        response_mean.append(transformed_X2 * channel_b)

    intercept = pm.Normal("intercept",mu = np.mean(data_transformed[target].values), sigma = 3)
    sigma = pm.HalfNormal("sigma", 4)
    likelihood = pm.Normal("outcome", mu = intercept + sum(response_mean), sigma = sigma,
                           observed = data_transformed[target].values)

with model_2:
    trace = pmjax.sample_numpyro_nuts(1000, tune=1000, target_accept=0.95)
    
    trace_summary = az.summary(trace)

I tried the following

• Replacing max_lag with int(max_lag.eval()) in the adstock_transform function.
• Replacing x = data_transformed[channel_name].values with x = pm.ConstantData("data_{channel_name",data_transformed[channel_name].values)

None of these worked. Should I change the type of sampler that is used? Are there different implementations of the Discrete Uniform Distribution that help with these types of problems? Your help is greatly appreciated!

Answer 1

I found a trivial workaround.

When using pmjax.sample_numpyro_nuts() for probabilistic modeling, it's important to note that this function is designed for continuous distributions and doesn't support discrete ones directly. One workaround for this limitation could be to substitute the discreteUniform distribution with a Uniform distribution. However, this substitution results in a continuous output, which may not be suitable for functions like adstock_transform that expect a discrete input. To address this, you can convert the continuous output back to a discrete value by using max_lag = int(max_lag.eval()). This line effectively casts the floating-point number into an integer, making it compatible with the adstock_transform function. This answer is open to feedback.