18 Multiple Indexing and Range Indexing
Stan allows multiple indexes to be provided for containers (i.e., arrays, vectors, and matrices) in a single position, using either an array of integer indexes or range bounds. In many cases, there are functions that provide similar behavior.
Allowing multiple indexes supports inline vectorization of models. For instance, consider the likelihood for a varying-slope, varying-intercept hierarchical linear regression, which could be coded as
for (n in 1:N) {
y[n] ~ normal(alpha[ii[n]] + beta[ii[n]] * x[n], sigma);
}
With multiple indexing, this can be coded in one line, leading to more efficient vectorized code.
y ~ normal(alpha[ii] + rows_dot_product(beta[ii], x), sigma);
This latter version is faster than the loop version; it is equivalent in speed to the clunky assignment to a local variable.
{
vector[N] mu;
for (n in 1:N) {
mu[n] = alpha[ii[n]] + beta[ii[n]] * x[n];
}
y ~ normal(mu, sigma);
}
The boost in speed compared to the original version is because the single call to the normal log density in the sampling statement will be much more memory efficient than the original version.