parameter recovery issue when nonmonotonic relationship exists between HSSM parameters and covariate

[hyang336]

Describe the bug
First of all, I am not sure if this is a bug or am intrinsic limitation of generalized linear model, PyMC, Bambi, or HSSM. Here is what I am trying to do:

I want to model a scenario where several competing (racing) accumulators whose drift rates depend on a 1-d signal (e.g. confidence), such that the accumulators' drift rates reach maximum at different values of the 1-d signal. For example, one accumulator may represent a confidence rating of 1 on a 5-point scale, and thus should have its max drift rate when a 1-d confidence signal is at its lowest value. Whereas another accumulator representing 2 on the 5-point scale will attain its max drift rate when the confidence signal is at its intermediate value.

To achieve these, I need a generalized linear model where the independent variables and dependent variables can have monotonic or nonmonotonic relationship. Inspired by the beta distribution, I came up with the following model using one of the race_4 models in HSSM. For each accumulator, its drift rate v is a function of the sum of the log transform of the covariate and the log transform of one minus the covariate (the covariate takes the value between 0 and 1), with an exponential link function. In other words, v = exp(intercept + alog(cov) + blog(1-cov)). This is basically taking log of the beta density and rearrange some terms.

Having selected a and b and the intercept, I generated some v's and then some data from those v's. But when I attempted parameter recovery, even though the chain appeared converged, the recovered parameters are way off. I suspect that this may have something to do with the fact that the two independent variables are not really independent. If that is the case, is there an alternative approach to achieve what I intended?

See attached code and figure for more details.

HSSM version
0.1.5

To Reproduce

#generate data using race_4 model, there are separate v and z for each accumulator, but a and t are shared
from ssms.basic_simulators import simulator
import numpy as np
import pandas as pd
from scipy.special import softmax
from scipy.special import beta
import hssm
import bambi as bmb
import arviz as az
from matplotlib import pyplot as plt
import pymc as pm

outdir='/scratch/hyang336/working_dir/HSSM_dev/race_4_LAN/'
#--------------------------------------We can try several generative model--------------------------------###
#fake trialwise neural data, the 4 accumulators are simulated to have monotonic or nonmonotonic relationships with 
#(log-transformed) neural data. This is controlled by take the beta distribution and log transform it, making it a 
#simple linear regression on the log-transformed neural data. The intercept becomes the log transform of 1 over beta 
#function, evaluated at the parameters for the beta distribution (a and b). This intercept was the normalizing factor
#before the log transformation to make sure the function was a distibution (i.e. integrate to 1). Note, larger values
#of a or b will result in large positive or negative value of v down the line, may need to find a way to rescale it...
a0=0.4
b0=1
intercept0=np.log(1/beta(a0,b0))
a1=1.5
b1=3
intercept1=np.log(1/beta(a1,b1))
a2=3
b2=1.5
intercept2=np.log(1/beta(a2,b2))
a3=1
b3=0.4
intercept3=np.log(1/beta(a3,b3))
simneural = np.random.uniform(0, 1, size=2000)

#simulate linear relationship between v and log-transformed neural data, following HSSM tutorial (i.e. no added noise at this step)
#now we need a log link function since we model log(v)=a+b*log(x)+c*log(1-x)
v0=np.exp(intercept0 + (a0-1)*np.log(simneural) + (b0-1)*np.log(1-simneural))
v1=np.exp(intercept1 + (a1-1)*np.log(simneural) + (b1-1)*np.log(1-simneural))
v2=np.exp(intercept2 + (a2-1)*np.log(simneural) + (b2-1)*np.log(1-simneural))
v3=np.exp(intercept3 + (a3-1)*np.log(simneural) + (b3-1)*np.log(1-simneural))

###IMPORTANT: for interpretable param rec test, make sure generate params within training bounds of LAN###

