Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Factor analysis sandbox
# Simulation script for factor analysis ala Leung & Drton (2016) ----------
library(rstan)
library(bayesplot)
m <- 5 # dimension of observed data (e.g., # traits)
k <- 2 # number of latent factors
n <- 100 # number of sample units (e.g., # species)
# residual variance matrix (is diagonal)
Omega <- diag(.3 + abs(rnorm(m, sd = .3)))
# hyperparameter for loading terms
C_0 <- 2
# Beta is the loading matrix, with zeros in upper triangle
Beta <- matrix(0, nrow = m, ncol = k)
for (i in 1:m) {
for (j in 1:k) {
if (i == j) {
Beta[i, i] <- abs(rnorm(1, mean = 0, sd = C_0))
} else if (i > j) {
Beta[i, j] <- rnorm(1, mean = 0, sd = C_0)
}
}
}
diag(Beta) <- rev(sort(diag(Beta)))
# construct covariance matrix and it's Cholesky factor
Sigma <- Omega + Beta %*% t(Beta)
L_Sigma <- t(chol(Sigma))
# generate observations
# y ~ Normal(0, Sigma)
y <- matrix(nrow = n, ncol = m)
for (i in 1:n) {
y[i, ] <- L_Sigma %*% rnorm(m)
}
pairs(y)
# view prior for diagonal terms in loading matrix (unscaled)
beta_diag <- seq(1E-6, 10, .1)
i <- 1
p_beta_diag <- (k - i) * log(beta_diag) - .5 * beta_diag ^ 2 / C_0
plot(beta_diag, exp(p_beta_diag) / max(exp(p_beta_diag)), type = "l",
ylab = "p(beta_diag)")
for (i in 2:k) {
p_beta_diag <- (k - i) * log(beta_diag) - .5 * beta_diag ^ 2 / C_0
lines(beta_diag, exp(p_beta_diag) / max(exp(p_beta_diag)),
type = "l", col = i)
}
legend("topright", lty = 1, col = 1:k, legend = paste("i =", 1:k))
# fake some missing data
old_y <- y
p_miss <- .5
n_miss <- floor(n * m * p_miss)
row_miss <- sample(n, n_miss, replace = TRUE)
col_miss <- sample(m, n_miss, replace = TRUE)
while (anyDuplicated(data.frame(row_miss, col_miss))) {
to_replace <- anyDuplicated(data.frame(row_miss, col_miss))
row_miss[to_replace] <- sample(n, 1)
col_miss[to_replace] <- sample(m, 1)
}
for (i in 1:n_miss) {
y[row_miss[i], col_miss[i]] <- NA
}
y_is_na <- is.na(y)
y_obs <- y[!y_is_na]
n_obs <- length(y_obs)
row_obs <- rep(1:n, times = m)[!y_is_na]
col_obs <- rep(1:m, each = n)[!y_is_na]
# fit model ---------------------------------------------------------------
stan_d <- list(k = k,
m = m,
n = n,
n_obs = n_obs,
y_obs = y_obs,
row_obs = row_obs,
col_obs = col_obs,
n_miss = n_miss,
row_miss = row_miss,
col_miss = col_miss)
library(rstan)
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())
m_mod <- stan_model("leung_drton_factor_analysis.stan")
m_fit <- sampling(object = m_mod, data = stan_d, chains = 3, iter = 2000,
pars = c("Sigma", "sigma_L", "y_miss"), init_r = .1)
# evaluate convergence
m_fit
post_array <- extract(m_fit, inc_warmup = TRUE, permuted = FALSE)
mcmc_trace(post_array, regex_pars = "Sigma")
mcmc_trace(post_array, regex_pars = "sigma_L")
mcmc_trace(post_array, regex_pars = "lp__")
# Compare predictions to true values --------------------------------------
post <- rstan::extract(m_fit)
y_ci <- apply(post$y_miss, 2,
FUN = function(x) quantile(x, probs = c(0.025, 0.975)))
y_df <- data.frame(lo = y_ci[1, ],
hi = y_ci[2, ],
med = c(apply(post$y_miss, 2, median)))
y_df$true <- NA
for (i in 1:n_miss) {
y_df$true[i] <- old_y[row_miss[i], col_miss[i]]
}
y_df$trait <- col_miss
ggplot(y_df, aes(x = true, y = med)) +
geom_point() +
geom_segment(aes(xend = true, y = lo, yend = hi), alpha = .5, col = "blue") +
geom_abline(intercept = 0, slope = 1, linetype = "dashed") +
facet_wrap(~trait, labeller = "label_both", scales = "free") +
theme_minimal()
data {
int<lower = 0> m; // dimension of observed data
int<lower = 0> k; // number of latent factors
int<lower = 0> n; // number of sample units (species)
int<lower = 0> n_obs; // number of observations
int<lower = 1, upper = n> row_obs[n_obs]; // row indices for observations
int<lower = 1, upper = m> col_obs[n_obs]; // column indices for observations
vector[n_obs] y_obs;
int<lower=0> n_miss; // number of missing observations
int<lower = 1, upper = n> row_miss[n_miss];
int<lower = 1, upper = m> col_miss[n_miss];
}
transformed data {
int<lower = 1> p; // number of nonzero lower triangular factor loadings
vector[m] zeros;
zeros = rep_vector(0, m);
p = k * (m - k) + k * (k - 1) / 2;
}
parameters {
vector<lower = 0>[k] beta_diag;
vector[p] beta_lower_tri;
vector<lower = 0>[m] sigma_y; // residual error
real<lower = 0> sigma_L; // hyperparameter for loading matrix elements
vector[n_miss] y_miss;
}
transformed parameters {
matrix[m, k] L;
cov_matrix[m] Sigma;
matrix[m, m] L_Sigma;
vector[m] y[n];
// build n by m matrix y by combining observed and missing observations
for (i in 1:n_obs) {
y[row_obs[i]][col_obs[i]] = y_obs[i];
//y[row_obs[i], col_obs[i]] = y_obs[i];
}
for (i in 1:n_miss) {
y[row_miss[i]][col_miss[i]] = y_miss[i];
//y[row_miss[i], col_miss[i]] = y_miss[i];
}
{ // temporary scope to define loading matrix L
int idx1;
int idx2;
for (i in 1:m){
for (j in (i + 1):k){
idx1 = idx1 + 1;
L[i, j] = 0; // constrain upper tri elements to zero
}
}
for (j in 1:k){
L[j, j] = beta_diag[j];
for (i in (j + 1):m){
idx2 = idx2 + 1;
L[i, j] = beta_lower_tri[idx2];
}
}
}
Sigma = multiply_lower_tri_self_transpose(L);
for (i in 1:m) Sigma[i, i] = Sigma[i, i] + sigma_y[i];
L_Sigma = cholesky_decompose(Sigma);
}
model {
// priors
beta_lower_tri ~ normal(0, sigma_L);
sigma_L ~ normal(0, 2);
sigma_y ~ normal(0, 2);
// priors for diagonal entries (Leung and Drton 2016)
for (i in 1:k) {
target += (k - i) * log(beta_diag[i]) - .5 * beta_diag[i] ^ 2 / sigma_L;
}
// likelihood
y ~ multi_normal_cholesky(zeros, L_Sigma);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment