Does to_vector have a performance hit?

I’ve been playing with a simple model that performs the equivalent of a paired-sample T-test using neural spike count data where the observation window varies for each cell. I believe I’ve managed to vectorize the code as much as possible, but was wondering if there were additional performance gains i could achieve. For example:

  • Instead of using to_vector, should I be declaring the data and parameters as vectors instead of arrays? If I declare the parameters as arrays, how do I set bounds on them?
  • Does having real<lower=0> in the parameters block and T[0, ] on the corresponding prior in the model block lead to a performance hit? Should I prefer one over the other?
  • Would it be better to try and concatenate the A and B arrays when calculating the poisson likelihood (so poisson is only calculated once instead of twice)?

Any other advice would be greatly appreciated. I’m new to STAN, so I’m trying to learn various tips and tricks to improve my code.

data {
    int<lower=1> n_cells;
    real<lower=0> sample_time_A[n_cells];
    real<lower=0> sample_time_B[n_cells];
    int<lower=0> spike_count_A[n_cells];
    int<lower=0> spike_count_B[n_cells];
}

parameters {
    real<lower=0> rate_alpha;
    real<lower=0> rate_beta;
    real<lower=0> rate_cell[n_cells];
    
    real<lower=0> rate_ratio_mean;
    real<lower=0> rate_ratio_sd;
    real<lower=0> rate_ratio_cell[n_cells];
}

transformed parameters {
    real rate_A;
    real rate_B;
    real rate_change;
    
    vector[n_cells] rate_cell_A;
    vector[n_cells] rate_cell_B;
    vector[n_cells] rate_cell_change;
    
    rate_A = rate_alpha/rate_beta;
    rate_B = rate_A * rate_ratio_mean;
    rate_change = rate_B - rate_A;
    
    rate_cell_A = to_vector(rate_cell);
    rate_cell_B = to_vector(rate_cell) .* to_vector(rate_ratio_cell);
    rate_cell_change = rate_cell_B - rate_cell_A;
}

model {
    vector[n_cells] lambda_A;
    vector[n_cells] lambda_B;
    
    rate_alpha ~ gamma(0.5, 0.1);
    rate_beta ~ gamma(0.1, 0.1);
    
    rate_ratio_mean ~ normal(1, 1)T[0, ];
    rate_ratio_sd ~ normal(0, 1)T[0, ];
    
    rate_cell ~ gamma(rate_alpha, rate_beta);
    rate_ratio_cell ~ normal(rate_ratio_mean, rate_ratio_sd);
    
    lambda_A = to_vector(rate_cell) .* to_vector(sample_time_A);
    lambda_B = to_vector(rate_cell) .* to_vector(rate_ratio_cell) .* to_vector(sample_time_B);
    spike_count_A ~ poisson(lambda_A);
    spike_count_B ~ poisson(lambda_B);
}

Yes

The bounds are specified the same way regardless of whether they are vectors or arrays

The real<lower = 0> is necessary while the truncation is only necessary if you need the correct posterior log-density for something post-estimation like Bayes factors or something. During the sampling, you only need the posterior log-density to be correct up to a constant, and in your case, the truncation only affects the posterior log-density by a constant.

Doing so probably won’t make a noticeable difference to the runtime.

Thank you! I was missing the syntax for bounded vectors (none of the examples I looked at had it, so I played around and finally figured it out):

vector<lower=0>[n_cells] rate_cell_B;

As an added bonus, the model is easier to read and compiles much faster as well (which is great for iterative testing). Here’s the latest version of the model:

data {
    int<lower=1> n_cells;
    vector<lower=0>[n_cells] sample_time_A;
    vector<lower=0>[n_cells] sample_time_B;
    int<lower=0> spike_count_A[n_cells];
    int<lower=0> spike_count_B[n_cells];
}

parameters {
    real<lower=0> rate_alpha;
    real<lower=0> rate_beta;
    vector<lower=0>[n_cells] rate_cell_A;
    
    real<lower=0> rate_ratio_mean;
    real<lower=0> rate_ratio_sd;
    vector<lower=0>[n_cells] rate_ratio_cell;
}

transformed parameters {
    real rate_A;
    real rate_B;
    real rate_change;
    
    vector[n_cells] rate_cell_B;
    vector[n_cells] rate_cell_change;
    
    rate_A = rate_alpha/rate_beta;
    rate_B = rate_A * rate_ratio_mean;
    rate_change = rate_B - rate_A;
    
    rate_cell_B = rate_cell_A .* rate_ratio_cell;
    rate_cell_change = rate_cell_B - rate_cell_A;
}

model {
    vector[n_cells] lambda_A;
    vector[n_cells] lambda_B;
    
    rate_alpha ~ gamma(0.5, 0.1);
    rate_beta ~ gamma(0.1, 0.1);
    
    rate_ratio_mean ~ normal(1, 1);
    rate_ratio_sd ~ normal(0, 1);
    
    rate_cell_A ~ gamma(rate_alpha, rate_beta);
    rate_ratio_cell ~ normal(rate_ratio_mean, rate_ratio_sd);
    
    spike_count_A ~ poisson(rate_cell_A .* sample_time_A);
    spike_count_B ~ poisson(rate_cell_A .* rate_ratio_cell .* sample_time_B);
}