Conditional subsetting/slicing of vectors in Stan

I have a model where the different observations may be specified by different distributions. I’m trying to vectorize this in order to get the performance to a reasonable level. The data is a ragged array and I am specifying the model to use with an integer vector. I can’t seem to figure out how to accomplish this in Stan.

Basically the data looks like this:

model_to_use = [1, 1, 1, 1, 2, 2, 2 ];
observations = [1.6, 1.7, 1.7, 1.7, 80, 91, 99];

what I want to accomplish is:

xxx[1:4] = bernoulli_lpmf( observations[model_to_use ==1] | 1.75);
xxx[5:7] = bernoulli_lpmf( observations[model_to_use ==2] | 91);

Is there a way to accomplish this through vectorizing, or some other method? Thanks!


The Stan language does not support subsetting by logical conditions. You can do

bernoulli_lpmf(observations | {1.75, 91}[model_to_use]);

Thanks Ben! That makes sense for my example.

Taking it one step further, what about a different type of distribution?

xxx[1:4] = bernoulli_lpmf( observations[model_to_use ==1] | 1.75);
xxx[5:7] = poisson_lpmf( observations[model_to_use ==2] | 91);

Best to declare them differently in the data block.

Got it. That’s what I ended up doing so I’m glad to hear I’m on the right track.

Cool syntax!