Reduce_sum indexing


Could someone please confirm reduce_sum indexing is correct?y_slice it is indexed between (1: end-start+1); the rest of vectors as parameters (pe,pc,lambda) are indexed between start:end. They are all describing data or parameters of the same participant.

functions {
  real partial_sum(array[] int y_slice,
                   int start, int end,
                   vector pe, vector pc, vector lambda){
    real summ=0;
    for (i in 1:(end-start+1)) {
    summ += log_mix(lambda[i+start-1], poisson_lpmf(y_slice[i] | pe[i+start-1]),poisson_lpmf(y_slice[i] | pc[i+start-1]));
    }
    return summ;
  }
}

model {

 //blind trial a mixture poisson distribution;

  int grainsize = 1;
  target += reduce_sum(partial_sum, y,
                       grainsize,
                       pe,pc,lambda);
}

I suggest you to test this by using the partial_sum function to add the total log likelihood in - say - 4 chunks to the model… or all at once. You should get matching results (so drop reduce_sum for doing that).

I think you might find it easier to verify correctness by taking slices and then using direct indexing.

  real partial_sum(array[] int y_slice,
                   int start, int end,
                   vector pe, vector pc, vector lambda){
    int len = size(y_slice);  // == end - start + 1;
    vector[len] lamba_slice = lambda[start:end];
    vector[len] pe_slice = pe[start:end];
    vector[len] pc_slice = pc[start:end];
    real lp = 0;
    for (i in 1:len) {
      lp += log_mix(lambda[i], 
                    poisson_lpmf(y_slice[i] | pe[i]),
                    poisson_lpmf(y_slice[i] | pc[i]));
    }
    return lp;
  }
}

This looks like it’s going to be problematic because pe, pc, and lambda all vary per observation. In a standard mixture model, all of pe, pc, and lambda would be scalars with single values. Without something else going on in the likelihood, this is going to degenerate to lambda[i] going to 0 or 1 and one of pe[i] or pc[i] going to y_slice[i] (if both go to y_slice[i] then lambda[i] is not identified).