Minor enhancement for the pst_Q model

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))

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)
 }

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)

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)

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)")))

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))
+})

commons/models/pst_Q.yml

@@ -38,5 +38,6 @@ parameters:
     info: [0, 1, 10]
 regressors:
 postpreds:
+- p_pred
 - y_pred
 additional_args:

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:

commons/stan_files/pst_Q.stan

@@ -31,24 +31,23 @@ transformed parameters {
   vector<lower=0,upper=1>[N] alpha;
   vector<lower=0,upper=10>[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<lower=0,upper=1>  mu_alpha;
   real<lower=0,upper=10> 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<lower=0,upper=1> p_pred[N, T];
+
+  // For posterior predictive check
+  real<lower=0,upper=1> 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];
       }
     }
   }
 }
-