Correlation matrix with positively constrained off-diagonals

Modeling
3

I had most of this code already, just constraining tanh(x) to positive.

functions {
  vector corr_constrain_lp(vector x) {
    int N = num_elements(x);
    vector[N] tanh_x = tanh(x);
    
    target += sum(log1m(square(tanh_x)));
    
    return tanh_x;
  }

  matrix cholesky_corr_constrain_lp (vector y, int K) {
    int k_choose_2 = (K * (K - 1)) / 2;
    vector[k_choose_2] z = corr_constrain_lp(y);
    matrix[K, K] x = rep_matrix(0, K, K);
    int iter = 0;
    
    x[1, 1] = 1;

    for (i in 2:K) {
      real sum_sqs;
      iter += 1;
      x[i, 1] = z[iter];
      sum_sqs = square(x[i, 1]);
     if (i > 2) {
      for (j in 2:i - 1) {
        iter += 1;
        target += 0.5 * log1m(sum_sqs);
        x[i, j] = z[iter] * sqrt(1.0 - sum_sqs);
        sum_sqs += square(x[i, j]);
      }
     }
    x[i, i] = sqrt(1.0 - sum_sqs);
  }
  return x;
  }
}
data {
  int K;
  real eta;
}
transformed data {
  int k_choose_2 = (K * (K - 1)) / 2;
}
parameters {
  vector<lower=0>[k_choose_2] y_raw;
}
transformed parameters {
  matrix[K, K] y = cholesky_corr_constrain_lp(y_raw, K);
}
model {
  y ~ lkj_corr_cholesky(eta);
}
generated quantities {
  matrix[K, K] Omega = multiply_lower_tri_self_transpose(y);
}

Run with

library(cmdstanr)

fp <- file.path("../positive_corr.stan")
mod <- cmdstan_model(fp)

K <- 10
mod_out <- mod$sample(
  data = list(K = K,
              eta = 100),
  chains = 2,
  adapt_delta = 0.8,
  parallel_chains = 2,
  iter_warmup = 200,
  iter_sampling = 200
)

matrix(mod_out$summary("Omega")$mean, K, K)
3 Likes