v0_inb= np.where(np.logical_and(v0>= 0,v0<= 2.5))
v1_inb= np.where(np.logical_and(v1>= 0,v1<= 2.5))
v01_inb=np.intersect1d(v0_inb[0],v1_inb[0])

v2_inb= np.where(np.logical_and(v2>= 0,v2<= 2.5))
v3_inb= np.where(np.logical_and(v3>= 0,v3<= 2.5))
v23_inb=np.intersect1d(v2_inb[0],v3_inb[0])

v0123_inb=np.intersect1d(v01_inb,v23_inb)#indices of elements that are in bound for all 4 arrays

#only keep inbound elements
simneural=simneural[v0123_inb]
v0=v0[v0123_inb]
v1=v1[v0123_inb]
v2=v2[v0123_inb]
v3=v3[v0123_inb]


# def replace_outbound(cov_ref,arr_in,low_bound,high_bound,intercept,beta1,beta2):
#     arr_ob=np.where(np.logical_or(arr_in<low_bound,arr_in>high_bound))
#     for idx in arr_ob:
#         outbound=True
#         while outbound:
#             new_cov=np.random.uniform(0, 1, size=1)
#             new_elem=np.exp(intercept+beta1*np.log(new_cov)+beta2*np.log(1-new_cov))
#             if new_elem>=low_bound and new_elem<=high_bound:
#                 cov_ref[idx]=new_cov
#                 arr_in[idx]=new_elem
#                 outbound=False
#     return arr_in,cov_ref


###########################################################################################################

#generate trial-wise parameters with fixed a, z, and t, and bnoundary_param, assumed to take the form theta in radian
true_values = np.column_stack(
    [v0,v1,v2,v3, np.repeat([[1.5, 0.0, 0.5,0.1]], axis=0, repeats=len(simneural))]
)


# Get mode simulations
race4_v = simulator.simulator(true_values, model="race_no_bias_angle_4", n_samples=1)

dataset_race4_v = pd.DataFrame(
    {
        "rt": race4_v["rts"].flatten(),
        "response": race4_v["choices"].flatten(),
        "x": np.log(simneural),
        "y": np.log(1-simneural)
    }
)


#estimate parameters based on data
model_race4_v = hssm.HSSM(
    data=dataset_race4_v,
    model='race_no_bias_angle_4',
    a=1.5,
    include=[
        {
            "name": "v0",
            "prior":{"name": "Uniform", "lower": -12, "upper": 12},
            "formula": "v0 ~ 1 + x + y",
            "link": "log",
        },
        {
            "name": "v1",
            "prior":{"name": "Uniform", "lower": -12, "upper": 12},
            "formula": "v1 ~ 1 + x + y",
            "link": "log",
        },
        {
            "name": "v2",
            "prior":{"name": "Uniform", "lower": -12, "upper": 12},
            "formula": "v2 ~ 1 + x + y",
            "link": "log",
        },
        {
            "name": "v3",
            "prior":{"name": "Uniform", "lower": -12, "upper": 12},
            "formula": "v3 ~ 1 + x + y",
            "link": "log",
        }
    ],
)

#model graph


#sample from the model, 2500-2500 is not enough for the chain the converge
infer_data_race4_v = model_race4_v.sample(
    step=pm.Slice(model=model_race4_v.pymc_model), sampler="mcmc", chains=1, cores=None, draws=5000, tune=10000
)

#save model
az.to_netcdf(infer_data_race4_v,outdir+'sample_5000_10000_trace_ParamInbound_Fixed_a_SliceSampler.nc4')

#load model
#infer_data_race4_v=az.from_netcdf('/home/hyang336/HSSM_race5_dev/HY_dev/race_4_LAN/sample50_trace.nc4')

#diagnostic plots
az.plot_trace(
    infer_data_race4_v,
    var_names="~log_likelihood",  # we exclude the log_likelihood traces here
)
plt.savefig(outdir+'posterior_diagnostic_5000_10000_trace_ParamInbound_Fixed_a_SliceSampler.png')

