Skip to content

Commit

Permalink
STAN models for ExpPower. #147 #162 (not working reliably)
Browse files Browse the repository at this point in the history
  • Loading branch information
Ben Vincent committed Dec 23, 2016
1 parent 99f6149 commit 1ed1970
Show file tree
Hide file tree
Showing 5 changed files with 223 additions and 14 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
for chain = 1:nchains
initialParams(chain).groupKmu = normrnd(0.001,0.1);
initialParams(chain).groupKsigma = rand*5;
initialParams(chain).groupTAUmu = normrnd(00,0.1); % <---- check
initialParams(chain).groupTAUmu = normrnd(0,0.1); % <---- check
initialParams(chain).groupTAUsigma = rand*5;% <---- check
initialParams(chain).groupW = rand;
initialParams(chain).groupALPHAmu = rand*10;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
% Generate initial values of the root nodes
nExperimentFiles = obj.data.getNExperimentFiles();
for chain = 1:nchains
initialParams(chain).k = unifrnd(0, 0.5, [nExperimentFiles,1]);
initialParams(chain).k = unifrnd(0.01, 0.5, [nExperimentFiles,1]);
initialParams(chain).tau = unifrnd(0.01, 2, [nExperimentFiles,1]);
initialParams(chain).epsilon = 0.1 + rand([nExperimentFiles,1])/10;
initialParams(chain).alpha = abs(normrnd(0.01,10,[nExperimentFiles,1]));
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
functions {
vector matrix_pow_elementwise(vector delay, vector tau){
// can't (currently) do elementwise matrix power operation, so manually loop
vector[rows(delay)] output;
for (i in 1:num_elements(delay)){
output[i] = pow(delay[i], tau[i]);
}
return output;
}

real psychometric_function(real alpha, real epsilon, real VA, real VB){
// returns probability of choosing B (delayed reward)
return epsilon + (1-2*epsilon) * Phi( (VB-VA) / alpha);
}

vector df_exp_power(vector reward, vector k, vector tau, vector delay){
//return reward .*( exp( -k .* (delay ^ tau) ) );
vector[rows(delay)] delay_to_power_tau;
delay_to_power_tau = matrix_pow_elementwise(delay,tau);
return reward .*( exp( -k .* delay_to_power_tau ) );
}

vector discounting(vector A, vector B, vector DA, vector DB, vector k, vector tau, vector epsilon, vector alpha){
vector[rows(A)] VA;
vector[rows(B)] VB;
vector[rows(A)] P;
// calculate present subjective values
VA = df_exp_power(A, k, tau, DA);
VB = df_exp_power(B, k, tau, DB);
// calculate probability of choosing delayed reward (B; coded as R=1)
for (t in 1:rows(A)){
P[t] = psychometric_function(alpha[t], epsilon[t], VA[t], VB[t]);
}
return P;
}
}

data {
int <lower=1> totalTrials;
int <lower=1> nRealExperimentFiles;
vector[totalTrials] A;
vector[totalTrials] B;
vector<lower=0>[totalTrials] DA;
vector<lower=0>[totalTrials] DB;
int <lower=0,upper=1> R[totalTrials];
int <lower=0,upper=nRealExperimentFiles> ID[totalTrials];
}

parameters {
real k_mu;
real<lower=0> k_sigma;
vector[nRealExperimentFiles+1] k; // +1 for unobserved participant

real tau_mu;
real<lower=0> tau_sigma;
vector<lower=0>[nRealExperimentFiles+1] tau; // +1 for unobserved participant

real alpha_mu;
real <lower=0> alpha_sigma;
vector<lower=0>[nRealExperimentFiles+1] alpha; // +1 for unobserved participant

real <lower=0,upper=1> omega;
real <lower=0> kappa;
vector<lower=0,upper=0.5>[nRealExperimentFiles+1] epsilon; // +1 for unobserved participants
}

transformed parameters {
vector[totalTrials] P;
P = discounting(A, B, DA, DB, k[ID], tau[ID], epsilon[ID], alpha[ID]);
}

model {
k_mu ~ normal(0.01, 2.5); // TODO : pick this in a more meaningul manner
k_sigma ~ inv_gamma(0.1,0.1); // TODO : pick this in a more meaningul manner
k ~ normal(k_mu, k_sigma);

tau_mu ~ normal(0.01, 2.5); // TODO : pick this in a more meaningul manner
tau_sigma ~ inv_gamma(0.1,0.1); // TODO: pick this in a more meaningul manner
tau ~ normal(tau_mu, tau_sigma);

alpha_mu ~ uniform(0,100);
alpha_sigma ~ inv_gamma(0.01,0.01);
alpha ~ normal(alpha_mu, alpha_sigma);

omega ~ beta(1.1, 10.9); // mode for lapse rate
kappa ~ gamma(0.1,0.1); // concentration parameter
epsilon ~ beta(omega*(kappa-2)+1 , (1-omega)*(kappa-2)+1 );

R ~ bernoulli(P);
}

