Calculating the log likelihood for the Boxes model in Statistical Rethinking (2nd Ed)

Modeling
,
1

In Section 16.2 of Statistical Rethinking (2nd Ed), Richard McElreath introduces the “boxes” model to estimate the frequencies of different cognitive strategies in a developmental psychology study. The model is as follows:

library(rethinking)
data(Boxes_model)
cat(Boxes_model)

data{
    int N;
    int y[N];
    int majority_first[N];
}
parameters{
    simplex[5] p;
}
model{
    vector[5] phi;
    
    // prior
    p ~ dirichlet( rep_vector(4,5) );
    
    // probability of data
    for ( i in 1:N ) {
        if ( y[i]==2 ) phi[1]=1; else phi[1]=0; // majority
        if ( y[i]==3 ) phi[2]=1; else phi[2]=0; // minority
        if ( y[i]==1 ) phi[3]=1; else phi[3]=0; // maverick
        phi[4]=1.0/3.0;                         // random
        if ( majority_first[i]==1 )             // follow first
            if ( y[i]==2 ) phi[5]=1; else phi[5]=0;
        else
            if ( y[i]==3 ) phi[5]=1; else phi[5]=0;
        
        // compute log( p_s * Pr(y_i|s )
        for ( j in 1:5 ) phi[j] = log(p[j]) + log(phi[j]);
        // compute average log-probability of y_i
        target += log_sum_exp( phi );
    }
}

I would like to add a predictor to this model (e.g., gender) and compare the different models using WAIC or LOO. To do this, I need to calculate the log-likelihood within the generated quantities block of the Stan code. But since this Stan code has quite an unusual model block, I’m not entirely sure how to do this correctly. Please can someone help?

2

There is a loop

    for ( i in 1:N ) {
...
        target += log_sum_exp( phi )
    }

where target += adds the log likelihood contribution given each observation.
The corresponding generated quantities block would include

    vector[5] phi;
    vector[N] log_lik;
    for ( i in 1:N ) {
...
        log_lik[i] = log_sum_exp( phi )
    }

In place of ... include the same computation for phi as in the model block. You need to repeat the computation of phi as phi was declared in the model block and thus not visible in the generated quantities block. If you would compute phi in transformed parameters block, you would not need to repeat the computation, but then the Stan output would include draws of phi, which you may not want to have.

We don’t recommend using WAIC, see more in CV-FAQ: 21 How are LOO and WAIC related?

3

Thanks so much @avehtari. I would prefer to avoid repetition in the code, and I’m okay with having phi in the output, so I have refactored the code as follows. It returns the same answers as before and I’m able to compute LOO from the log-likelihood.

data{
  int N;
  array[N] int y;
  array[N] int majority_first;
}
parameters{
  simplex[5] p;
}
transformed parameters {
  array[N,5] real phi;
  
  for ( i in 1:N ) {
    
    // compute phi
    if ( y[i]==2 ) phi[i,1]=1; else phi[i,1]=0; // majority
    if ( y[i]==3 ) phi[i,2]=1; else phi[i,2]=0; // minority
    if ( y[i]==1 ) phi[i,3]=1; else phi[i,3]=0; // maverick
    phi[i,4]=1.0/3.0;                           // random
    if ( majority_first[i]==1 )                 // follow first
        if ( y[i]==2 ) phi[i,5]=1; else phi[i,5]=0;
    else
        if ( y[i]==3 ) phi[i,5]=1; else phi[i,5]=0;
        
    // compute log( p_s * Pr(y_i|s )
    for ( j in 1:5 ) phi[i,j] = log(p[j]) + log(phi[i,j]);
    
  }
}
model{
  // prior
  p ~ dirichlet( rep_vector(4,5) );
  // compute average log-probability of y_i
  for ( i in 1:N ) target += log_sum_exp( phi[i,] );
}
generated quantities{
  // log-likelihood
  vector[N] log_lik;
  for ( i in 1:N ) log_lik[i] = log_sum_exp( phi[i,] );
}

Computed from 4000 by 629 log-likelihood matrix

         Estimate   SE
elpd_loo   -635.0 10.5
p_loo         2.8  0.1
looic      1270.0 21.0
------
Monte Carlo SE of elpd_loo is 0.0.

All Pareto k estimates are good (k < 0.5).
See help('pareto-k-diagnostic') for details.
1 Like