#fit summary
res_sum=az.summary(model_race4_v.traces)
res_sum.to_csv(outdir+'summary_5000_10000_trace_ParamInbound_Fixed_a_SliceSampler.csv')
#res_slope=res_sum[res_sum.iloc[:,0].str.contains("_x|_y")]
#res_sum.loc[['v0_x','v0_y','v1_x','v1_y','v2_x','v2_y','v3_x','v3_y']]


#parameter recovery is pretty bad...

Screenshots

Additional context
This model was run on a Linux cluster with its own Jax (for bug-fix reasons).

[hyang336]

I reran this with more draws and 2 chains, got similar results:

[AlexanderFengler]

Hi @hyang336,

as an initial reaction to this.

Could you be a bit more specific about the parameters that you used to generate your simulated data here? (I assume the fit it to simulated?)

One thing that strikes me as off, is the fact that the z parameter is hugging 0.
This is likely the case because the parameters of the race models are really hard to identify due to some tradeoffs. Specifically z and a parameters trade off quite severely which in turn affects the estimates of v as well. I suggest to fix the z parameters here (because it doesn't really fulfill the role of a bias in the traditional 2-choice sense anyways).

Which z parameter did you use for you simulations?

Also, but this is separate from the validity of my statements above:

I don't think I fully understood how the racing particles utilize the confidence signals and how that maps onto the choice of GLM in your case.

Best,
Alex

[AlexanderFengler]

Sidenote, this post might be better hosted under discussions.

[hyang336]

Hi, Alex,

Thanks for your reply.

I first chose some intercept and slope values for the GLM, I also randomly generated 2000 trials of neural (confidence) data.
The slope values were chosen more or less arbitrarily as long as they give me the desired (non)monotonicity:

a0=0.4
b0=1
intercept0=np.log(1/beta(a0,b0))
a1=1.5
b1=3
intercept1=np.log(1/beta(a1,b1))
a2=3
b2=1.5
intercept2=np.log(1/beta(a2,b2))
a3=1
b3=0.4
intercept3=np.log(1/beta(a3,b3))
simneural = np.random.uniform(0, 1, size=2000)

Then I used these to generate drift rates for the 4 particles, 2000 trials each, according to the GLM I mentioned:

v0=np.exp(intercept0 + (a0-1)*np.log(simneural) + (b0-1)*np.log(1-simneural))
v1=np.exp(intercept1 + (a1-1)*np.log(simneural) + (b1-1)*np.log(1-simneural))
v2=np.exp(intercept2 + (a2-1)*np.log(simneural) + (b2-1)*np.log(1-simneural))
v3=np.exp(intercept3 + (a3-1)*np.log(simneural) + (b3-1)*np.log(1-simneural))

I also removed trials where any of the v's fall outside of [0 , 2.5].

Then I used these v's and some other prespecified parameters to generate simulated data:

true_values = np.column_stack(
    [v0,v1,v2,v3, np.repeat([[1.5, 0.0, 0.5,0.1]], axis=0, repeats=len(simneural))]
)

race4_v = simulator.simulator(true_values, model="race_no_bias_angle_4", n_samples=1)

As you can see the ground truth z was set at 0 (second value in the np.repeat() above), so in a sense that parameter was recovered correctly. I didn't fix z in the model estimating step, but I did fix a to be the ground truth value. Do you think fixing z instead will lead to better performance?

I want particle one to reflect evidence of low confidence, particle two to reflect evidence of slightly higher confidence, and so on. Assuming we only have a 1-d array representing confidence covariate across trials, then particle one's drift rate can be modeled as having a monotonically decreasing relationship with the covariate (i.e. negative slope), while for particle two, it needs to be nonmonotonic because particle two accumulate evidence the fastest when the covariate is not at its lowest, but at a slightly higher level. And I want 1 type of GLM (i.e. one set of parameters) to be able to accommodate these complex relationship. Does this make sense?

PS. Is there a way to move this whole thread to Discussion?