Parameterized text macros in Stan

Once of the challenges I’ve encountered when using Stan is that many common motifs (e.g., non-central parameterizations, horseshoe priors, splines) require coordinated code that belongs in multiple blocks. While #include statements can help with this somewhat, they have certain limitations: the defined variable names are fixed and the they require replacing the entire line.

I’ve thrown together an R package that addresses these limitations (tentatively titled macroStan).

The basic idea is to define the macro in R, write a Stan file with special notation that indicates where the various components belong, and then use an R function to combine them into a valid Stan model. Additional details and examples are given the the readme on Github.

I’m hoping some people here may find this useful and would welcome any suggestions for changes in syntax, approach, etc.

5 Likes

I’ve rewritten the macro interface to be substantially easier to use. Macros are now written as stan files with an extra macro args program block, and they’re used in other stan files with a use macros block.

For example, a macro for the regularized horseshoe would look like:

macro args {
  N_local = N_group;
  value = beta_hs;
}
functions {
  // horseshoe computation
  vector horseshoe(vector zb, vector[] local, real[] global,
                 real scale_global, real c2) {
    int K = rows(zb);
    vector[K] lambda = local[1] .* sqrt(local[2]);
    vector[K] lambda2 = square(lambda);
    real tau = global[1] * sqrt(global[2]) * scale_global;
    vector[K] lambda_tilde = sqrt(c2 * lambda2 ./ (c2 + tau^2 * lambda2));
    return zb .* lambda_tilde * tau;
  }
}
data {
  // data for horseshoe prior
  real<lower=0> hs_df_{|value|};
  real<lower=0> hs_df_global_{|value|};  // global degrees of freedom
  real<lower=0> hs_df_slab_{|value|};  // slab degrees of freedom
  real<lower=0> hs_scale_global_{|value|};  // global prior scale
  real<lower=0> hs_scale_slab_{|value|};  // slab prior scale",
  }
parameters{
  // horseshoe shrinkage parameters, global
  real<lower=0> hs_global_{|value|}[2];  // global shrinkage parameters
  real<lower=0> hs_c2_{|value|};  // slab regularization parameter
  // local parameters for horseshoe
  vector[{|N_local|}] hs_z_{|value|};
  vector<lower=0>[{|N_local|}] hs_local_{|value|}[2];

}
transformed parameters{
 // horseshoe regression coefs
  vector[{|N_local|}] {|value|} =
    horseshoe(hs_z_{|value|}, hs_local_{|value|},
    hs_global_{|value|}, hs_scale_global_{|value|},
    hs_scale_slab_{|value|}^2 * hs_c2_{|value|}  );
}
model{
 // horseshoe prior, global
  target += std_normal_lpdf(hs_global_{|value|}[1]) - 1 * log(0.5) +
            inv_gamma_lpdf(hs_global_{|value|}[2] |
              0.5 * hs_df_global_{|value|},0.5 * hs_df_global_{|value|} ) +
            inv_gamma_lpdf(hs_c2_{|value|} |
              0.5 * hs_df_slab_{|value|},0.5 * hs_df_slab_{|value|} );
  //horseshoe prior, local
  target += std_normal_lpdf(hs_z_{|value|}) +
            std_normal_lpdf(hs_local_{|value|}[1]) -  {|N_local|} * log(0.5) +
            inv_gamma_lpdf(hs_local_{|value|}[2] |
               0.5 * hs_df_{|value|}, 0.5 * hs_df_{|value|});
}

It could be used by a stan model like this:

// Horseshoe prior example
use macros {
  beta = horseshoe(D);
}
data {
  int N;
  int D;
  matrix[N,D] x;
  vector[N] y;
}
parameters {
  real<lower=0> sigma;
  real alpha;
}
model {
  vector[N] eta = alpha + x * beta ;
  target += normal_lpdf(y | eta, sigma);
}

And it would get parsed into this:

functions { 
// horseshoe computation
  vector horseshoe(vector zb, vector[] local, real[] global,
                 real scale_global, real c2) {
    int K = rows(zb);
    vector[K] lambda = local[1] .* sqrt(local[2]);
    vector[K] lambda2 = square(lambda);
    real tau = global[1] * sqrt(global[2]) * scale_global;
    vector[K] lambda_tilde = sqrt(c2 * lambda2 ./ (c2 + tau^2 * lambda2));
    return zb .* lambda_tilde * tau;
  } 
} 
data { 
int N;
  int D;
  matrix[N,D] x;
  vector[N] y;

// data for horseshoe prior
real<lower=0> hs_df_beta;
real<lower=0> hs_df_global_beta;  // global degrees of freedom
real<lower=0> hs_df_slab_beta;  // slab degrees of freedom
real<lower=0> hs_scale_global_beta;  // global prior scale
real<lower=0> hs_scale_slab_beta;  // slab prior scale", 
} 
parameters { 
real<lower=0> sigma;
  real alpha;

  // horseshoe shrinkage parameters, global
  real<lower=0> hs_global_beta[2];  // global shrinkage parameters
  real<lower=0> hs_c2_beta;  // slab regularization parameter
  // local parameters for horseshoe
  vector[D] hs_z_beta;
  vector<lower=0>[D] hs_local_beta[2]; 
} 
transformed parameters { 
// horseshoe regression coefs
 vector[D] beta =
   horseshoe(hs_z_beta, hs_local_beta,
   hs_global_beta, hs_scale_global_beta,
   hs_scale_slab_beta^2 * hs_c2_beta  ); 
} 
model { 
vector[N] eta = alpha + x * beta ;

 // horseshoe prior, global
  target += std_normal_lpdf(hs_global_beta[1]) - 1 * log(0.5) +
            inv_gamma_lpdf(hs_global_beta[2] |
              0.5 * hs_df_global_beta,0.5 * hs_df_global_beta ) +
            inv_gamma_lpdf(hs_c2_beta |
              0.5 * hs_df_slab_beta,0.5 * hs_df_slab_beta );
  //horseshoe prior, local
  target += std_normal_lpdf(hs_z_beta) +
            std_normal_lpdf(hs_local_beta[1]) -  D * log(0.5) +
            inv_gamma_lpdf(hs_local_beta[2] |
               0.5 * hs_df_beta, 0.5 * hs_df_beta);

  target += normal_lpdf(y | eta, sigma); 
} 

More details and examples are in the GitHub repo. I’m planning to slowly more macros as I need them for my own analyses.

4 Likes

I wanted to bump this because I think it’s a really good idea. Would love to see more active development on it.

2 Likes

yeah, this looks great!
Is this still in active development?

I think there was a lot of talk on how to add more modularization to Stan and it was hard for people to agree which is the best way… See Stan++/Stan3 Preliminary Design (there might have been some followups, but I don’t think there was a broad agreement)

I’m glad to see some interest in the project! If anyone uses this and wants to share their macros, I’d be happy to add or link to them in the repository.

I haven’t really done much on the project since my latest post, and my near-future plans are focused almost entirely on trying to finish my dissertation. That being said, I’d love to hear it if anyone has any interesting thoughts/ideas for future directions.

1 Like