diff --git a/inst/stan/bmsma.stan b/inst/stan/bmsma.stan new file mode 100644 index 0000000..0c6a661 --- /dev/null +++ b/inst/stan/bmsma.stan @@ -0,0 +1,74 @@ +// A multi-level sma model fit +data { + int n_obs; + int n_groups; + int group[n_obs]; + vector[2] x[n_obs]; + } + +parameters { + // hyper parameters + real mu_mu_x; + real sigma_mu_x; + real mu_b_0; + real sigma_b_0; + real mu_log_b_1; + real sigma_log_b_1; + real a_sigma_u1; + real b_sigma_u1; + real a_sigma_u2; + real b_sigma_u2; + + // group-level effects + real mu_x1[n_groups]; + real b_0[n_groups]; + real b_1[n_groups]; + real sigma_u1[n_groups]; + real sigma_u2[n_groups]; +} + +model { + // Define variables used + real u_c; // constant + vector[2] uv[n_obs]; // (u1,u2) data in vector form + + matrix[2,2] U[n_groups]; // Rotation matrices + vector[2] mu_x[n_groups]; // vector means for x + vector[2] sigma_u[n_groups]; // covariance matrices for u + + // Sample group parameters from distributions + mu_x1 ~ normal(mu_mu_x, sigma_mu_x); + b_0 ~ normal(mu_b_0, sigma_b_0); + b_1 ~ lognormal(mu_log_b_1, sigma_log_b_1); + sigma_u1 ~ normal(a_sigma_u1, b_sigma_u1); + sigma_u2 ~ normal(a_sigma_u2, b_sigma_u2); + // NB: sigmas should be sampled from gamma (or inverse gamma), + // but that doesn't converge. Gelman says gamma performs badly + // when variance is close to zero + + // Now calculate vectors and matrices from proposed parameters, indexed by group + for (i in 1:n_groups) { + mu_x[i, 1] = mu_x1[i]; + mu_x[i, 2] = b_0[i] + b_1[i] * mu_x1[i]; + + // Rotation matrix, U + u_c = 1 / sqrt(2) / b_1[i]; + U[i,1,1] = b_1[i]^2 * u_c; + U[i,1,2] = b_1[i] * u_c; + U[i,2,1] = -b_1[i] * u_c; + U[i,2,2] = 1 * u_c; + + // Covariance matrix of U + sigma_u[i,1] = sigma_u1[i]; + sigma_u[i,2] = sigma_u2[i]; + } + + for (j in 1:n_obs) { + // Calculate u by centering then rotating x + uv[j] = U[group[j]] * (log(x[j]) - mu_x[group[j]]); + + // Likelihood of data + // consider re-parameterising as described on pg 218 + uv[j] ~ multi_normal(rep_vector(0,2), diag_matrix(sigma_u[group[j]])); + } +}