R/bfa.R

In vsartor/bnsa: Bayesian Non-Structured Analysis

#' Bayesian Factor Analysis
#'
#' Simulates from the posterior distribution of the parameters of a
#' simple Bayesian Factor Analysis using the Gibbs algorithm,
#' considering vague prior distributions for all parameters.
#'
#' @param Y A T by m matrix containing each m-dimensional observation
#'          in each row.
#' @param k The number of factors to be considered.
#' @param N The number of iterations for the MCMC algorithm to run.
#'
#' @return A list with the chains `factor`, `beta` and `vars`.
#'
#' @importFrom stats rgamma
#' @importFrom MASS mvrnorm
#'
#' @export
#'
bfa = function(Y, k, N = 5000) {
    # Dimensions of the problem
    n = nrow(Y)
    m = ncol(Y)

    # List to be returned
    ch         = list()
    ch$factors = array(0, dim = c(n, k, N))
    ch$beta    = array(0, dim = c(m, k, N))
    ch$vars    = array(0, dim = c(m, N))

    # Hyperparameters
    nu = 1
    s  = 1
    m0 = 0
    C0 = 1e8

    # Set up non-zero initial values
    if (k == 1) {
        ch$beta[ ,1,1] = 1
    } else {
        diag(ch$beta[ , ,1]) = 1
    }
    ch$vars[ ,1]  = 1

    # Initialize loop variables
    factors = matrix(ch$factors[ , ,1], nrow = n)
    beta    = matrix(ch$beta[ , ,1], nrow = m)
    vars    = ch$vars[ ,1]

    # Initialize constants
    I = vector("list", k)
    for (i in 1:k) I[[i]] = diag(i)

    for (it in 2:N) {
        # Update user on progress
        if (it %% 250 == 0)
            cat(sprintf("[%d/%d] Computing...\n", it, N))

        # Draw for the factors
        Cov = solve(I[[k]] + t(beta) %*% diag(1/vars) %*% beta)
        Fac = t(beta) %*% diag(1/vars)
        for (j in 1:n) {
            factors[j, ] = mvrnorm(1, Cov %*% Fac %*% Y[j, ], Cov)
        }
        ch$factors[ , ,it] = factors

        # Draw for the variances
        for (i in 1:m) {
            di      = sum((Y[ ,i] - factors %*% beta[i, ])^2)
            vars[i] = 1 / rgamma(1, (nu + n)/2, (nu*s + di)/2)
        }
        ch$vars[ ,it] = vars

        # Draw for loadings
        for (i in 1:k) {
            beta[i,1:i] = -1
            safeout = 0
            while (beta[i,i] <= 0) {
                safeout = safeout + 1
                if (safeout >= 1000) stop("Failed to generate from truncated normal")

                Fac         = factors[ ,1:i]
                Ci          = solve(I[[i]] / C0 + t(Fac) %*% Fac / vars[i])
                mi          = Ci %*% (m0 * rep(1, i) / C0 + t(Fac) %*% Y[ ,i] / vars[i])
                beta[i,1:i] = mvrnorm(1, mi, Ci)
            }
        }
        for (i in (k+1):m) {
            Fac       = factors
            Ci        = solve(I[[k]] / C0 + t(Fac) %*% Fac / vars[i])
            mi        = Ci %*% (m0 * rep(1, k) / C0 + t(Fac) %*% Y[ ,i] / vars[i])
            beta[i, ] = mvrnorm(1, mi, Ci)
        }
        ch$beta[ , ,it] = beta
    }

    # Return the chains
    ch
}