chronic_array_hierarchical_regression.stan

data {
  // # of data points
  int<lower=1> N;
  // Observed x
  vector[N] x_obs;
  // Observed y
  vector[N] y_obs;

  // Estimate of measurement error on x_obs
  vector<lower=0>[N] err_x_obs;
  // Estimate of measurement error on y_obs
  vector<lower=0>[N] err_y_obs;

  // number of sessions. sess_id must contain all of 1:num_sess
  int<lower=1> num_sess;
  // number of electrodes. elec_id must contain all of 1:num_elec
  int<lower=1> num_elec;

  // session ID of each measurement
  int<lower=1, upper=num_sess> sess_id[N];
  // electrode ID of each measurement
  int<lower=1, upper=num_elec> elec_id[N];

  // Hyper-prior configuration. sigma_sigma means ``scale of hyper-prior
  // over sigma parameter'', i.e. sigma ~ cauchy(0, sigma_sigma)
  real<lower=0> sigma_sigma_dx_elec;
  real<lower=0> sigma_sigma_dy_elec;
  real<lower=0> sigma_sigma_dm_elec;
  real<lower=0> sigma_sigma_dx_sess;
  real<lower=0> sigma_sigma_dy_sess;
  real<lower=0> sigma_sigma_dm_sess;
  real<lower=0> sigma_sigma_dx_priv;
  real<lower=0> sigma_sigma_dy_priv;
}

parameters {
  // The main regression parameter (slope) we're interested in.
  real beta;

  // Infer the overall scales of various sources of variance.
  // variation in true d' (away from x0) across the population
  real<lower=0> sigma_dx_elec;
  // defines the range of delta CP across electrodes
  real<lower=0> sigma_dy_elec;
  // defines the range of delta slope across electrodes
  real<lower=0> sigma_dm_elec;
  // defines the range of delta d' across sessions
  real<lower=0> sigma_dx_sess;
  // defines the range of delta CP across sessions
  real<lower=0> sigma_dy_sess;
  // defines the range of delta slope across sessions
  real<lower=0> sigma_dm_sess;
  // defines the range of delta d' private to each measurement
  real<lower=0> sigma_dx_priv;
  // defines the range of delta CP private to each measurement
  real<lower=0> sigma_dy_priv;

  // Overall center offset in x and y
  real x0;
  real y0;

  // Infer 'true' x, y, and slope for each electrode. Prefix z_ means 
  // 'z-transformed' -- see 'reparameterization' note in transformed
  // parameters block.
  // ``baseline'' d' value per electrode is x0 plus this value
  vector[num_elec] z_delta_x_elec;
  // actual delta CP away from the regression line per electrode
  vector[num_elec] z_delta_y_elec;
  // actual delta slope per electrode
  vector[num_elec] z_delta_m_elec;

  // Infer 'true' x, y, and slope for each session.
  // actual delta d' per session
  vector[num_sess] z_delta_x_sess;
  // actual delta CP per session away from line
  vector[num_sess] z_delta_y_sess;
  // actual delta slope per session
  vector[num_sess] z_delta_m_sess;

  // Infer additional 'private' deltas per measurement.
  vector[N] z_delta_x_priv;
  vector[N] z_delta_y_priv;
}

transformed parameters {
  // Reparameterization to link hierarchical 'deltas' to their 'sigmas'.
  // This reduces the dependence between the '_z' prefixed variables
  // and the 'sigma_' variables in the posterior. See the STAN docs on
  // reparameterization here:
  // https://mc-stan.org/docs/2_18/stan-users-guide/reparameterization-section.html
  vector[num_elec] delta_x_elec = z_delta_x_elec * sigma_dx_elec;
  vector[num_elec] delta_y_elec = z_delta_y_elec * sigma_dy_elec;
  vector[num_elec] delta_m_elec = z_delta_m_elec * sigma_dm_elec;
  vector[num_sess] delta_x_sess = z_delta_x_sess * sigma_dx_sess;
  vector[num_sess] delta_y_sess = z_delta_y_sess * sigma_dy_sess;
  vector[num_sess] delta_m_sess = z_delta_m_sess * sigma_dm_sess;
  vector[N] delta_x_priv = z_delta_x_priv * sigma_dx_priv;
  vector[N] delta_y_priv = z_delta_y_priv * sigma_dy_priv;


  // Define 'measurement means' in x and y, which include all
  // structure up to per-observation measurement error
  vector[N] mu_x_obs;
  vector[N] mu_y_obs;
  for (i in 1:N) {
    // Construct x coordinate as the baseline x0 plus multiple 'delta's
    mu_x_obs[i] = x0 + delta_x_elec[elec_id[i]] + 
      delta_x_sess[sess_id[i]] + delta_x_priv[i];
    
    // Construct y from line m*x+b, plus additional deltas
    real net_intercept = y0 + delta_y_elec[elec_id[i]] + 
      delta_y_sess[sess_id[i]] + delta_y_priv[i];
    real net_slope = beta + delta_m_elec[elec_id[i]] + 
      delta_m_sess[sess_id[i]];
    mu_y_obs[i] = net_intercept + net_slope * (mu_x_obs[i] - x0);
  }
}

model {
  // Cut-cauchy priors on all standard deviations so we can infer the
  // overall amount of each type of variation.
  sigma_dx_elec ~ cauchy(0, sigma_sigma_dx_elec);
  sigma_dy_elec ~ cauchy(0, sigma_sigma_dy_elec);
  sigma_dm_elec ~ cauchy(0, sigma_sigma_dm_elec);
  sigma_dx_sess ~ cauchy(0, sigma_sigma_dx_sess);
  sigma_dy_sess ~ cauchy(0, sigma_sigma_dy_sess);
  sigma_dm_sess ~ cauchy(0, sigma_sigma_dm_sess);
  sigma_dx_priv ~ cauchy(0, sigma_sigma_dx_priv);
  sigma_dy_priv ~ cauchy(0, sigma_sigma_dy_priv);

  // Generate all 'deltas' independently from Gaussians. Thanks to
  // reparameterization, these are all normal(0, 1).
  z_delta_x_elec ~ normal(0, 1);
  z_delta_y_elec ~ normal(0, 1);
  z_delta_m_elec ~ normal(0, 1);
  z_delta_x_sess ~ normal(0, 1);
  z_delta_y_sess ~ normal(0, 1);
  z_delta_m_sess ~ normal(0, 1);
  z_delta_x_priv ~ normal(0, 1);
  z_delta_y_priv ~ normal(0, 1);

  // All the magic of the actuall regression model happens
  // inside 'transformed parameters' where mu_x_obs and mu_y_obs are
  // calculated. Here, we simply state that the observed data are
  // normally distributed around these highly- structured means with 
  // some residual observation noise of a roughly-known magnitude
  x_obs ~ normal(mu_x_obs, err_x_obs);
  y_obs ~ normal(mu_y_obs, err_y_obs); 
}