Instantly share code, notes, and snippets.

# mbjoseph/leung_drton_factor_analysis.R Last active Jun 22, 2018

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 m; // dimension of observed data int k; // number of latent factors int n; // number of sample units (species) int n_obs; // number of observations int row_obs[n_obs]; // row indices for observations int col_obs[n_obs]; // column indices for observations vector[n_obs] y_obs; int n_miss; // number of missing observations int row_miss[n_miss]; int col_miss[n_miss]; } transformed data { int 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[k] beta_diag; vector[p] beta_lower_tri; vector[m] sigma_y; // residual error real 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 idx = 0; for (i in 1:m){ for (j in (i + 1):k){ 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){ idx = idx + 1; L[i, j] = beta_lower_tri[idx]; } } } 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); }