feat(pst_Q): add p_pred, y_pred

Author: bugoverdose

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)

commons/models/pst_Q.yml

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

commons/stan_files/pst_Q.stan

@@ -70,11 +70,26 @@ 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
+  
+  real pe[N, T]; // prediction error
 
   // For log-likelihood calculation
   real log_lik[N];
 
+  // 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;
 
@@ -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];
       }
     }
   }
 }
-