sum_to_zero_vector case study

jupyter/radon/README.md

@@ -0,0 +1,6 @@
+Notebook based on Gelman and Hill 2007, Radon case study.
+Introduces hierarchical models and Python packages `cmdstanpy`, and `plotnine`.
+
+Author:  Mitzi Morris
+
+

jupyter/sum-to-zero/README.md

@@ -0,0 +1,9 @@
+Notebooks to demonstrate correctness and efficiency of constrained parameter type `sum\_to\_zero\_vector`,
+introduced in Stan 2.36.
+
+HTML page has all necessary resources.  To rebuild,  add `stan-dev/quarto-config` submodule.
+
+Author:  Mitzi Morris
+
+
+

jupyter/sum-to-zero/stan/binomial_4preds_hard.stan

@@ -0,0 +1,53 @@
+// multi-level model for binomial data with 4 categorical predictors.
+data {
+  int<lower=1> N; // number of strata
+  int<lower=1> N_age;
+  int<lower=1> N_eth;
+  int<lower=1> N_edu;
+
+  array[N] int<lower=0> pos_tests;
+  array[N] int<lower=0> tests;
+  array[N] int<lower=1, upper=2> sex;
+  array[N] int<lower=1, upper=N_age> age;
+  array[N] int<lower=1, upper=N_eth> eth;
+  array[N] int<lower=1, upper=N_edu> edu;
+
+  // hyperparameters
+  real<lower=0, upper=1> sens;
+  real<lower=0, upper=1> spec;
+}
+parameters {
+  real beta_0;
+  real beta_sex_raw;
+  real<lower=0> sigma_age, sigma_eth, sigma_edu;
+  vector[N_age - 1] beta_age_raw;
+  vector[N_eth - 1] beta_eth_raw;
+  vector[N_edu - 1] beta_edu_raw;
+}
+transformed parameters {
+  vector[2] beta_sex = [beta_sex_raw, -beta_sex_raw]';
+
+  vector[N_age] beta_age = append_row(beta_age_raw, -sum(beta_age_raw));
+  vector[N_eth] beta_eth = append_row(beta_eth_raw, -sum(beta_eth_raw));
+  vector[N_edu] beta_edu = append_row(beta_edu_raw, -sum(beta_edu_raw));
+
+  vector[N] eta =  inv_logit(beta_0 + beta_sex[sex] + beta_age[age] + beta_eth[eth] +  beta_edu[edu]);
+  vector[N] prob_pos_test = eta * sens + (1 - eta) * (1 - spec);
+}
+model {
+  pos_tests ~ binomial(tests, prob_pos_test);  // likelihood
+
+  // priors
+  beta_0 ~ normal(0, 2.5);
+  beta_sex ~ std_normal();
+  // centered parameterization
+  beta_age_raw ~ normal(0, sigma_age);
+  beta_eth_raw ~ normal(0, sigma_eth);
+  beta_edu_raw ~ normal(0, sigma_edu);
+  sigma_eth ~ std_normal();
+  sigma_age ~ std_normal();
+  sigma_edu ~ std_normal();
+}
+generated quantities {
+  array[N] int<lower=0>y_rep = binomial_rng(tests, prob_pos_test);
+}

jupyter/sum-to-zero/stan/binomial_4preds_hard_ppc.stan

@@ -0,0 +1,38 @@
+// multi-level model for binomial data with 4 categorical predictors.
+data {
+  int<lower=1> N; // number of strata
+  int<lower=1> N_age;
+  int<lower=1> N_eth;
+  int<lower=1> N_edu;
+
+  // hyperparameters
+  real<lower=0, upper=1> sens;
+  real<lower=0, upper=1> spec;
+}
+parameters {
+  real beta_0;
+  real beta_sex_raw;
+  real<lower=0> sigma_age, sigma_eth, sigma_edu;
+  vector[N_age - 1] beta_age_raw;
+  vector[N_eth - 1] beta_eth_raw;
+  vector[N_edu - 1] beta_edu_raw;
+}
+transformed parameters {
+  vector[2] beta_sex = [beta_sex_raw, -beta_sex_raw]';
+
+  vector[N_age] beta_age = append_row(beta_age_raw, -sum(beta_age_raw));
+  vector[N_eth] beta_eth = append_row(beta_eth_raw, -sum(beta_eth_raw));
+  vector[N_edu] beta_edu = append_row(beta_edu_raw, -sum(beta_edu_raw));
+}
+model {
+  // priors
+  beta_0 ~ normal(0, 2.5);
+  beta_sex ~ std_normal();
+  // centered parameterization
+  beta_age_raw ~ normal(0, sigma_age);
+  beta_eth_raw ~ normal(0, sigma_eth);
+  beta_edu_raw ~ normal(0, sigma_edu);
+  sigma_eth ~ std_normal();
+  sigma_age ~ std_normal();
+  sigma_edu ~ std_normal();
+}

