Issues with indexing and performing mathematical operations on three-dimensional array

Modeling
,
1

I am trying to operate on a three-dimensional array and perform a simple mathematical formula by selecting and indexing a single vector from this 80x131x296 matrix. I declare my three dimensional matrix as below

  int dists[nSubjects, maxTrials, 296];

And I attempt to perform an operation using the a vector of 296 elements as below:

 v = v + ( A[s] ./ ( 1 + dists[s, t, ] ) ) * pe;

However, for this I receive the error “No matches for: Available argument signatures for operator +: Expression is ill formed.”

I am not sure why Stan seems unhappy with me using an indexed vector of the three dimensional array in the formula. Stan remains unhappy even if I attempt to assign this indexed vector to a one-dimensional vector of 296 elements, as below, despite this seeming to be the way that the stan manual recommends indexing from three dimensional arrays.

 curDists = dists[s,t,];

The only way that I can get Stan to accept this formula and allow the indexing is for me to (inefficeintly) iterate through all 296 elements of dists and assign them to the curDists vector one-by-one as follows:

      for(i in 1:296){
        curDists[i] = dists[s,t,i];
      }
      v = v + ( A[s] ./ ( 1 + curDists ) ) * pe;

Stan is finally happy and allows me to index and perform the formula once I assign them one-by-one. This is very inefficient though and I would like to vectorize. Does anyone know why I am having these issues? Why I cannot index a single vector of 296 elements from the 3-dimensional array? Why I need to assign the elements one-by-one to the vector of 296, rather than just assigning the vector all at once? etc.

My full Stan code is below:

data {
  int<lower=1> nSubjects;
  int<lower = 1> nArms; //number of bandit arms
  int<lower = 1> maxTrials; //number of trials
  int<lower = 1> nTrials[nSubjects]; //number of trials per sub
  int<lower = 0, upper = 296> cueChoice[nSubjects, maxTrials]; // which cue chosen
  int<lower = 0, upper = 2> sideChoice[nSubjects,maxTrials]; //index of which side chosen
  real reward[nSubjects,maxTrials]; //outcome of bandit arm pull
  int<lower = 0> leftOption[nSubjects, maxTrials]; // left option
  int<lower = 0> rightOption[nSubjects, maxTrials]; // right option
  
  //int<lower = 0> dists[nSubjects, maxTrials, 296];
  int dists[nSubjects, maxTrials, 296];
}

transformed data {
  vector[296] initV;  // initial values for V
  initV = rep_vector(0.5, 296);
}

parameters {
  // Declare all parameters as vectors for vectorizing
  // Hyper(group)-parameters
  vector[2] mu_pr;
  vector<lower=0>[2] sigma;

  // Subject-level raw parameters (for Matt trick)
  vector[nSubjects] A_pr;    // learning rate
  vector[nSubjects] tau_pr;  // inverse temperature
}

transformed parameters {
  // subject-level parameters
  vector<lower=0, upper=1>[nSubjects] A;
  vector<lower=0, upper=10>[nSubjects] tau;

  for (i in 1:nSubjects) {
    A[i]   = Phi_approx(mu_pr[1]  + sigma[1]  * A_pr[i]);
    tau[i] = Phi_approx(mu_pr[2] + sigma[2] * tau_pr[i]) * 10;
  }
}

model {
  // Hyperparameters
  mu_pr  ~ normal(0, 1);
  sigma ~ normal(0, 0.2);

  // individual parameters
  A_pr   ~ normal(0, 1);
  tau_pr ~ normal(0, 1);
  
  
  
  for (s in 1:nSubjects) {
    vector[296] v; 
    real pe;
    vector[2] w;
    vector[296] curDists;
    vector[296] Avec;
    v = initV;

    for (t in 1:nTrials[s]) {
      w[1] = v[leftOption[s,t]];
      w[2] = v[rightOption[s,t]];
      
      sideChoice[s,t] ~ categorical_logit( tau[s] * w );
      pe = reward[s,t] - v[cueChoice[s,t]];      
      
      for(i in 1:296){
        curDists[i] = dists[s,t,i];
      }
      v = v + ( A[s] ./ ( 1 + curDists ) ) * pe;
      //v = v + ( A[s] ./ ( 1 + dists[s,t,] ) ) * pe;
      
    }
  }    
}

generated quantities {
  // For group level parameters
  real<lower=0, upper=1> mu_A;
  real<lower=0, upper=10> mu_tau;

  // For log likelihood calculation
  real log_lik[nSubjects];

  // For posterior predictive check
  real y_pred[nSubjects, maxTrials];

  // Set all posterior predictions to 0 (avoids NULL values)
  for (i in 1:nSubjects) {
    for (t in 1:maxTrials) {
      y_pred[i, t] = -1;
    }
  }

  mu_A   = Phi_approx(mu_pr[1]);
  mu_tau = Phi_approx(mu_pr[2]) * 10;

  { // local section, this saves time and space
    
    for (s in 1:nSubjects) {
      vector[296] v; 
      real pe;
      vector[2] w;
      vector[296] curDists;
      v = initV;

      log_lik[s] = 0;

      for (t in 1:(nTrials[s])) {
        w[1] = v[leftOption[s,t]];
        w[2] = v[rightOption[s,t]];
        
        // compute log likelihood of current trial
        log_lik[s] += categorical_logit_lpmf(sideChoice[s, t] | tau[s] * w);

        // generate posterior prediction for current trial
        y_pred[s, t] = categorical_rng(softmax(tau[s] * w));

        // prediction error
        pe = reward[s, t] - v[cueChoice[s, t]];
        
        for(i in 1:296){
          curDists[i] = dists[s,t,i];
        }

        // value updating (learning)
        v += ( A[s] ./ ( 1 + curDists ) ) * pe;
        //v += (A[s]/ ( 1 + dists[s, t, 1:2] ) ) * pe;
      }

  }   
    
  }
}
2

I believe the issue is not due to your indexing, but in in the addition:

1 + dists[s, t, ]

There is not a valid signature for:

int + int[]

And so the parser errors.

Additionally, you’re mixing different container types, as v and A[s] are both vector and dists[s,t,] is an int[], which is also not a valid operation.

If you want to use dists in operations with vectors, then you’ll need to declare it as an array of vectors instead:

vector[296] dists[nSubjects, maxTrials];
}
3

Thank you so much! It looks like this worked. Quick clarification question, you said there is no valid signature for int + int[]. I was wondering why this is the case? For example, at this link, it does indicate that vectors and not compatible with arrays and vice versa. But shouldn’t int be compatible with int since they are of the same type/class? What am I misssing, in terms of why they are not compatible with one another?

1 Like
4

It’s because stan doesn’t do automatic broadcasting nor recycling. So it’s not about whether int is compatible with int, but whether scalar is compatible with array.

There are some cases where Stan code is “vectorized”, so you can use scalars together with vectors (instead of arrays). But whether this is supported varies between different functions/operators. Also note that some functions/operators don’t interprete vectors as a container of unrelated elements, but in their vector-mathematics sense.