Factor analysis sandbox

leung_drton_factor_analysis.R

# 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()

leung_drton_factor_analysis.stan

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 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); 
}