How to speed up my Stan code?

Hi Stan community,
I’m new to Stan. The following paragraph is my stan code for the hierarchical survival model.
I used this code for large data (N=35000).
To run, It took ~10 days to even for 2000 iteration/500 warmup.
I really want to know how can I speed up my code. So would be greatly appreciated if there are any tips!

Thank you.

My stan code is as below (N=35000, n_grid=350, province=31, num=12).

fit ← stan(“final model.stan”,data = da,iter=iter, warmup=warmup, cores =4,chains = 2, thin = 1, init = “random”,control=list(max_treedepth=15))

data {
  int <lower=0> province;
  int <lower=0> n_grid;
  int <lower=0> N;
  int <lower=0> num;
  int <lower=0> number;
  vector <lower=0> [province] p;
  matrix [province,n_grid] kernel;
  matrix [N,num] predictor;
  matrix [N,number] respon;
}
parameters{
  vector [num] beta;
  real <lower=0> sigma0;
  real <lower=0> sigma;
  vector [n_grid] x;
}
transformed parameters{
  real alpha = 1 / sigma;
  vector [N] lambdaa;
  vector[province] a; // `a` saved to output
     lambdaa = exp ((-1 * predictor * beta) / sigma);
           { // not part of output
    vector[province] z;
for (k in 1 : province) {
     vector[n_grid] landa_k;
for (j in 1 : n_grid) {
     landa_k[j] = kernel[k, j] * exp(x[j]);
           }
     z[k] = sum(landa_k) * p[k];
           }
     a = z / sum(z); // assign `a`
           }
           }
model{
      matrix [N,province] f;
      matrix [N,province] s;
     target += normal_lpdf(x|0,1);
    target += normal_lpdf(beta | 0, sigma0);
    target += inv_gamma_lpdf(sigma0| 0.001, 0.001);
    target += normal_lpdf(log (alpha) | 0, 5);
for (i in 1:N) {
for (k in 1:province){
    f[i,k]= a[k]*((lambdaa [i]*alpha*respon[i,1]^(alpha-1))/((1+(lambdaa [i]*respon[i,1]^alpha))^2));
    s[i,k]= a[k]*(1- ((lambdaa [i]*respon[i,1]^alpha)/ (1+(lambdaa [i]*respon[i,1]^alpha))));
    target += (respon[i,2] *  log (f[i,k])) + ((1 - respon[i,2])* log(s[i, k]));
    }
  }
}

Afternoon!

Could you provide the version of rstan, R, and your OS? Can you also provide a snippet of your data or some simulated data.

thanks

I think you can calculate z with a matrix multiply and elementwise multiplication on p

data {
  int<lower=0> province;
  int<lower=0> n_grid;
  int<lower=0> N;
  int<lower=0> num;
  int<lower=0> number;
  vector<lower=0> [province] p;
  matrix[province,n_grid] kernel;
  matrix[N,num] predictor;
  matrix[N,number] respon;
}
parameters {
  vector[num] beta;
  real<lower=0> sigma0;
  real<lower=0> sigma;
  vector[n_grid] x;
}
transformed parameters{
  real alpha = 1 / sigma;
  vector[N] lambda;
  vector[province] a; // `a` saved to output
  lambda = exp((-1 * predictor * beta) / sigma);
  { // not part of output
    vector[province] z;
    /**
    for these sizes
    //   matrix[province,n_grid] kernel;
    // vector[n_grid] x;
    I think you can just do a mat mul since
    the result dimensions are by province 
    // result[province] = kernel[province, n_grid] * x[n_grid]
    for (k in 1:province) {
      vector[n_grid] lambda_k;
      for (j in 1:n_grid) {
        lambda_k[j] = kernel[k, j] * exp(x[j]);
      }
    // then this just becomes an elementwise multiply
      z[k] = sum(lambda_k) * p[k];
    }
    
    */
    z = (kernel * exp(x)) .* p[k];
   // is this supposed to be a simplex?
    a = z / sum(z); // assign `a`
  }
}

model {
  matrix[N,province] f;
  matrix[N,province] s;
  target += normal_lpdf(x| 0, 1);
  target += normal_lpdf(beta | 0, sigma0);
  target += inv_gamma_lpdf(sigma0| 0.001, 0.001);
  target += normal_lpdf(log(alpha) | 0, 5);
  for (i in 1:N) {
    for (k in 1:province) {
      f[i, k]= a[k] * ((lambda[i] * alpha * respon[i, 1]^(alpha - 1)) / ((1 + (lambda[i] * respon[i, 1]^alpha))^2));
      s[i, k]= a[k] * (1 - ((lambda[i] * respon[i, 1]^alpha) / (1 + (lambda[i] * respon[i, 1]^alpha))));
      target += (respon[i, 2] * log(f[i, k])) + ((1 - respon[i, 2]) * log(s[i, k]));
    }
  }
}

In the future I’d recommend cleaning up spacing and indentation so it’s a little simpler to follow in the future. Do you have a paper associated with the model?

3 Likes

you can also pull out and vectorize operations like the above that are done multiple times

2 Likes

I really appreciated your information and help. I’m working with R 4.0.3 and the last version of RSTAN.
Actually, I made a lot of changes to the previous code and the latest version of my code is as below. Please check and let me know if you have any suggestions for improving its speed. I will be very grateful.

