model with an arbitrary n of partially pooled parameters

[wwiecek]

Flexible extension of mu & tau into P dimensions. Below is an example. This would need more work and the main problem is that you can't have full pooling on some parameters and non-full on others.

functions {
#include /functions/prior_increment.stan
}

data {
  int<lower=0> P;  //number of parameters -- CURRENTLY IT HAS TO BE 2, more will fail
  int<lower=0> N;  // total number of observations
  int<lower=0> K;  // number of sites
  int<lower=0> Nc; //number of covariates (fixed effects)
  matrix[N,Nc] X;  //covariate values (design matrix for FE)
  int pooling_type[P]; //0 if none, 1 if partial, 2 if full
  int joint_prior;     //0 if no, 1 if yes, but 1 requires all pooling == 1
  int<lower=0,upper=K> site[N];
  // vector<lower=0,upper=1>[N] treatment;
  matrix[N, P] U;
  
  // priors:
  
  // joint prior for the mean
  int prior_hypermean_fam;
  matrix[P, 3] prior_hypermean_val;
  vector[P] prior_hypermean_mean;
  matrix<lower=0>[P, P] prior_hypermean_scale;
  
  //dispersions for effect and baselines:
  int prior_hypersd_fam;
  matrix[P, 3] prior_hypersd_val;
  
  //correlations (only LKJ allowed for now)
  int prior_hypercor_fam; 
  real prior_hypercor_val;
  
  //fixed effects (same dist. for all betas)
  int prior_beta_fam;
  matrix[Nc, 3] prior_beta_val;
  
    //cross-validation variables:
  int<lower=0> N_test;
  int<lower=0> K_test;
  matrix[N_test, Nc] X_test;
  int<lower=0, upper=K> test_site[N_test];
  int<lower=0, upper=1> test_treatment[N_test];

  // NORMAL specific:
  vector[N] y;
  vector[N_test] test_y;
  vector[K_test] test_sigma_y_k;
}

transformed data {
  int K_pooled; // number of modelled sites if we take pooling into account
  if(pooling_type == 2)
    K_pooled = 0;
  if(pooling_type != 2)
    K_pooled = K;
}

parameters {
  // corr_matrix[P] Omega;        // prior correlation
  cholesky_factor_corr[P] L_Omega[joint_prior? 1: 0]; 
  vector<lower=0>[P] hypersd[pooling_type == 1? 1: 0];        // prior scale
  vector[P] mu[pooling_type != 0? 1: 0];      
  matrix[P,K] eta[pooling_type != 2? 1: 0];
  vector[Nc] beta;
  // NORMAL specific:
  vector<lower=0>[K] sigma_y_k;
}
transformed parameters {
  matrix[P,K] theta_k[pooling_type != 2? 1: 0];
  matrix[P,P] tau[pooling_type == 1? 1: 0];
  if(pooling_type == 0)
    theta_k[1] = eta[1];
  if(pooling_type == 1){
    if(joint_prior)
      tau[1] = diag_pre_multiply(hypersd[1], L_Omega[1]);
    else
      tau[1] = diag_matrix(hypersd[1]);
    theta_k[1] = rep_matrix(mu[1], K) + tau[1] * eta[1];
  }
}
model {
  vector[N] y_mean;
  if(N > 0){
    if(Nc == 0)
      y_mean = rep_vector(0.0, N);
    else{
      y_mean = X*beta;
      for(p in 1:Nc)
        target += prior_increment_vec(prior_beta_fam, beta, to_vector(prior_beta_val[p]));
    }
    if(pooling_type != 2){
      for(p in 1:P)
        y_mean += U[, p] .* to_vector(theta_k[1][p, site]);
    }else{
      for(p in 1:P)
        y_mean += U[, p] * mu[1][p];
    }
  }
  
  if(joint_prior){
    mu[1] ~ multi_normal(prior_hypermean_mean, prior_hypermean_scale);
    L_Omega[1] ~ lkj_corr_cholesky(prior_hypercor_val);
  }else {
    for(p in 1:P){
      if(pooling_type > 0)
        target += prior_increment_real(prior_hypermean_fam, mu[1][p], to_vector(prior_hypermean_val[p]));
      else
        target += prior_increment_vec(prior_hypermean_fam, to_vector(eta[1][p,]), to_vector(prior_hypermean_val[p]));
      if(pooling_type == 1) {
        to_vector(eta[1][p,]) ~ std_normal();
        target += prior_increment_real(prior_hypersd_fam, hypersd[1][p], to_vector(prior_hypersd_val[p]));
      }
    }
  }
  
  
  y ~ normal(y_mean, sigma_y_k[site]);
}

generated quantities {
  // to do this, we must first (outside of Stan) calculate SEs in each test group,
  // i.e. test_sigma_y_k`
}