Non-linear drift rates

[Garrid0]

Hello everyone,

I am trying to use HSSM for the first time. I'm interested in modelling the case in which the drift rate is a function of two covariates (say x and y). I've seen in the tutorial that it is possible to do a simple linear regression like this one:

# Make some hierarchical data

n_subjects = 15  # number of subjects
n_trials = 200  # number of trials per subject

sd_v = 0.5  # sd for v-intercept
mean_v = 0.5  # mean for v-intercept

data_list = []
for i in range(n_subjects):
    # Make parameters for subject i
    intercept = np.random.normal(mean_v, sd_v, size=1)
    x = np.random.uniform(-1, 1, size=n_trials)
    y = np.random.uniform(-1, 1, size=n_trials)
    v = intercept + (0.8 * x) + (0.3 * y)

    true_values = np.column_stack(
        [v, np.repeat([[1.5, 0.5, 0.5, 0.0]], axis=0, repeats=n_trials)]
    )

    # Simulate data
    obs_ddm_reg_v = simulator(true_values, model="ddm", n_samples=1)

    # Append simulated data to list
    data_list.append(
        pd.DataFrame(
            {
                "rt": obs_ddm_reg_v["rts"].flatten(),
                "response": obs_ddm_reg_v["choices"].flatten(),
                "x": x,
                "y": y,
                "subject": i,
            }
        )
    )

# Make single dataframe out of subject-wise datasets
dataset_reg_v_hier = pd.concat(data_list)
dataset_reg_v_hier

However, what if we want to model a drift rate that depends non-linearly on these covariates? For example:

v = v_0 + mu_0 * ( x ** alpha + y ** beta ),

Where alpha, beta, mu_0 and v_0 are free parameters

Is it possible to do this with HSSM?

Thank you in advance! :)