Constrained multinormal

Modeling
1

The toy example below is an attempt to do a constrained version of b ~ multi_normal(mu, quad_form_diag(Omega, tau)), where the off diagonals of Omega are set to rho.

data {
  vector[2] mu;
  vector[2] tau;
  vector[2] lb;
  vector[2] ub;
  real rho;
}

parameters {
  real g_std;      // group effect
  vector[2] b_std; // individual effects
}

transformed parameters {
  real<lower=lb[1], upper=ub[1]> b1 = mu[1] + tau[1] * (sqrt(1-rho) * b_std[1] + sqrt(rho) * g_std);
  real<lower=lb[2], upper=ub[2]> b2 = mu[2] + tau[2] * (sqrt(1-rho) * b_std[2] + sqrt(rho) * g_std);
}

model {
  g_std ~ std_normal();
  b_std ~ std_normal();
}

It does seem produce reasonable distributions. However, I get a lot of “divergent transitions after warmup” messages. Is there better way to keep individual parameters within constraints when there exist a group effect?

image

2
3

Thanks @bgoodri, that works nicely. Based on your paper, I added the case with both lower and upper bounds

functions {
  vector[] make_stuff(vector mu, matrix L, vector lb, vector ub, vector u) {
    int K = rows(mu); vector[K] d; vector[K] z; vector[K] out[2];
    for (k in 1:K) {
      int km1 = k - 1;
      
      if (lb[k] == negative_infinity() && ub[k] == positive_infinity()) { // no bounds
        z[k] = inv_Phi(u[k]);
        d[k] = 1;
      } else {
        real v; 
        
        if (lb[k] != negative_infinity() && ub[k] != positive_infinity()) { // lower & upper bounds
          real z_star_lb = (lb[k] - (mu[k] + ((k > 1) ? L[k,1:km1] * head(z, km1) : 0))) / L[k,k];
          real z_star_ub = (ub[k] - (mu[k] + ((k > 1) ? L[k,1:km1] * head(z, km1) : 0))) / L[k,k];
          real u_star_lb = Phi(z_star_lb);
          real u_star_ub = Phi(z_star_ub);
          v = u_star_lb + (u_star_ub - u_star_lb) * u[k];
          d[k] = u_star_ub - u_star_lb;          
        } else if (ub[k] != positive_infinity()) { // upper bound
          real z_star = (ub[k] - (mu[k] + ((k > 1) ? L[k,1:km1] * head(z, km1) : 0))) / L[k,k];
          real u_star = Phi(z_star);
          v = u_star * u[k];
          d[k] = u_star;
        } else { // lower bound
          real z_star = (lb[k] - (mu[k] + ((k > 1) ? L[k,1:km1] * head(z, km1) : 0))) / L[k,k];
          real u_star = Phi(z_star);
          d[k] = 1 - u_star;
          v = u_star + d[k] * u[k];
        }
        z[k] = inv_Phi(v);
      }
    }
    out[1] = z;
    out[2] = d;
    return out;
  }
}
data {
  int<lower=2> K;             // number of dimensions
  vector[K] mu;
  vector[K] sigma;
  vector[K] lb;              // lower bounds, -Inf for no bound
  vector[K] ub;              // upper bounds,  Inf for no bound
  real rho;
}

transformed data {
  matrix[K, K] Omega;
  matrix[K, K] Sigma;
  cholesky_factor_cov[K,K] L;
  for (m in 1:K) for (n in 1:K) Omega[m, n] = (m==n) ? 1.0 : rho;
  Sigma = quad_form_diag(Omega, sigma);
  L = cholesky_decompose(Sigma);
}

parameters {
  vector<lower=0,upper=1>[K] u;
}
model {
  target += log(make_stuff(mu, L, lb, ub, u)[2]); // Jacobian adjustments
  // implicit: u ~ uniform(0,1)
}
generated quantities {
  vector[K] y = mu + L * make_stuff(mu, L, lb, ub, u)[1];
}