fit < - stan(“final model.stan”,data = da,iter=iter, warmup=warmup, cores =4,chains = 4, thin = 1, init = “random”)

// My Stan codes for hierarchical survival model for large dataset(N=35000, n_grid=354, province=31, num=12).

data {
  int <lower=0> province;
  int <lower=0> n_grid;
  int <lower=0> N;
  int <lower=0> num;
  int <lower=0> number;
  int county[N];
  vector <lower=0> [province] p;
  //row_vector[n_grid] kernel[province];
  matrix [province,n_grid] kernel;
  matrix [N,number] respon;
}
parameters{
  vector [num] beta;
  real <lower=0> sigma;
  vector [n_grid] x;
}
transformed parameters{
  vector [num] time_ratio;
  vector [N] lambdaa;
  real alpha = 1 / sigma;
  vector [n_grid] exp_x;
  vector[province] a; // `a` saved to output
     time_ratio  = exp(beta);
  exp_x = exp(x);
     lambdaa = exp ((-1 * respon[,6:17] * beta) / sigma);
{ // `z` not part of output
    vector[province] z;
    z = (kernel * exp_x) .* p;
    a = z / sum(z);
   // vector[province] landa;
    //for (k in 1 : province) {
      //  landa[k] = kernel[k] * exp_x * p[k];
   // }
   //a = landa / sum(landa); // assign `a`
  }
  }
model{
  vector [N] f;
  vector [N] s;
  real log_alpha;
  log_alpha = log (alpha);
    target += normal_lpdf(x|0,5);
    target += normal_lpdf(beta | 0, pow(10,5));
    target += normal_lpdf(log_alpha | 0, 5);
for (i in 1:N) {
    f[i]= a[county[i]]*((lambdaa [i]*alpha*respon[i,4]^(alpha-1))/((1+(lambdaa [i]*respon[i,4]^alpha))^2));
    s[i]= a[county[i]]*(1- ((lambdaa [i]*respon[i,4]^alpha)/ (1+(lambdaa [i]*respon[i,4]^alpha))));
    target += respon[i,5] *  (log (f[i])-log(s[i])) +  log(s[i]) - log_alpha;
  }
}

if f and s are being logged you can propogate that log through the model to make the multiplications and exponentiations into additions and multiplication

1 Like

Thank you for your help and reply. Now, I modified my code as below according to your instructions. I would greatly appreciate it if you could let me know when you have checked my code.

data {
  int <lower=0> province;
  int <lower=0> n_grid;
  int <lower=0> N;
  int <lower=0> num;
  int <lower=0> number;
  int county[N];
  vector <lower=0> [province] p;
  //row_vector[n_grid] kernel[province];
  matrix [province,n_grid] kernel;
  matrix [N,number] respon;
}
parameters{
  vector [num] beta;
  real <lower=0> sigma;
  vector [n_grid] x;
}
transformed parameters{
  vector [num] time_ratio;
  vector [N] lambdaa;
  real alpha = 1 / sigma;
  vector [n_grid] exp_x;
  vector[province] a; // `a` saved to output
     time_ratio  = exp(beta);
  exp_x = exp(x);
     lambdaa = exp ((-1 * respon[,6:17] * beta) / sigma);
{ // `z` not part of output
    vector[province] z;
    z = (kernel * exp_x) .* p;
    a = z / sum(z);
   // vector[province] landa;
    //for (k in 1 : province) {
      //  landa[k] = kernel[k] * exp_x * p[k];
   // }
   //a = landa / sum(landa); // assign `a`
  }
  }
model{
  // vector [N] f;
    vector [N] log_f;
  //vector [N] s;
   vector [N] log_s;
  real log_alpha;
  log_alpha = log (alpha);
    target += normal_lpdf(x|0,5);
    target += normal_lpdf(beta | 0, pow(10,5));
    target += normal_lpdf(log_alpha | 0, 5);
for (i in 1:N) {
 //  f[i]= a[county[i]]*((lambdaa [i]*alpha*respon[i,4]^(alpha-1))/((1+(lambdaa 
    [i]*respon[i,4]^alpha))^2));
     log_f[i]= log(a[county[i]])+log(lambdaa[i])+log_alpha+alpha*log(respon[i,4])- 
     log(respon[i,4])- 2 * log(1+(lambdaa[i]*respon[i,4]^alpha));

// s[i]= a[county[i]]*(1- ((lambdaa [i]*respon[i,4]^alpha)/ (1+(lambdaa [i]*respon[i,4]^alpha))));
     log_s[i]=log(a[county[i]])+log(1- ((lambdaa [i]*respon[i,4]^alpha)/ (1+(lambdaa 
     [i]*respon[i,4]^alpha))));
   // target += respon[i,5] *  (log (f[i])-log(s[i])) +  log(s[i]) - log_alpha;
  }
     target += respon[,5] .*  (log_f - log_s) +  log_s - log_alpha;
}

I can’t check your code because I don’t have data associated with your model. I recommend following the SBC approach and building code for your data generating process and using that while you optimize the model to make sure you are getting the same results

2 Likes

Ok. Thanks a lot for your information and help.