From 1ed1970b5a0d9c1788d9cd25d663b3aad3862637 Mon Sep 17 00:00:00 2001 From: Ben Vincent Date: Fri, 23 Dec 2016 10:58:57 +0000 Subject: [PATCH] STAN models for ExpPower. #147 #162 (not working reliably) --- .../ModelHierarchicalExpPower.m | 2 +- .../bens_new_model/ModelSeparateExpPower.m | 2 +- .../bens_new_model/hierarchicalExpPower.stan | 97 +++++++++++++++++++ .../bens_new_model/mixedExpPower.stan | 91 +++++++++++++++++ .../bens_new_model/separateExpPower.stan | 45 ++++++--- 5 files changed, 223 insertions(+), 14 deletions(-) create mode 100644 ddToolbox/models/parametric_models/bens_new_model/hierarchicalExpPower.stan create mode 100644 ddToolbox/models/parametric_models/bens_new_model/mixedExpPower.stan diff --git a/ddToolbox/models/parametric_models/bens_new_model/ModelHierarchicalExpPower.m b/ddToolbox/models/parametric_models/bens_new_model/ModelHierarchicalExpPower.m index b8d07556..07a107e2 100644 --- a/ddToolbox/models/parametric_models/bens_new_model/ModelHierarchicalExpPower.m +++ b/ddToolbox/models/parametric_models/bens_new_model/ModelHierarchicalExpPower.m @@ -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; diff --git a/ddToolbox/models/parametric_models/bens_new_model/ModelSeparateExpPower.m b/ddToolbox/models/parametric_models/bens_new_model/ModelSeparateExpPower.m index e06eeaf2..fa66cda6 100644 --- a/ddToolbox/models/parametric_models/bens_new_model/ModelSeparateExpPower.m +++ b/ddToolbox/models/parametric_models/bens_new_model/ModelSeparateExpPower.m @@ -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])); diff --git a/ddToolbox/models/parametric_models/bens_new_model/hierarchicalExpPower.stan b/ddToolbox/models/parametric_models/bens_new_model/hierarchicalExpPower.stan new file mode 100644 index 00000000..a5a1c8f3 --- /dev/null +++ b/ddToolbox/models/parametric_models/bens_new_model/hierarchicalExpPower.stan @@ -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 totalTrials; + int nRealExperimentFiles; + vector[totalTrials] A; + vector[totalTrials] B; + vector[totalTrials] DA; + vector[totalTrials] DB; + int R[totalTrials]; + int ID[totalTrials]; +} + +parameters { + real k_mu; + real k_sigma; + vector[nRealExperimentFiles+1] k; // +1 for unobserved participant + + real tau_mu; + real tau_sigma; + vector[nRealExperimentFiles+1] tau; // +1 for unobserved participant + + real alpha_mu; + real alpha_sigma; + vector[nRealExperimentFiles+1] alpha; // +1 for unobserved participant + + real omega; + real kappa; + vector[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 Rpostpred[totalTrials]; + for (t in 1:totalTrials){ + Rpostpred[t] = bernoulli_rng(P[t]); + } +} diff --git a/ddToolbox/models/parametric_models/bens_new_model/mixedExpPower.stan b/ddToolbox/models/parametric_models/bens_new_model/mixedExpPower.stan new file mode 100644 index 00000000..2ac8b519 --- /dev/null +++ b/ddToolbox/models/parametric_models/bens_new_model/mixedExpPower.stan @@ -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 totalTrials; + int nRealExperimentFiles; + vector[totalTrials] A; + vector[totalTrials] B; + vector[totalTrials] DA; + vector[totalTrials] DB; + int R[totalTrials]; + int ID[totalTrials]; +} + +parameters { + // Discounting parameters + vector[nRealExperimentFiles] k; + vector[nRealExperimentFiles] tau; + + // Psychometric function parameters + real alpha_mu; + real alpha_sigma; + vector[nRealExperimentFiles+1] alpha; + + real omega; + real kappa; + vector[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 Rpostpred[totalTrials]; + + for (t in 1:totalTrials){ + Rpostpred[t] = bernoulli_rng(P[t]); + } +} diff --git a/ddToolbox/models/parametric_models/bens_new_model/separateExpPower.stan b/ddToolbox/models/parametric_models/bens_new_model/separateExpPower.stan index 741fcb3e..a2d94e97 100644 --- a/ddToolbox/models/parametric_models/bens_new_model/separateExpPower.stan +++ b/ddToolbox/models/parametric_models/bens_new_model/separateExpPower.stan @@ -1,4 +1,14 @@ 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); @@ -6,8 +16,25 @@ functions { 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 { @@ -22,23 +49,18 @@ data { } parameters { - vector[nRealExperimentFiles] k; + // Discounting parameters + vector[nRealExperimentFiles] k; vector[nRealExperimentFiles] tau; + + // Psychometric function parameters vector[nRealExperimentFiles] alpha; vector[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 { @@ -52,7 +74,6 @@ model { generated quantities { // NO VECTORIZATION IN THIS BLOCK ? int Rpostpred[totalTrials]; - for (t in 1:totalTrials){ Rpostpred[t] = bernoulli_rng(P[t]); }