From efeaa104b4e7051208edce7e0230d17b6bf6149f Mon Sep 17 00:00:00 2001 From: Jeong Jinwoo Date: Wed, 17 Dec 2025 19:48:00 +0900 Subject: [PATCH 1/6] feat(pst_Q): add plot_pst_Q --- R/inst/plotting/plot_functions.R | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/R/inst/plotting/plot_functions.R b/R/inst/plotting/plot_functions.R index dbcf2b3e..20b7f8eb 100644 --- a/R/inst/plotting/plot_functions.R +++ b/R/inst/plotting/plot_functions.R @@ -405,6 +405,14 @@ plot_cra_exp <- function(obj, fontSize = 10, ncols = 3, binSize = 30) { return(h_all) } +plot_pst_Q <- function(obj, fontSize = 10, ncols = 2, binSize = 30) { + pars = obj$parVals + h1 = plotDist(sample = pars$mu_alpha, fontSize = fontSize, binSize = binSize, xLim = c(0,1), xLab = expression(paste(alpha, " (Learning Rate)"))) + h2 = plotDist(sample = pars$mu_beta, fontSize = fontSize, binSize = binSize, xLim = c(0,10), xLab = expression(paste(beta, " (Inverse Temp.)"))) + h_all = multiplot(h1, h2, cols = ncols) + return(h_all) +} + plot_pst_gainloss_Q <- function(obj, fontSize = 10, ncols = 3, binSize = 30) { pars = obj$parVals h1 = plotDist(sample = pars$mu_alpha_pos, fontSize = fontSize, binSize = binSize, xLim = c(0,2), xLab = expression(paste(alpha[pos], " (+Learning Rate)"))) From 32937f19996d2e5de71e76a6ec757a3bc6fa1e66 Mon Sep 17 00:00:00 2001 From: Jeong Jinwoo Date: Wed, 17 Dec 2025 19:48:25 +0900 Subject: [PATCH 2/6] feat(pst_Q): add prediction error as generated quantities --- R/R/hBayesDM_model.R | 3 +++ commons/stan_files/pst_Q.stan | 20 +++++++++----------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/R/R/hBayesDM_model.R b/R/R/hBayesDM_model.R index 603ac626..c94acaed 100644 --- a/R/R/hBayesDM_model.R +++ b/R/R/hBayesDM_model.R @@ -298,6 +298,9 @@ hBayesDM_model <- function(task_name = "", if ((model_name == "hgf_ibrb") && (model_type == "single")) { pars <- c(pars, paste0("logit_", names(parameters))) } + if ((task_name == "pst") && (model_name == "Q")) { + pars <- c(pars, "pe") + } pars <- c(pars, "log_lik") if (modelRegressor) { pars <- c(pars, names(regressors)) diff --git a/commons/stan_files/pst_Q.stan b/commons/stan_files/pst_Q.stan index 4a6cbb37..d3111606 100644 --- a/commons/stan_files/pst_Q.stan +++ b/commons/stan_files/pst_Q.stan @@ -31,24 +31,23 @@ transformed parameters { vector[N] alpha; vector[N] beta; - alpha = Phi_approx(mu_pr[1] + sigma[1] * alpha_pr); - beta = Phi_approx(mu_pr[2] + sigma[2] * beta_pr) * 10; + alpha = Phi_approx(mu_pr[1] + sigma[1] * alpha_pr); + beta = Phi_approx(mu_pr[2] + sigma[2] * beta_pr) * 10; } model { // Priors for group-level parameters - mu_pr ~ normal(0, 1); + mu_pr ~ normal(0, 1); sigma ~ normal(0, 0.2); // Priors for subject-level parameters alpha_pr ~ normal(0, 1); - beta_pr ~ normal(0, 1); + beta_pr ~ normal(0, 1); for (i in 1:N) { int co; // Chosen option real delta; // Difference between two options real pe; // Prediction error - //real alpha; vector[6] ev; // Expected values ev = initial_values; @@ -71,19 +70,18 @@ generated quantities { // For group-level parameters real mu_alpha; real mu_beta; + real pe[N, T]; // Prediction error // For log-likelihood calculation real log_lik[N]; - mu_alpha = Phi_approx(mu_pr[1]); - mu_beta = Phi_approx(mu_pr[2]) * 10; + mu_alpha = Phi_approx(mu_pr[1]); + mu_beta = Phi_approx(mu_pr[2]) * 10; { for (i in 1:N) { int co; // Chosen option real delta; // Difference between two options - real pe; // Prediction error - //real alpha; vector[6] ev; // Expected values ev = initial_values; @@ -97,8 +95,8 @@ generated quantities { delta = ev[option1[i, t]] - ev[option2[i, t]]; log_lik[i] += bernoulli_logit_lpmf(choice[i, t] | beta[i] * delta); - pe = reward[i, t] - ev[co]; - ev[co] += alpha[i] * pe; + pe[i, t] = reward[i, t] - ev[co]; + ev[co] += alpha[i] * pe[i, t]; } } } From 1ae9f346ad4b31e16c9667831f74c0c3baff8023 Mon Sep 17 00:00:00 2001 From: Jeong Jinwoo Date: Fri, 30 Jan 2026 12:57:37 +0900 Subject: [PATCH 3/6] feat(pst_gainloss_Q): add prediction error as generated quantities --- R/R/hBayesDM_model.R | 2 +- commons/stan_files/pst_gainloss_Q.stan | 11 +++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/R/R/hBayesDM_model.R b/R/R/hBayesDM_model.R index c94acaed..7908a72a 100644 --- a/R/R/hBayesDM_model.R +++ b/R/R/hBayesDM_model.R @@ -298,7 +298,7 @@ hBayesDM_model <- function(task_name = "", if ((model_name == "hgf_ibrb") && (model_type == "single")) { pars <- c(pars, paste0("logit_", names(parameters))) } - if ((task_name == "pst") && (model_name == "Q")) { + if ((task_name == "pst") && (model_name == "Q" || model_name == "gainloss_Q")) { pars <- c(pars, "pe") } pars <- c(pars, "log_lik") diff --git a/commons/stan_files/pst_gainloss_Q.stan b/commons/stan_files/pst_gainloss_Q.stan index 788b9a4e..a813ed60 100644 --- a/commons/stan_files/pst_gainloss_Q.stan +++ b/commons/stan_files/pst_gainloss_Q.stan @@ -40,7 +40,7 @@ transformed parameters { model { // Priors for group-level parameters - mu_pr ~ normal(0, 1); + mu_pr ~ normal(0, 1); sigma ~ normal(0, 0.2); // Priors for subject-level parameters @@ -77,6 +77,7 @@ generated quantities { real mu_alpha_pos; real mu_alpha_neg; real mu_beta; + real pe[N, T]; // Prediction error // For log-likelihood calculation real log_lik[N]; @@ -89,7 +90,6 @@ generated quantities { for (i in 1:N) { int co; // Chosen option real delta; // Difference between two options - real pe; // Prediction error real alpha; vector[6] ev; // Expected values @@ -104,11 +104,10 @@ generated quantities { delta = ev[option1[i, t]] - ev[option2[i, t]]; log_lik[i] += bernoulli_logit_lpmf(choice[i, t] | beta[i] * delta); - pe = reward[i, t] - ev[co]; - alpha = (pe >= 0) ? alpha_pos[i] : alpha_neg[i]; - ev[co] += alpha * pe; + pe[i, t] = reward[i, t] - ev[co]; + alpha = (pe[i, t] >= 0) ? alpha_pos[i] : alpha_neg[i]; + ev[co] += alpha * pe[i, t]; } } } } - From adf9a90f4cb097179ab81cbca8648c5e3f3e533a Mon Sep 17 00:00:00 2001 From: Jeong Jinwoo Date: Mon, 23 Feb 2026 17:23:44 +0900 Subject: [PATCH 4/6] feat(pst_Q): add p_pred, y_pred --- R/R/pst_Q.R | 4 ++-- commons/models/pst_Q.yml | 1 + commons/stan_files/pst_Q.stan | 22 ++++++++++++++++++++-- 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/R/R/pst_Q.R b/R/R/pst_Q.R index 8c4a502d..f85de370 100644 --- a/R/R/pst_Q.R +++ b/R/R/pst_Q.R @@ -10,7 +10,7 @@ #' @templateVar DATA_COLUMNS "subjID", "type", "choice", "reward" #' @templateVar PARAMETERS \code{alpha} (learning rate), \code{beta} (inverse temperature) #' @templateVar REGRESSORS -#' @templateVar POSTPREDS "y_pred" +#' @templateVar POSTPREDS "p_pred", "y_pred" #' @templateVar LENGTH_DATA_COLUMNS 4 #' @templateVar DETAILS_DATA_1 \item{subjID}{A unique identifier for each subject in the data-set.} #' @templateVar DETAILS_DATA_2 \item{type}{Two-digit number indicating which pair of stimuli were presented for that trial, e.g. 12, 34, or 56. The digit on the left (tens-digit) indicates the presented stimulus for option1, while the digit on the right (ones-digit) indicates that for option2. Code for each stimulus type (1~6) is defined as for 80\% (type 1), 20\% (type 2), 70\% (type 3), 30\% (type 4), 60\% (type 5), 40\% (type 6). The modeling will still work even if different probabilities are used for the stimuli; however, the total number of stimuli should be less than or equal to 6.} @@ -40,5 +40,5 @@ pst_Q <- hBayesDM_model( ), additional_args = NULL, regressors = NULL, - postpreds = c("y_pred"), + postpreds = c("p_pred", "y_pred"), preprocess_func = pst_preprocess_func) diff --git a/commons/models/pst_Q.yml b/commons/models/pst_Q.yml index 67432fd2..6c746f57 100644 --- a/commons/models/pst_Q.yml +++ b/commons/models/pst_Q.yml @@ -38,5 +38,6 @@ parameters: info: [0, 1, 10] regressors: postpreds: +- p_pred - y_pred additional_args: diff --git a/commons/stan_files/pst_Q.stan b/commons/stan_files/pst_Q.stan index d3111606..454c55cd 100644 --- a/commons/stan_files/pst_Q.stan +++ b/commons/stan_files/pst_Q.stan @@ -70,11 +70,26 @@ generated quantities { // For group-level parameters real mu_alpha; real mu_beta; - real pe[N, T]; // Prediction error + + real pe[N, T]; // prediction error // For log-likelihood calculation real log_lik[N]; + // predicted probability of choosing option1 + real p_pred[N, T]; + + // For posterior predictive check + real y_pred[N, T]; + + // Set all posterior predictions to 0 (avoids NULL values) + for (i in 1:N) { + for (t in 1:T) { + p_pred[i, t] = 0.5; + y_pred[i, t] = 0; + } + } + mu_alpha = Phi_approx(mu_pr[1]); mu_beta = Phi_approx(mu_pr[2]) * 10; @@ -95,10 +110,13 @@ generated quantities { delta = ev[option1[i, t]] - ev[option2[i, t]]; log_lik[i] += bernoulli_logit_lpmf(choice[i, t] | beta[i] * delta); + // generate posterior prediction for current trial + p_pred[i, t] = inv_logit(beta[i] * delta); + y_pred[i, t] = bernoulli_rng(p_pred[i, t]); + pe[i, t] = reward[i, t] - ev[co]; ev[co] += alpha[i] * pe[i, t]; } } } } - From a067486bfd8338ee18291c4d2d69b10ebb9154a2 Mon Sep 17 00:00:00 2001 From: Jeong Jinwoo Date: Tue, 24 Feb 2026 19:28:41 +0900 Subject: [PATCH 5/6] feat(pst_decay_Q): implement pst_Q with decay rate --- R/R/hBayesDM_model.R | 2 +- R/R/preprocess_funcs.R | 14 ++- R/R/pst_decay_Q.R | 46 ++++++++++ R/tests/testthat/test_pst_decay_Q.R | 10 +++ commons/models/pst_decay_Q.yml | 46 ++++++++++ commons/stan_files/pst_decay_Q.stan | 128 ++++++++++++++++++++++++++++ 6 files changed, 244 insertions(+), 2 deletions(-) create mode 100644 R/R/pst_decay_Q.R create mode 100644 R/tests/testthat/test_pst_decay_Q.R create mode 100644 commons/models/pst_decay_Q.yml create mode 100644 commons/stan_files/pst_decay_Q.stan diff --git a/R/R/hBayesDM_model.R b/R/R/hBayesDM_model.R index 7908a72a..d68ad861 100644 --- a/R/R/hBayesDM_model.R +++ b/R/R/hBayesDM_model.R @@ -298,7 +298,7 @@ hBayesDM_model <- function(task_name = "", if ((model_name == "hgf_ibrb") && (model_type == "single")) { pars <- c(pars, paste0("logit_", names(parameters))) } - if ((task_name == "pst") && (model_name == "Q" || model_name == "gainloss_Q")) { + if (task_name == "pst") { pars <- c(pars, "pe") } pars <- c(pars, "log_lik") diff --git a/R/R/preprocess_funcs.R b/R/R/preprocess_funcs.R index 2605c236..826a2542 100644 --- a/R/R/preprocess_funcs.R +++ b/R/R/preprocess_funcs.R @@ -789,6 +789,10 @@ pst_preprocess_func <- function(raw_data, general_info) { option2 <- array(-1, c(n_subj, t_max)) choice <- array(-1, c(n_subj, t_max)) reward <- array(-1, c(n_subj, t_max)) + has_trial <- "trialnum" %in% names(raw_data) + if (has_trial) { + trial_gap <- array(0, c(n_subj, t_max)) + } # Write from raw_data to the data arrays for (i in 1:n_subj) { @@ -800,6 +804,12 @@ pst_preprocess_func <- function(raw_data, general_info) { option2[i, 1:t] <- DT_subj$type %% 10 choice[i, 1:t] <- DT_subj$choice reward[i, 1:t] <- DT_subj$reward + if (has_trial) { + if (any(diff(DT_subj$trialnum) <= 0)) { + stop(sprintf("Subject %s: trial indices must be strictly increasing", subj)) + } + trial_gap[i, 1:t] <- c(0, diff(DT_subj$trialnum) - 1) + } } # Wrap into a list for Stan @@ -812,7 +822,9 @@ pst_preprocess_func <- function(raw_data, general_info) { choice = choice, reward = reward ) - + if (has_trial) { + data_list$trial_gap <- trial_gap + } # Returned data_list will directly be passed to Stan return(data_list) } diff --git a/R/R/pst_decay_Q.R b/R/R/pst_decay_Q.R new file mode 100644 index 00000000..71071141 --- /dev/null +++ b/R/R/pst_decay_Q.R @@ -0,0 +1,46 @@ +#' @templateVar MODEL_FUNCTION pst_decay_Q +#' @templateVar CONTRIBUTOR \href{https://github.com/bugoverdose}{Jinwoo Jeong} <\email{jwjeong96@@gmail.com}> +#' @templateVar TASK_NAME Probabilistic Selection Task +#' @templateVar TASK_CODE pst +#' @templateVar TASK_CITE +#' @templateVar MODEL_NAME Q Learning Model with Decay +#' @templateVar MODEL_CODE decay_Q +#' @templateVar MODEL_CITE (Frank et al., 2007) +#' @templateVar MODEL_TYPE Hierarchical +#' @templateVar DATA_COLUMNS "subjID", "trialNum", "type", "choice", "reward" +#' @templateVar PARAMETERS \code{alpha} (learning rate), \code{beta} (inverse temperature), \code{gamma} (decay) +#' @templateVar REGRESSORS +#' @templateVar POSTPREDS "y_pred" +#' @templateVar LENGTH_DATA_COLUMNS 5 +#' @templateVar DETAILS_DATA_1 \item{subjID}{A unique identifier for each subject in the data-set.} +#' @templateVar DETAILS_DATA_2 \item{trialNum}{Nominal integer representing the trial number: 1, 3, 4, 6, ...} +#' @templateVar DETAILS_DATA_3 \item{type}{Two-digit number indicating which pair of stimuli were presented for that trial, e.g. 12, 34, or 56. The digit on the left (tens-digit) indicates the presented stimulus for option1, while the digit on the right (ones-digit) indicates that for option2. Code for each stimulus type (1~6) is defined as for 80\% (type 1), 20\% (type 2), 70\% (type 3), 30\% (type 4), 60\% (type 5), 40\% (type 6). The modeling will still work even if different probabilities are used for the stimuli; however, the total number of stimuli should be less than or equal to 6.} +#' @templateVar DETAILS_DATA_4 \item{choice}{Whether the subject chose the left option (option1) out of the given two options (i.e. if option1 was chosen, 1; if option2 was chosen, 0).} +#' @templateVar DETAILS_DATA_5 \item{reward}{Amount of reward earned as a result of the trial.} +#' @templateVar LENGTH_ADDITIONAL_ARGS 0 +#' +#' @template model-documentation +#' +#' @export +#' @include hBayesDM_model.R +#' @include preprocess_funcs.R + +#' @references +#' Frank, M. J., Moustafa, A. A., Haughey, H. M., Curran, T., & Hutchison, K. E. (2007). Genetic triple dissociation reveals multiple roles for dopamine in reinforcement learning. Proceedings of the National Academy of Sciences, 104(41), 16311-16316. +#' + + +pst_decay_Q <- hBayesDM_model( + task_name = "pst", + model_name = "decay_Q", + model_type = "", + data_columns = c("subjID", "trialNum", "type", "choice", "reward"), + parameters = list( + "alpha" = c(0, 0.5, 1), + "beta" = c(0, 1, 10), + "gamma" = c(0, 0.5, 1) + ), + additional_args = NULL, + regressors = NULL, + postpreds = c("y_pred"), + preprocess_func = pst_preprocess_func) diff --git a/R/tests/testthat/test_pst_decay_Q.R b/R/tests/testthat/test_pst_decay_Q.R new file mode 100644 index 00000000..bc7b8baf --- /dev/null +++ b/R/tests/testthat/test_pst_decay_Q.R @@ -0,0 +1,10 @@ +context("Test pst_decay_Q") +library(hBayesDM) + +test_that("Test pst_decay_Q", { + # Do not run this test on CRAN + skip_on_cran() + + expect_output(pst_decay_Q( + data = "example", niter = 10, nwarmup = 5, nchain = 1, ncore = 1)) +}) diff --git a/commons/models/pst_decay_Q.yml b/commons/models/pst_decay_Q.yml new file mode 100644 index 00000000..0d6f8fff --- /dev/null +++ b/commons/models/pst_decay_Q.yml @@ -0,0 +1,46 @@ +task_name: + code: pst + desc: Probabilistic Selection Task + cite: +model_name: + code: decay_Q + desc: Q Learning Model with Decay + cite: + - Frank, M. J., Moustafa, A. A., Haughey, H. M., Curran, T., & Hutchison, K. E. + (2007). Genetic triple dissociation reveals multiple roles for dopamine in reinforcement + learning. Proceedings of the National Academy of Sciences, 104(41), 16311-16316. +model_type: + code: + desc: Hierarchical +notes: +contributors: +- name: Jinwoo Jeong + email: jwjeong96@gmail.com + link: https://github.com/bugoverdose +data_columns: + subjID: A unique identifier for each subject in the data-set. + trialNum: "Nominal integer representing the trial number: 1, 3, 4, 6, ..." + type: Two-digit number indicating which pair of stimuli were presented for that + trial, e.g. 12, 34, or 56. The digit on the left (tens-digit) indicates the presented + stimulus for option1, while the digit on the right (ones-digit) indicates that + for option2. Code for each stimulus type (1~6) is defined as for 80\% (type 1), + 20\% (type 2), 70\% (type 3), 30\% (type 4), 60\% (type 5), 40\% (type 6). The + modeling will still work even if different probabilities are used for the stimuli; + however, the total number of stimuli should be less than or equal to 6. + choice: Whether the subject chose the left option (option1) out of the given two + options (i.e. if option1 was chosen, 1; if option2 was chosen, 0). + reward: Amount of reward earned as a result of the trial. +parameters: + alpha: + desc: learning rate + info: [0, 0.5, 1] + beta: + desc: inverse temperature + info: [0, 1, 10] + gamma: + desc: decay + info: [0, 0.5, 1] +regressors: +postpreds: +- y_pred +additional_args: diff --git a/commons/stan_files/pst_decay_Q.stan b/commons/stan_files/pst_decay_Q.stan new file mode 100644 index 00000000..e0969185 --- /dev/null +++ b/commons/stan_files/pst_decay_Q.stan @@ -0,0 +1,128 @@ +#include /pre/license.stan + +data { + int N; // Number of subjects + int T; // Maximum # of trials + int Tsubj[N]; // # of trials for acquisition phase + + int option1[N, T]; + int option2[N, T]; + int choice[N, T]; + real reward[N, T]; + real trial_gap[N, T]; // Gap between trials. 0 if trials were consecutive. +} + +transformed data { + // Default values to initialize the vector of expected values + vector[6] initial_values; + initial_values = rep_vector(0, 6); +} + +parameters { + // Group-level parameters + vector[3] mu_pr; + vector[3] sigma; + + // Subject-level parameters for Matt trick + vector[N] alpha_pr; + vector[N] beta_pr; + vector[N] gamma_pr; +} + +transformed parameters { + vector[N] alpha; + vector[N] beta; + vector[N] gamma; + + alpha = Phi_approx(mu_pr[1] + sigma[1] * alpha_pr); + beta = Phi_approx(mu_pr[2] + sigma[2] * beta_pr) * 10; + gamma = Phi_approx(mu_pr[3] + sigma[3] * gamma_pr); +} + +model { + // Priors for group-level parameters + mu_pr ~ normal(0, 1); + sigma ~ normal(0, 0.2); + + // Priors for subject-level parameters + alpha_pr ~ normal(0, 1); + beta_pr ~ normal(0, 1); + gamma_pr ~ normal(0, 1); + + for (i in 1:N) { + int co; // Chosen option + real delta; // Difference between two options + real pe; // Prediction error + vector[6] ev; // Expected values + real decay_rate; // How much expected values decayed + + ev = initial_values; + + // Acquisition Phase + for (t in 1:Tsubj[i]) { + co = (choice[i, t] > 0) ? option1[i, t] : option2[i, t]; + decay_rate = pow(gamma[i], trial_gap[i, t]); + + // Luce choice rule + delta = (ev[option1[i, t]] * decay_rate) - (ev[option2[i, t]] * decay_rate); + target += bernoulli_logit_lpmf(choice[i, t] | beta[i] * delta); + + pe = reward[i, t] - ev[co]; + ev[co] += alpha[i] * pe; + } + } +} + +generated quantities { + // For group-level parameters + real mu_alpha; + real mu_beta; + real mu_gamma; + + real pe[N, T]; // prediction error + + // For log-likelihood calculation + real log_lik[N]; + + // For posterior predictive check + real y_pred[N, T]; + + // Set all posterior predictions to 0 (avoids NULL values) + for (i in 1:N) { + for (t in 1:T) { + y_pred[i, t] = 0; + } + } + + mu_alpha = Phi_approx(mu_pr[1]); + mu_beta = Phi_approx(mu_pr[2]) * 10; + mu_gamma = Phi_approx(mu_pr[3]); + + { + for (i in 1:N) { + int co; // Chosen option + real delta; // Difference between two options + vector[6] ev; // Expected values + real decay_rate; // How much expected values decayed + + ev = initial_values; + log_lik[i] = 0; + + // Acquisition Phase + for (t in 1:Tsubj[i]) { + co = (choice[i, t] > 0) ? option1[i, t] : option2[i, t]; + decay_rate = pow(gamma[i], trial_gap[i, t]); + + // Luce choice rule + delta = (ev[option1[i, t]] * decay_rate) - (ev[option2[i, t]] * decay_rate); + log_lik[i] += bernoulli_logit_lpmf(choice[i, t] | beta[i] * delta); + + // generate posterior prediction for current trial + y_pred[i, t] = bernoulli_rng(inv_logit(beta[i] * delta)); + + pe[i, t] = reward[i, t] - ev[co]; + ev[co] += alpha[i] * pe[i, t]; + } + } + } +} From 09a22f85157eec6ec08455f1b2e91ea0154913ec Mon Sep 17 00:00:00 2001 From: Jeong Jinwoo Date: Fri, 27 Feb 2026 11:55:07 +0900 Subject: [PATCH 6/6] feat(pst_gainloss_Q): add y_pred --- commons/stan_files/pst_gainloss_Q.stan | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/commons/stan_files/pst_gainloss_Q.stan b/commons/stan_files/pst_gainloss_Q.stan index a813ed60..74038605 100644 --- a/commons/stan_files/pst_gainloss_Q.stan +++ b/commons/stan_files/pst_gainloss_Q.stan @@ -82,6 +82,16 @@ generated quantities { // For log-likelihood calculation real log_lik[N]; + // For posterior predictive check + real y_pred[N, T]; + + // Set all posterior predictions to 0 (avoids NULL values) + for (i in 1:N) { + for (t in 1:T) { + y_pred[i, t] = 0; + } + } + mu_alpha_pos = Phi_approx(mu_pr[1]); mu_alpha_neg = Phi_approx(mu_pr[2]); mu_beta = Phi_approx(mu_pr[3]) * 10; @@ -104,6 +114,9 @@ generated quantities { delta = ev[option1[i, t]] - ev[option2[i, t]]; log_lik[i] += bernoulli_logit_lpmf(choice[i, t] | beta[i] * delta); + // generate posterior prediction for current trial + y_pred[i, t] = bernoulli_rng(inv_logit(beta[i] * delta);); + pe[i, t] = reward[i, t] - ev[co]; alpha = (pe[i, t] >= 0) ? alpha_pos[i] : alpha_neg[i]; ev[co] += alpha * pe[i, t];