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.

2 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.

2 Likes