diff --git a/R/R/hBayesDM_model.R b/R/R/hBayesDM_model.R index 603ac626..d68ad861 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") { + pars <- c(pars, "pe") + } pars <- c(pars, "log_lik") if (modelRegressor) { pars <- c(pars, names(regressors)) 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_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/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/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)"))) 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_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/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_Q.stan b/commons/stan_files/pst_Q.stan index 4a6cbb37..454c55cd 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,33 @@ 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; + // 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; { 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,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); - pe = reward[i, t] - ev[co]; - ev[co] += alpha[i] * pe; + // 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]; } } } } - 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]; + } + } + } +} diff --git a/commons/stan_files/pst_gainloss_Q.stan b/commons/stan_files/pst_gainloss_Q.stan index 788b9a4e..74038605 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,10 +77,21 @@ 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]; + // 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; @@ -89,7 +100,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 +114,13 @@ 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; + // 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]; } } } } -