Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
170 lines (147 sloc) 5.28 KB
data {
int<lower=1> N; // Number of subjects
int<lower=1> T; // Max trials per subject
int<lower=1> B; // Number of bet options
int<lower=1,upper=T> Tsubj[N]; // number of trials/subject
int<lower=0,upper=2> col_chosen[N,T]; // chosen color index
int<lower=0,upper=B> bet_chosen[N,T]; // chosen bet indexs
vector[B] bet_delay; // vector of bet delays
real gain[N,T,B]; // gain: (capital + capital * bet_prop)
real loss[N,T,B]; // loss: (capital - capital * bet_prop)
real prop_red[N,T]; // proportion of red boxes
real prop_chosen[N,T]; // proportion of chosen boxes
}
parameters {
// Declare all parameters as vectors for vectorizing
// Hyper(group)-parameters
vector[5] mu_pr;
vector<lower=0>[5] sigma;
// Subject-level raw parameters (for Matt trick)
vector[N] alpha_pr;
vector[N] rho_pr;
vector[N] gamma_pr;
vector[N] c_pr;
vector[N] beta_pr;
}
transformed parameters {
// subject-level parameters
vector<lower=0,upper=5>[N] alpha;
vector<lower=0>[N] rho;
vector<lower=0>[N] gamma;
vector<lower=0,upper=1>[N] c;
vector<lower=0>[N] beta;
for (i in 1:N) {
alpha[i] = Phi_approx( mu_pr[1] + sigma[1] * alpha_pr[i] ) * 5;
c[i] = Phi_approx( mu_pr[4] + sigma[4] * c_pr[i] );
}
rho = exp(mu_pr[2] + sigma[2] * rho_pr);
gamma = exp(mu_pr[3] + sigma[3] * gamma_pr);
beta = exp(mu_pr[5] + sigma[5] * beta_pr);
}
model {
// Hyperpriors (vectorized)
mu_pr ~ normal(0, 1);
sigma ~ normal(0, 0.2);
// Individual parameters
alpha_pr ~ normal(0,1);
rho_pr ~ normal(0,1);
gamma_pr ~ normal(0,1);
c_pr ~ normal(0,1);
beta_pr ~ normal(0,1);
// subject loop and trial loop
for (i in 1:N) {
// Define vectors
vector[B] gain_util;
vector[B] loss_util;
vector[B] bet_util;
vector[2] col_util;
// Model
for (t in 1:Tsubj[i]) { // Need to set trial/vectorized parts correctly
// Assign probability of choosing color (fix bias red)
col_util[1] = (c[i] * pow(prop_red[i,t], alpha[i])) / (c[i] * pow(prop_red[i,t], alpha[i]) + ((1 - c[i]) * pow(1 - prop_red[i,t], alpha[i])));
col_util[2] = 1 - col_util[1];
// Increment log likelihood for color choice
col_chosen[i,t] ~ categorical(col_util);
// For each bet option
for (b in 1:B) {
// Assign gain/loss utilities
gain_util[b] = log(1 + gain[i,t,b] * 1);
loss_util[b] = log(1 + loss[i,t,b] * rho[i]);
}
// Utility of all bets
bet_util = gain_util * prop_chosen[i,t] + loss_util * (1 - prop_chosen[i,t]);
// Utility of bet with delays
// bet_util = ((beta[i] + bet_util) / beta[i]) - beta[i] * bet_delay;
bet_util = bet_util - beta[i] * bet_delay;
// Increment log likelihood for choosing bet
bet_chosen[i,t] ~ categorical_logit(bet_util * gamma[i]);
}
}
}
generated quantities {
// Define group level parameters
real<lower=0,upper=5> mu_alpha;
real<lower=0> mu_rho;
real<lower=0> mu_gamma;
real<lower=0,upper=5> mu_c;
real<lower=0> mu_beta;
// Define log likelihood vector
real log_lik[N];
real y_hat_col[N,T];
real y_hat_bet[N,T];
real bet_utils[N,T,B];
for (j in 1:N) {
for (k in 1:T) {
y_hat_col[j,k] = 0;
y_hat_bet[j,k] = 0;
for (b in 1:B) {
bet_utils[j,k,b] = 0;
}
}
}
// Assign group level parameters
mu_alpha = Phi_approx(mu_pr[1]) * 5;
mu_rho = exp(mu_pr[2]);
mu_gamma = exp(mu_pr[3]);
mu_c = Phi_approx(mu_pr[4]);
mu_beta = exp(mu_pr[5]);
{ // local section, this saves time and space
for (i in 1:N) {
// Define vectors
vector[B] gain_util;
vector[B] loss_util;
vector[B] bet_util;
vector[2] col_util;
log_lik[i] = 0;
// Model
for (t in 1:Tsubj[i]) { // Need to set trial/vectorized parts correctly
// Assign probability of choosing color (fix bias red)
col_util[1] = (c[i] * pow(prop_red[i,t], alpha[i])) / (c[i] * pow(prop_red[i,t], alpha[i]) + ((1 - c[i]) * pow(1 - prop_red[i,t], alpha[i])));
col_util[2] = 1 - col_util[1];
// Increment log likelihood for color choice
log_lik[i] = log_lik[i] + categorical_lpmf(col_chosen[i,t] | col_util);
// Posterior prediction for bet
y_hat_col[i,t] = categorical_rng(col_util);
// For each bet option
for (b in 1:B) {
// Assign gain/loss utilities
gain_util[b] = log(1 + gain[i,t,b] * 1);
loss_util[b] = log(1 + loss[i,t,b] * rho[i]);
}
// Utility of all bets
bet_util = gain_util * prop_chosen[i,t] + loss_util * (1 - prop_chosen[i,t]);
// Utility of bet with delays
// bet_util = ((beta[i] + bet_util) / beta[i]) - beta[i] * bet_delay;
bet_util = bet_util - beta[i] * bet_delay;
// Increment log likelihood for choosing bet
log_lik[i] += categorical_logit_lpmf(bet_chosen[i,t] | bet_util * gamma[i]);
// Posterior prediction for bet
y_hat_bet[i,t] = categorical_rng(softmax(bet_util * gamma[i]));
// Save bet utility
for (b in 1:B) {
bet_utils[i,t,b] = bet_util[b];
}
}
}
}
}
You can’t perform that action at this time.