jupyter/sum-to-zero/stan/binomial_4preds_ozs.stan

@@ -0,0 +1,62 @@
+// multi-level model for binomial data with 4 categorical predictors.
+data {
+  int<lower=1> N;  // number of strata
+  int<lower=1> N_age;
+  int<lower=1> N_eth;
+  int<lower=1> N_edu;
+
+  array[N] int<lower=0> pos_tests;
+  array[N] int<lower=0> tests;
+  array[N] int<lower=1, upper=2> sex;
+  array[N] int<lower=1, upper=N_age> age;
+  array[N] int<lower=1, upper=N_eth> eth;
+  array[N] int<lower=1, upper=N_edu> edu;
+
+  // hyperparameters
+  real<lower=0, upper=1> sens;
+  real<lower=0, upper=1> spec;
+}
+transformed data {
+  real mean_sex = mean(sex);
+  vector[N] sex_c = to_vector(sex) - mean_sex;
+  // scaling factors for marginal variances of sum_to_zero_vectors
+  // https://discourse.mc-stan.org/t/zero-sum-vector-and-normal-distribution/38296
+  real s_age = sqrt(N_age * inv(N_age - 1));
+  real s_eth = sqrt(N_eth * inv(N_eth - 1));
+  real s_edu = sqrt(N_edu * inv(N_edu - 1));
+}
+parameters {
+  real beta_0;
+  real beta_sex;
+  real<lower=0> sigma_age, sigma_eth, sigma_edu;
+  sum_to_zero_vector[N_age] beta_age;
+  sum_to_zero_vector[N_eth] beta_eth;
+  sum_to_zero_vector[N_edu] beta_edu;
+}
+transformed parameters {
+  // true prevalence
+  vector[N] p = inv_logit(beta_0 + beta_sex * sex_c + beta_age[age]
+			  + beta_eth[eth] + beta_edu[edu]);
+  // incorporate test sensitivity and specificity.
+  vector[N] p_sample = p * sens + (1 - p) * (1 - spec);
+}
+model {
+  pos_tests ~ binomial(tests, p_sample);  // likelihood
+
+  // priors
+  beta_0 ~ normal(0, 2.5);
+  beta_sex ~ std_normal();
+  sigma_eth ~ std_normal();
+  sigma_age ~ std_normal();
+  sigma_edu ~ std_normal();
+
+  // centered parameterization
+  // scale normal priors on sum_to_zero_vectors
+  beta_age ~ normal(0, s_age * sigma_age);
+  beta_eth ~ normal(0, s_eth * sigma_eth);
+  beta_edu ~ normal(0, s_edu * sigma_edu);
+}
+generated quantities {
+  real beta_intercept = beta_0 - mean_sex * beta_sex;
+  array[N] int<lower=0>y_rep = binomial_rng(tests, p_sample);
+}

jupyter/sum-to-zero/stan/binomial_4preds_ozs_ppc.stan

@@ -0,0 +1,38 @@
+// generate sample from model priors, (before seeing any data)
+data {
+  int<lower=1> N; // number of strata
+  int<lower=1> N_age;
+  int<lower=1> N_eth;
+  int<lower=1> N_edu;
+  // omit observational data
+}
+transformed data {
+  // scaling factors for marginal variances of sum_to_zero_vectors
+  // https://discourse.mc-stan.org/t/zero-sum-vector-and-normal-distribution/38296
+  real s_age = sqrt(N_age * inv(N_age - 1));
+  real s_eth = sqrt(N_eth * inv(N_eth - 1));
+  real s_edu = sqrt(N_edu * inv(N_edu - 1));
+}
+parameters {
+  real beta_0;
+  real beta_sex;
+  real<lower=0> sigma_age, sigma_eth, sigma_edu;
+  sum_to_zero_vector[N_age] beta_age;
+  sum_to_zero_vector[N_eth] beta_eth;
+  sum_to_zero_vector[N_edu] beta_edu;
+}
+model {
+  // omit likelihood
+  // priors
+  beta_0 ~ normal(0, 2.5);
+  beta_sex ~ std_normal();
+  sigma_eth ~ std_normal();
+  sigma_age ~ std_normal();
+  sigma_edu ~ std_normal();
+
+  // centered parameterization
+  // scale normal priors on sum_to_zero_vectors
+  beta_age ~ normal(0, s_age * sigma_age);
+  beta_eth ~ normal(0, s_eth * sigma_eth);
+  beta_edu ~ normal(0, s_edu * sigma_edu);
+}

jupyter/sum-to-zero/stan/binomial_4preds_soft.stan