generated quantities { // NO VECTORIZATION IN THIS BLOCK
int <lower=0,upper=1> Rpostpred[totalTrials];
for (t in 1:totalTrials){
Rpostpred[t] = bernoulli_rng(P[t]);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
functions {
vector matrix_pow_elementwise(vector delay, vector tau){
// can't (currently) do elementwise matrix power operation, so manually loop
vector[rows(delay)] output;
for (i in 1:num_elements(delay)){
output[i] = pow(delay[i], tau[i]);
}
return output;
}

real psychometric_function(real alpha, real epsilon, real VA, real VB){
// returns probability of choosing B (delayed reward)
return epsilon + (1-2*epsilon) * Phi( (VB-VA) / alpha);
}

vector df_exp_power(vector reward, vector k, vector tau, vector delay){
//return reward .*( exp( -k .* (delay ^ tau) ) );
vector[rows(delay)] delay_to_power_tau;
delay_to_power_tau = matrix_pow_elementwise(delay,tau);
return reward .*( exp( -k .* delay_to_power_tau ) );
}

vector discounting(vector A, vector B, vector DA, vector DB, vector k, vector tau, vector epsilon, vector alpha){
vector[rows(A)] VA;
vector[rows(B)] VB;
vector[rows(A)] P;
// calculate present subjective values
VA = df_exp_power(A, k, tau, DA);
VB = df_exp_power(B, k, tau, DB);
// calculate probability of choosing delayed reward (B; coded as R=1)
for (t in 1:rows(A)){
P[t] = psychometric_function(alpha[t], epsilon[t], VA[t], VB[t]);
}
return P;
}
}

data {
int <lower=1> totalTrials;
int <lower=1> nRealExperimentFiles;
vector[totalTrials] A;
vector[totalTrials] B;
vector<lower=0>[totalTrials] DA;
vector<lower=0>[totalTrials] DB;
int <lower=0,upper=1> R[totalTrials];
int <lower=0,upper=nRealExperimentFiles> ID[totalTrials];
}

parameters {
// Discounting parameters
vector[nRealExperimentFiles] k;
vector<lower=0>[nRealExperimentFiles] tau;

// Psychometric function parameters
real alpha_mu;
real <lower=0> alpha_sigma;
vector<lower=0>[nRealExperimentFiles+1] alpha;

real <lower=0,upper=1> omega;
real <lower=0> kappa;
vector<lower=0,upper=0.5>[nRealExperimentFiles+1] epsilon;
}

transformed parameters {
vector[totalTrials] P;
P = discounting(A, B, DA, DB, k[ID], tau[ID], epsilon[ID], alpha[ID]);
}

model {
alpha_mu ~ uniform(0,100);
alpha_sigma ~ inv_gamma(0.01,0.01);
alpha ~ normal(alpha_mu, alpha_sigma);

omega ~ beta(1.1, 10.9); // mode for lapse rate
kappa ~ gamma(0.1,0.1); // concentration parameter
epsilon ~ beta(omega*(kappa-2)+1 , (1-omega)*(kappa-2)+1 );

// no hierarchical inference for k
k ~ normal(0, 2);
tau ~ normal(1, 1);

R ~ bernoulli(P);
}

generated quantities { // NO VECTORIZATION IN THIS BLOCK
int <lower=0,upper=1> Rpostpred[totalTrials];

for (t in 1:totalTrials){
Rpostpred[t] = bernoulli_rng(P[t]);
}
}
Original file line number Diff line number Diff line change
@@ -1,13 +1,40 @@
functions {

vector matrix_pow_elementwise(vector delay, vector tau){
// can't (currently) do elementwise matrix power operation, so manually loop
vector[rows(delay)] output;
for (i in 1:num_elements(delay)){
output[i] = pow(delay[i], tau[i]);
}
return output;
}

real psychometric_function(real alpha, real epsilon, real VA, real VB){
// returns probability of choosing B (delayed reward)
return epsilon + (1-2*epsilon) * Phi( (VB-VA) / alpha);
}

vector df_exp_power(vector reward, vector k, vector tau, vector delay){
//return reward .*( exp( -k .* (delay ^ tau) ) );
return reward .*( exp( -k .* matrix_pow(delay,tau) ) );
vector[rows(delay)] delay_to_power_tau;
delay_to_power_tau = matrix_pow_elementwise(delay,tau);
return reward .*( exp( -k .* delay_to_power_tau ) );
}

vector discounting(vector A, vector B, vector DA, vector DB, vector k, vector tau, vector epsilon, vector alpha){
vector[rows(A)] VA;
vector[rows(B)] VB;
vector[rows(A)] P;
// calculate present subjective values
VA = df_exp_power(A, k, tau, DA);
VB = df_exp_power(B, k, tau, DB);
// calculate probability of choosing delayed reward (B; coded as R=1)
for (t in 1:rows(A)){
P[t] = psychometric_function(alpha[t], epsilon[t], VA[t], VB[t]);
}
return P;
}

}

data {
Expand All @@ -22,23 +49,18 @@ data {
}

parameters {
vector[nRealExperimentFiles] k;
// Discounting parameters
vector[nRealExperimentFiles] k;
vector<lower=0>[nRealExperimentFiles] tau;

// Psychometric function parameters
vector<lower=0>[nRealExperimentFiles] alpha;
vector<lower=0,upper=0.5>[nRealExperimentFiles] epsilon;
}

transformed parameters {
vector[totalTrials] VA;
vector[totalTrials] VB;
vector[totalTrials] P;

VA = df_exp_power(A, k[ID], tau[ID], DA);
VB = df_exp_power(B, k[ID], tau[ID], DB);

for (t in 1:totalTrials){
P[t] = psychometric_function(alpha[ID[t]], epsilon[ID[t]], VA[t], VB[t]);
}
P = discounting(A, B, DA, DB, k[ID], tau[ID], epsilon[ID], alpha[ID]);
}

model {
Expand All @@ -52,7 +74,6 @@ model {

generated quantities { // NO VECTORIZATION IN THIS BLOCK ?
int <lower=0,upper=1> Rpostpred[totalTrials];

for (t in 1:totalTrials){
Rpostpred[t] = bernoulli_rng(P[t]);
}
Expand Down

0 comments on commit 1ed1970

Please sign in to comment.