Hi, i’m trying to convert a JAGS model into a PyMC model. The JAGS model has a nested for loop that iterates over Nparticipants and then over Ntrials, to do that in PyMC im using a nested PyTensor scan. The data, so y, d_phy and r, are 2D numeric arrays that have shape (Nparticipants, Ntrials). This is the model:
import pymc as pm
import pytensor.tensor as pt
import numpy as np
from pytensor import scan
def model_1v_LG2D(data):
with pm.Model() as model:
y = pt.as_tensor_variable(data['y'])
d_phy = pt.as_tensor_variable(data['d_phy'])
r = pt.as_tensor_variable(data['r'])
# Hyperpriors
lambda_mu = pm.TruncatedNormal('lambda_mu', mu = 0.1, sigma = 1, lower = 0)
lambda_sigma = pm.Uniform('lambda_sigma', lower = 1e-9, upper = 1)
alpha_mu = pm.Deterministic('alpha_mu', 1e-9 + (1 - 2 * 1e-9) * pm.Beta('alpha_mu_raw', alpha = 1, beta = 1))
alpha_kappa = pm.Uniform('alpha_kappa', lower = 1, upper = 10)
w0_mu = pm.Normal('w0_mu', mu = 0, sigma = 10)
w0_sigma = pm.Exponential('w0_sigma', 1)
w1_a = pm.Exponential('w1_a', 1)
w1_b = pm.Exponential('w1_b', 1)
# Latent group indicators
gp = pm.Categorical('gp', p = pm.Dirichlet('pi', a = np.array([1, 1, 1, 1])), shape = data['Nparticipants'])
# Priors
lambda_1 = pm.TruncatedNormal('lambda_1', mu = lambda_mu, sigma = 1 / lambda_sigma, lower = 0.0052, shape = data['Nparticipants'])
lambda_2 = pm.TruncatedNormal('lambda_2', mu = lambda_mu, sigma = 1/ lambda_sigma, upper = 0.0052, shape = data['Nparticipants'])
lambda_ = pm.math.switch(pm.math.eq(gp, 1), 0, pm.math.switch(pm.math.eq(gp, 2), lambda_2, lambda_1))
w0 = pm.Normal('w0', mu = w0_mu, sigma = w0_sigma, shape = data['Nparticipants'])
w1_1 = pm.Gamma('w1_1', alpha = w1_a, beta = w1_b, shape = data['Nparticipants'])
w1 = pm.math.switch(pm.math.eq(gp, 1), 0, w1_1)
alpha_1 = pm.Beta('alpha_1', alpha = alpha_mu * alpha_kappa, beta = (1 - alpha_mu) * alpha_kappa, shape = data['Nparticipants'])
alpha = pm.math.switch(pm.math.eq(gp, 1), 0, alpha_1)
sigma = pm.Uniform('sigma', lower = [1e-9, 1.5], upper = [1.5, 3], shape = 2)
d_sigma = pm.Gamma('d_sigma', alpha = 2, beta = 1, shape = data['Nparticipants'])
bound = pm.Potential('bound', pm.math.switch(d_sigma < 1e-9, -np.inf, 0))
def main_trial_update(v_prev, d_per_prev, theta_prev, y_ij, d_phy_ij, r_ij, gp_i, lambda_i, w0_i, w1_i, alpha_i, d_sigma_i, sigma):
print("Before operation - v_prev shape:", v_prev.shape, "d_per_prev shape:", d_per_prev.shape)
# Generalization
d_per = pm.Normal('d_per', mu = d_phy_ij, sigma = pm.math.sqrt(1 / d_sigma_i)) # Generate missing perceptual distance between CS+ and S
s = pt.switch((v_prev > 0) & (gp_i > 1), pt.exp(-lambda_i * pt.switch(pt.eq(gp_i, 4), d_per, d_phy_ij)), 1)
g = v_prev * s
# Non-linear transfromation (latent - observed scale)
theta = 1 + (10 - 1) / (1 + pt.exp(-(w0_i + w1_i * g)))
# Likelihood
y_likelihood = pm.Normal('y_likelihood', mu = theta, sigma = pm.math.sqrt(1.0 / sigma[gp_i]), observed = y_ij)
# Learning
# Excitatory associative strength
v_next = pt.switch(pt.neq(gp_i, 1), pt.switch(pt.eq(r_ij, 1), v_prev + alpha_i * (r_ij - v_prev), v_prev), 0)
print("After operation - v_next shape:", v_next.shape)
return [v_next, d_per, theta]
def main_participant_update(i, y, d_phy, r, gp, lambda_, w0, w1, alpha, d_sigma, sigma):
v_init = pt.zeros((data['Ntrials'], ), dtype = 'float64') # The initial excitatory associative strength
d_per_init = pt.zeros((data['Ntrials'], ), dtype='float64')
theta_init = pt.zeros((data['Ntrials'], ), dtype='float64')
[v_sequence, d_per_sequence, theta_sequence], _ = scan(
fn = main_trial_update,
sequences = [y[i], d_phy[i], r[i]],
outputs_info = [v_init, d_per_init, theta_init],
non_sequences = [gp[i], lambda_[i], w0[i], w1[i], alpha[i], d_sigma[i], sigma]
)
return [v_sequence, d_per_sequence, theta_sequence]
# Scan over participants
[v_sequence, d_per_sequence, theta_sequence], _ = scan(
fn = main_participant_update,
sequences = [pt.arange(data['Nparticipants'])],
non_sequences = [y, d_phy, r, gp, lambda_, w0, w1, alpha, d_sigma, sigma]
)
return model
Im running the model through this function which also loads the data:
import pymc as pm
import data_preprocessing as data
import arviz as az
import importlib
import logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger('pymc')
logger.setLevel(logging.ERROR)
def sampling_fun(datafile, model_module_name, model_function_name):
# Dynamically import the model definition module
model_module = importlib.import_module(model_module_name)
# Get the specific model function from the module
model_func = getattr(model_module, model_function_name)
# Use the function to create the model
model = model_func(datafile)
with model:
# Trace
trace = pm.sample(draws = 25000, tune = 75000, chains = 4, cores = 4, return_inferencedata = True)
# Prediction
post_pred = pm.sample_posterior_predictive(trace)
result = az.concat(trace, post_pred, dim = "chain")
return result
Result_Study1_CLG2D = sampling_fun(data.Data1_pymcinput_CLG, "model_definitions", "model_1v_LG2D")
az.to_netcdf(Result_Study1_CLG2D, "Fitting_results/Result_Study1_CLG2D.nc")
The model is not complete but i wanted to test a simplified version of it, the problem i am facing is that i am receiving a MissingInputError about the first input of the computational graph that i think might be the index i of the outer scan which i wanted to use to index the data and the priors. This is the error:
pytensor.graph.utils.MissingInputError: Input 0 (<Scalar(int64, shape=())>) of the graph (indices start from 0), used to compute ScalarFromTensor(<Scalar(int64, shape=())>), was not provided and not given a value. Use the PyTensor flag exception_verbosity='high', for more information on this error.
Backtrace when that variable is created:
File "c:\Users\User\Documents\Fear Generalization\run_pymc.py", line 32, in <module>
Result_Study1_CLG2D = sampling_fun(data.Data1_pymcinput_CLG, "model_definitions copy", "model_1v_LG2D")
File "c:\Users\User\Documents\Fear Generalization\run_pymc.py", line 19, in sampling_fun
model = model_func(datafile)
File "c:\Users\User\Documents\Fear Generalization\model_definitions copy.py", line 76, in model_1v_LG2D
[v_sequence, d_per_sequence, theta_sequence], _ = scan(
My python version is 3.12.3, PyMC 5.15.0, PyTensor 2.20.0, numpy 1.26.4.
I also tried to modify the structure of the model by flattening the data and using a singular scan but i failed to do that and it doesn’t work aswell.
Im very new to PyMC, so any help would be very much appreciated.