@@ -0,0 +1,58 @@
+// multi-level model for binomial data with 4 categorical predictors.
+data {
+  int<lower=1> N; // number of strata
+  int<lower=1> N_age;
+  int<lower=1> N_eth;
+  int<lower=1> N_edu;
+
+  array[N] int<lower=0> pos_tests;
+  array[N] int<lower=0> tests;
+  array[N] int<lower=1, upper=2> sex;
+  array[N] int<lower=1, upper=N_age> age;
+  array[N] int<lower=1, upper=N_eth> eth;
+  array[N] int<lower=1, upper=N_edu> edu;
+
+  // hyperparameters
+  real<lower=0, upper=1> sens;
+  real<lower=0, upper=1> spec;
+}
+transformed data {
+  real mean_sex = mean(sex);
+  vector[N] sex_c = to_vector(sex) - mean_sex;
+}
+parameters {
+  real beta_0;
+  real beta_sex;
+  real<lower=0> sigma_age, sigma_eth, sigma_edu;
+  vector[N_age] beta_age;
+  vector[N_eth] beta_eth;
+  vector[N_edu] beta_edu;
+}
+transformed parameters {
+  // non-standard link function
+  vector[N] p =  inv_logit(beta_0 + beta_sex * sex_c + beta_age[age]
+			   + beta_eth[eth] +  beta_edu[edu]);
+  vector[N] p_sample = p * sens + (1 - p) * (1 - spec);
+}
+model {
+  pos_tests ~ binomial(tests, p_sample);  // likelihood
+
+  // priors
+  beta_0 ~ normal(0, 2.5);
+  beta_sex ~ std_normal();
+  // centered parameterization
+  beta_age ~ normal(0, sigma_age);
+  beta_eth ~ normal(0, sigma_eth);
+  beta_edu ~ normal(0, sigma_edu);
+  sigma_eth ~ std_normal();
+  sigma_age ~ std_normal();
+  sigma_edu ~ std_normal();
+  // soft sum-to-zero constraint
+  sum(beta_age) ~ normal(0, 0.001 * N_age); 
+  sum(beta_eth) ~ normal(0, 0.001 * N_eth);
+  sum(beta_edu) ~ normal(0, 0.001 * N_edu);
+}
+generated quantities {
+  real beta_intercept = beta_0 - mean_sex * beta_sex;
+  array[N] int<lower=0>y_rep = binomial_rng(tests, p_sample);
+}

jupyter/sum-to-zero/stan/bym2_hard.stan

@@ -0,0 +1,63 @@
+ data {
+  int<lower=0> N;
+  array[N] int<lower=0> y; // count outcomes
+  vector<lower=0>[N] E; // exposure
+  int<lower=1> K; // num covariates
+  matrix[N, K] xs; // design matrix
+
+  // spatial structure
+  int<lower = 0> N_edges;  // number of neighbor pairs
+  array[2, N_edges] int<lower = 1, upper = N> neighbors;  // columnwise adjacent
+
+  real tau; // scaling factor
+}
+transformed data {
+  vector[N] log_E = log(E);
+  // center continuous predictors 
+  vector[K] means_xs;  // column means of xs before centering
+  matrix[N, K] xs_centered;  // centered version of xs
+  for (k in 1:K) {
+    means_xs[k] = mean(xs[, k]);
+    xs_centered[, k] = xs[, k] - means_xs[k];
+  }
+}
+parameters {
+  real beta0; // intercept
+  vector[K] betas; // covariates
+  real<lower=0, upper=1> rho; // proportion of spatial variance
+  vector[N-1] phi_raw; // spatial random effects
+  vector[N] theta; // heterogeneous random effects
+  real<lower = 0> sigma;  // scale of combined effects
+}
+transformed parameters {
+  vector[N] phi = append_row(phi_raw, -sum(phi_raw));
+  vector[N] gamma = (sqrt(1 - rho) * theta + sqrt(rho / tau) * phi);  // BYM2
+}
+model {
+  y ~ poisson_log(log_E + beta0 + xs_centered * betas + gamma * sigma);
+  rho ~ beta(0.5, 0.5);
+  target += -0.5 * dot_self(phi[neighbors[1]] - phi[neighbors[2]]); // ICAR
+  beta0 ~ std_normal();
+  betas ~ std_normal();
+  theta ~ std_normal();
+  sigma ~ std_normal();
+}
+generated quantities {
+  real beta_intercept = beta0 - dot_product(means_xs, betas);  // adjust intercept
+  array[N] int y_rep;
+  {
+    vector[N] eta = log_E + beta0 + xs_centered * betas + gamma * sigma;
+    if (max(eta) > 26) {
+      // avoid overflow in poisson_log_rng
+      print("max eta too big: ", max(eta));
+      for (n in 1:N) {
+        y_rep[n] = -1;
+        // log_lik[n] = -1;
+      }
+    } else {
+      for (n in 1:N) {
+        y_rep[n] = poisson_log_rng(eta[n]);
+      }
+    }
+  }
+}