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 andT[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);
}