functions { real partial_sum(real[] dummy, int start1, int end, real[,] X, vector mu_prior_x, real sigma_v, matrix A, matrix B, real[,,] U, real Q, matrix C, real[] alpha_rdm_pr, real[] beta_rdm_pr, real[] delta_rdm_pr, real[] tau_rdm_pr, real sigma_r, real[,] RTu_rdm, real[,] RTl_rdm, real[,] Cohu_rdm, real[,] Cohl_rdm, real[] delta_rdm, real[] alpha_rdm, real[] tau_rdm, real[] beta_rdm, int w, int[] idx_rdm_obs, int[] Nu_rdm, int[] Nl_rdm, int[,] choice_itc, real[,] amount_later, real[,] delay_later, real[,] amount_sooner, real[] itc_k_pr, real[] itc_beta_pr, real[] itc_k, real[] itc_beta, int[] idx_itc_obs, int[] Tr_itc, int[,] choice_nc, int[,] deltaM, real[,] TotalS, real[] weber_nc, real[] weber_nc_pr, int[] idx_nc_obs, int[] Tr_nc, int[,] y_smb, real[,,] x_smb, real[,] params_smb, int[] idx_smb_obs, int[] Tr_smb, int[,] choice_cd, int[,] Tar_cd, int[,] Nb_cd, real[,] delta_cd, real[] criterion_cd_pr, real[] sigma_cd_pr, real[] criterion_cd, real[] sigma_cd, int[] idx_cd_obs, int[] Tr_cd, int[,] choice_lt, real[,] hi_p_lt, real[,] low_p_lt, real[,] hi_narr_lt, real[,] hi_wide_lt, real[,] low_narr_lt, real[,] low_wide_lt, real[] risk_lt_pr, real[] beta_lt_pr, real[] risk_lt, real[] beta_lt, int[] idx_lt_obs, int[] Tr_lt, int[,,] pressed_gng, int[,,] cue_gng, real[,,] outcome_gng, int[,] Tr_gng, int[] idx_gng_obs, real[] b_gng_pr, real[] pi_gng_pr, real[] xi_gng_pr, real[] ep_gng_pr, real[] rho_gng_pr, real[] b_gng, real[] pi_gng, real[] xi_gng, real[] ep_gng, real[] rho_gng, vector initV_gng, int Bl) { real lt = 0; for (start in start1:end) { if (w == 1) { lt += normal_lpdf(to_vector(X[start,])|mu_prior_x,sigma_v); // initialize X } else { lt += normal_lpdf(to_vector(X[start,])|(A * to_vector(X[start,]) + B * to_vector(U[start,w-1,])), Q); // update X } // update parameters lt += normal_lpdf(weber_nc_pr[start]|C[1,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(itc_k_pr[start]|C[2,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(itc_beta_pr[start]|C[3,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(alpha_rdm_pr[start]|C[4,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(beta_rdm_pr[start]|C[5,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(delta_rdm_pr[start]|C[6,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(tau_rdm_pr[start]|C[7,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(to_vector(params_smb[start,])|C[8:11,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(criterion_cd_pr[start]|C[12,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(sigma_cd_pr[start]|C[13,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(risk_lt_pr[start]|C[14,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(beta_lt_pr[start]|C[15,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(xi_gng_pr[start]|C[16,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(ep_gng_pr[start]|C[17,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(b_gng_pr[start]|C[18,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(pi_gng_pr[start]|C[19,] * to_vector(X[start]),sigma_r); lt += normal_lpdf(rho_gng_pr[start]|C[20,] * to_vector(X[start]),sigma_r); // TASK 1 if (idx_nc_obs[start] != 0) { for (t in 1:Tr_nc[start]) { lt += bernoulli_lpmf(choice_nc[start,t] | normal_cdf(0, deltaM[start, t], weber_nc[start] * TotalS[start, t])); } } // TASK 2 if (idx_itc_obs[start] != 0) { vector[Tr_itc[start]] ev_later = to_vector(amount_later[start,:Tr_itc[start]]) ./ (1 + itc_k[start] * to_vector(delay_later[start,:Tr_itc[start]])/7); vector[Tr_itc[start]] ev_sooner = to_vector(amount_sooner[start,:Tr_itc[start]]); lt += bernoulli_logit_lpmf(choice_itc[start,:Tr_itc[start]]|itc_beta[start] * (ev_later - ev_sooner)); } // TASK 3 if (idx_rdm_obs[start] != 0) { vector[Nu_rdm[start]] delta_cohu; vector[Nl_rdm[start]] delta_cohl; delta_cohu = delta_rdm[start]*to_vector(Cohu_rdm[start,:Nu_rdm[start]]); lt += wiener_lpdf(RTu_rdm[start, :Nu_rdm[start]] | alpha_rdm[start], tau_rdm[start], beta_rdm[start], delta_cohu); delta_cohl = delta_rdm[start]*to_vector(Cohl_rdm[start,:Nl_rdm[start]]); lt += wiener_lpdf(RTl_rdm[start,:Nl_rdm[start]] | alpha_rdm[start], tau_rdm[start], 1-beta_rdm[start], -delta_cohl); } // TASK 4 if (idx_smb_obs[start] != 0) { matrix[Tr_smb[start], 4] X_smb; X_smb = to_matrix(x_smb[start,:Tr_smb[start],]); lt += bernoulli_lpmf(y_smb[start,:Tr_smb[start]] | Phi(X_smb * to_vector(params_smb[start,]))); } // TASK 5 if (idx_cd_obs[start] != 0) { real f; real h; real p; f = 0.5 * (erfc((criterion_cd[start])/(sqrt(2)*sigma_cd[start])) + erfc((criterion_cd[start])/(sqrt(2)*sigma_cd[start]))); for (t in 1:Tr_cd[start]) { if (delta_cd[start,t] > 0) { h = 0.5 * (erfc((criterion_cd[start]-delta_cd[start,t])/(sqrt(2)*sigma_cd[start])) + erfc((criterion_cd[start]+delta_cd[start,t])/(sqrt(2)*sigma_cd[start]))); p = pow((1-h), Tar_cd[start,t]) * pow((1-f),(Nb_cd[start,t]-Tar_cd[start,t])); } else if (delta_cd[start,t] == 0) { p = pow((1-f),Nb_cd[start,t]); } lt += bernoulli_lpmf(choice_cd[start,t]|1-p); } } // TASK 6 if (idx_lt_obs[start] != 0) { real Utility1; real Utility2; for (t in 1:Tr_lt[start]) { Utility1 = hi_p_lt[start,t] * pow(hi_narr_lt[start,t], risk_lt[start]) + low_p_lt[start,t] * pow(low_narr_lt[start,t], risk_lt[start]); Utility2 = hi_p_lt[start,t] * pow(hi_wide_lt[start,t], risk_lt[start]) + low_p_lt[start,t] * pow(low_wide_lt[start,t], risk_lt[start]); lt += bernoulli_logit_lpmf(choice_lt[start,t]|beta_lt[start]*(Utility1 - Utility2)); } } // TASK 7 if (idx_gng_obs[start] != 0) { vector[4] wv_g; // action weight for go vector[4] wv_ng; // action weight for nogo vector[4] qv_g; // Q value for go vector[4] qv_ng; // Q value for nogo vector[4] sv; // stimulus value vector[4] pGo; // prob of go (press) for (bl in 1:Bl) { wv_g = initV_gng; wv_ng = initV_gng; qv_g = initV_gng; qv_ng = initV_gng; sv = initV_gng; for (t in 1:Tr_gng[start,bl]) { wv_g[cue_gng[start,bl, t]] = qv_g[cue_gng[start,bl, t]] + b_gng[start] + pi_gng[start] * sv[cue_gng[start,bl, t]]; wv_ng[cue_gng[start,bl, t]] = qv_ng[cue_gng[start,bl, t]]; // qv_ng is always equal to wv_ng (regardless of action) pGo[cue_gng[start,bl, t]] = inv_logit(wv_g[cue_gng[start,bl, t]] - wv_ng[cue_gng[start,bl, t]]); { // noise pGo[cue_gng[start,bl, t]] *= (1 - xi_gng[start]); pGo[cue_gng[start,bl, t]] += xi_gng[start]/2; } lt += bernoulli_lpmf(pressed_gng[start,bl, t]|pGo[cue_gng[start,bl, t]]); // after receiving feedback, update sv[t + 1] sv[cue_gng[start,bl, t]] += ep_gng[start] * (rho_gng[start] * outcome_gng[start,bl, t] - sv[cue_gng[start,bl, t]]); // update action values if (pressed_gng[start,bl, t]) { // update go value qv_g[cue_gng[start,bl, t]] += ep_gng[start] * (rho_gng[start] * outcome_gng[start,bl, t] - qv_g[cue_gng[start,bl, t]]); } else { // update no-go value qv_ng[cue_gng[start,bl, t]] += ep_gng[start] * (rho_gng[start] * outcome_gng[start,bl, t] - qv_ng[cue_gng[start,bl, t]]); } } // end of t loop } // end of b loop } } return lt; } } data { int W; int N; int Xdim; int exo_q_num; real U[N,W,exo_q_num]; //data TASK 1 int W_nc_obs[N]; int W_nc_mis[N]; int idx_nc_obs[N,W]; int P_nc; int T_max_nc; int Tr_nc[N, W]; int deltaM[N, W, T_max_nc]; real TotalS[N, W, T_max_nc]; int choice_nc[N, W, T_max_nc]; //data TASK 2 int idx_itc_obs[N,W]; int P_itc; int T_max_itc; int Tr_itc[N, W]; real delay_later[N, W, T_max_itc]; real amount_later[N, W, T_max_itc]; real amount_sooner[N, W, T_max_itc]; int choice_itc[N, W, T_max_itc]; //data TASK 3 int idx_rdm_obs[N,W]; int P_rdm; int Nu_max_rdm; int Nl_max_rdm; int Nu_rdm[N,W]; int Nl_rdm[N,W]; real RTu_rdm[N, W, Nu_max_rdm]; real RTl_rdm[N, W, Nl_max_rdm]; real Cohu_rdm[N, W, Nu_max_rdm]; real Cohl_rdm[N, W, Nl_max_rdm]; matrix[N,W] minRT_rdm; real RTbound_rdm; //data TASK 4 int idx_smb_obs[N,W]; int P_smb; int T_max_smb; int Tr_smb[N, W]; real x_smb[N, W, T_max_smb, P_smb]; int y_smb[N, W, T_max_smb]; //data TASK 5 int idx_cd_obs[N,W]; int P_cd; int T_max_cd; int Tr_cd[N, W]; int Nb_cd[N, W, T_max_cd]; int Tar_cd[N, W, T_max_cd]; real delta_cd[N, W, T_max_cd]; int choice_cd[N, W, T_max_cd]; //data TASK 6 int idx_lt_obs[N,W]; int P_lt; int T_max_lt; int Tr_lt[N,W]; real hi_p_lt[N,W, T_max_lt]; real low_p_lt[N,W, T_max_lt]; real hi_narr_lt[N,W, T_max_lt]; real low_narr_lt[N,W, T_max_lt]; real hi_wide_lt[N,W, T_max_lt]; real low_wide_lt[N,W, T_max_lt]; int choice_lt[N,W,T_max_lt]; //data TASK 7 int idx_gng_obs[N,W]; int P_gng; int Bl; int T_max_gng; int Tr_gng[N,W,Bl]; int cue_gng[N,W,Bl,T_max_gng]; int pressed_gng[N,W,Bl, T_max_gng]; real outcome_gng[N, W, Bl, T_max_gng]; } transformed data { real dummy[N]; real Q = 1; int cauchy_alpha = 5; // cauchy scale parameter hyperprior int num_par = P_nc + P_itc + P_rdm + P_smb + P_cd + P_lt + P_gng; int grainsize = 1; vector[4] initV_gng = rep_vector(0.0, 4); // initial values for TASK 7 } parameters { // parameters for dynamic process real sigma_x; real sigma_v; real sigma_r; real sigma_a; real sigma_b; real sigma_c; vector[Xdim] mu_prior_x; real X[N,W,Xdim]; matrix[Xdim, Xdim] A; matrix[Xdim, exo_q_num] B; matrix[num_par, Xdim] C; // task parameters real weber_nc_pr[N,W]; // parameters TASK 1 real itc_k_pr[N,W]; // parameters TASK 2 real itc_beta_pr[N,W]; // parameters TASK 2 real alpha_rdm_pr[N,W]; // parameters TASK 3 real beta_rdm_pr[N,W]; // parameters TASK 3 real delta_rdm_pr[N,W]; // parameters TASK 3 real tau_rdm_pr[N,W]; // parameters TASK 3 real params_smb[N, W, P_smb]; // parameters TASK 4 real criterion_cd_pr[N,W]; // parameters TASK 5 real sigma_cd_pr[N,W]; // parameters TASK 5 real risk_lt_pr[N,W]; // parameters TASK 6 real beta_lt_pr[N,W]; // parameters TASK 6 real xi_gng_pr[N,W]; // parameters TASK 7 real ep_gng_pr[N,W]; // parameters TASK 7 real b_gng_pr[N,W]; // parameters TASK 7 real pi_gng_pr[N,W]; // parameters TASK 7 real rho_gng_pr[N,W]; // parameters TASK 7 } transformed parameters { real weber_nc[N,W]; real itc_k[N,W]; real itc_beta[N,W]; real alpha_rdm[N,W]; real beta_rdm[N,W]; real delta_rdm[N,W]; real tau_rdm[N,W]; real criterion_cd[N,W]; real sigma_cd[N,W]; real risk_lt[N,W]; real beta_lt[N,W]; real xi_gng[N,W]; real ep_gng[N,W]; real b_gng[N,W]; real pi_gng[N,W]; real rho_gng[N,W]; for (n in 1:N) { weber_nc[n,] = exp(weber_nc_pr[n,]); itc_k[n,] = exp(itc_k_pr[n,]); itc_beta[n,] = exp(itc_beta_pr[n,]); alpha_rdm[n,] = exp(alpha_rdm_pr[n,]); delta_rdm[n,] = exp(delta_rdm_pr[n,]); criterion_cd[n,] = exp(criterion_cd_pr[n,]); sigma_cd[n,] = exp(sigma_cd_pr[n,]); risk_lt[n,] = exp(risk_lt_pr[n,]); beta_lt[n,] = exp(beta_lt_pr[n,]); rho_gng[n,] = exp(rho_gng_pr[n,]); } for (n in 1:N) { for (w in 1:W) { beta_rdm[n,w] = Phi_approx(beta_rdm_pr[n,w]); tau_rdm[n,w] = Phi_approx(tau_rdm_pr[n,w]) * (minRT_rdm[n,w] - RTbound_rdm) + RTbound_rdm; xi_gng[n,w] = Phi_approx(xi_gng_pr[n,w]); ep_gng[n,w] = Phi_approx(ep_gng_pr[n,w]); } } b_gng = b_gng_pr; pi_gng = pi_gng_pr; } model { // initial priors on dynamic parameters sigma sigma_x ~ cauchy(0, cauchy_alpha); // prior on R diagonal sigma_v ~ cauchy(0, cauchy_alpha); // prior on R diagonal sigma_r ~ cauchy(0, cauchy_alpha); // prior on R diagonal sigma_a ~ cauchy(0, cauchy_alpha); // prior on A variance sigma_b ~ cauchy(0, cauchy_alpha); // prior on B variance sigma_c ~ cauchy(0, cauchy_alpha); // prior on C variance mu_prior_x ~ normal(0,sigma_x); // prior on X mean // put priors an A, B, C to_vector(A) ~ normal(0,sigma_a); to_vector(B) ~ normal(0,sigma_b); to_vector(C) ~ normal(0,sigma_c); for (w in 1:W) { target += reduce_sum(partial_sum, dummy, grainsize, X[,w,], mu_prior_x, sigma_v, A, B, U, Q, C, alpha_rdm_pr[,w], beta_rdm_pr[,w], delta_rdm_pr[,w], tau_rdm_pr[,w], sigma_r, RTu_rdm[,w,], RTl_rdm[,w,], Cohu_rdm[,w,], Cohl_rdm[,w,], delta_rdm[,w], alpha_rdm[,w], tau_rdm[,w],beta_rdm[,w], w, idx_rdm_obs[,w], Nu_rdm[,w], Nl_rdm[,w], choice_itc[,w,], amount_later[,w,], delay_later[,w,], amount_sooner[,w,], itc_k_pr[,w], itc_beta_pr[,w], itc_k[,w], itc_beta[,w], idx_itc_obs[,w], Tr_itc[,w], choice_nc[,w,], deltaM[,w,], TotalS[,w,], weber_nc[,w], weber_nc_pr[,w], idx_nc_obs[,w], Tr_nc[,w], y_smb[,w,], x_smb[,w,,], params_smb[,w,], idx_smb_obs[,w], Tr_smb[,w], choice_cd[,w,], Tar_cd[,w,], Nb_cd[,w,], delta_cd[,w,], criterion_cd_pr[,w], sigma_cd_pr[,w], criterion_cd[,w], sigma_cd[,w], idx_cd_obs[,w], Tr_cd[,w], choice_lt[,w,], hi_p_lt[,w,], low_p_lt[,w,], hi_narr_lt[,w,], hi_wide_lt[,w,], low_narr_lt[,w,], low_wide_lt[,w,], risk_lt_pr[,w], beta_lt_pr[,w], risk_lt[,w], beta_lt[,w], idx_lt_obs[,w], Tr_lt[,w], pressed_gng[,w,,], cue_gng[,w,,], outcome_gng[,w,,], Tr_gng[,w,], idx_gng_obs[,w], b_gng_pr[,w], pi_gng_pr[,w], xi_gng_pr[,w], ep_gng_pr[,w], rho_gng_pr[,w], b_gng[,w], pi_gng[,w], xi_gng[,w], ep_gng[,w], rho_gng[,w], initV_gng, Bl); } }