# R/EMAlg.R In REMLA: Robust Expectation-Maximization Estimation for Latent Variable Models

#### Defines functions EMAlg

#' Expectation Maximization Algorithm
#' @param X numeric data frame containing all the variables needed in the factor analysis
#' @param k number of factors considered in the analysis
#' @param constraints p x k matrix of zeros and ones denoting the factors (rows) and observed variables (columns)
#' @param rotation factor rotation method (either 'oblimin' or 'varimax')
#' @param ctrREM provides control parameters (default: controlREM(steps = 25, tol = 1e-6, maxiter = 1e3))
#' @returns List of outputs that contains the following components:
#'  \item{mu}{intercept}
#'  \item{psi}{variance}
#'  \iten(phi){factor covariance matrix}
#'  \item{ind_lik}{likelihood value for each individual}
#' @author Bryan Ortiz-Torres (bortiztorres@wisc.edu); Kenneth Nieser (nieser@wisc.edu)
#' @references Nieser, K. J., & Cochran, A. L. (2021). Addressing heterogeneous populations in latent variable settings through robust estimation. Psychological Methods.
#' @importFrom stats factanal quantile rnorm varimax na.omit cov2cor pnorm
#' @importFrom GPArotation oblimin
#' @noRd

EMAlg <- function(X, k, constraints, rotation, ctrREM = controlREM()){

# data dimensions
n = nrow(X)
p = ncol(X)

# sample mean and covariance matrix
mu = apply(X, 2, mean)
X_centered = t(apply(X, 1, function(y) y - mu))
Cxx = 1/n * t(X_centered) %*% X_centered

if(k == 1)
{
intl_guess <- factanal(X, k) #Maximum likelihood factor analysis  (r function)
inv_rotmat <- 1
} else {
intl_guess <- factanal(X, k, rotation = 'oblimin') #Maximum likelihood factor analysis  (r function)
inv_rotmat <- solve(intl_guess$rotmat) } lambda <- matrix(intl_guess$loadings, ncol = k)
lambda[constraints==0] <- 0
psi <- intl_guess$uniquenesses #describes variances of the error terms (unique factors$\epsilon in the paper$) phi <- inv_rotmat %*% t(inv_rotmat) # tolerance parameters tol = ctrREM$tol
maxiter = ctrREM$maxiter # constraint patterns to be processed separately patterns = unique(constraints) pattern.id = 1:nrow(patterns) for (j in 1:nrow(patterns)){ idx = which(apply(constraints, 1, function(x) return(all(x == c(patterns[j,]))))) pattern.id[idx] = j } # EM algorithm obj = rep(-Inf, maxiter) for (iter in 2:maxiter){ # calculate log-likelihoods sigma = lambda %*% phi %*% t(lambda) + diag(psi) V = chol(sigma) inv_V = solve(V) logdetV = log(det(V)) Z = t(apply(X, 1, function(y) y - mu)) %*% inv_V ind_lik = -(1/2)*(p*log(2*pi) + 2*logdetV + apply(Z, 1, function(y) t(y) %*% y)) # calculate value for objective function obj[iter] = sum(ind_lik) # check for convergence if(abs(obj[iter] - obj[iter-1]) < tol*abs(obj[iter-1])){break} # E step beta = phi %*% (t(lambda) %*% inv_V) %*% t(inv_V) CxxB = Cxx %*% t(beta) Ezz = (diag(k) - beta %*% lambda) %*% phi + beta %*% CxxB # M step # estimate the subsets of lambda and psi for each constraint pattern for (i in 1:nrow(patterns)){ item_idx <- pattern.id==i factor_idx <- patterns[i,]!=0 Cxx.ii = diag(Cxx)[item_idx] CxxB.i = matrix(CxxB[item_idx, factor_idx], nrow = length(Cxx.ii)) Ezz.i = Ezz[factor_idx, factor_idx] chol_Ezz.i = chol(Ezz.i) inv_chol_Ezz.i = solve(chol_Ezz.i) lambda[item_idx, factor_idx] = (CxxB.i %*% inv_chol_Ezz.i) %*% t(inv_chol_Ezz.i) psi[item_idx] = Cxx.ii - diag(((CxxB.i %*% inv_chol_Ezz.i) %*% t(inv_chol_Ezz.i)) %*% t(CxxB.i)) } phi = cov2cor(phi - beta %*% lambda %*% phi + beta %*% CxxB) if(iter==maxiter){message( paste('maxiter reached; consider increasing maxiter or increasing convergence tolerance'))} } # rescale estimates sigma = lambda %*% phi %*% t(lambda) + diag(psi) total_var = diag(sigma) lambda = apply(lambda, 2, function(x) x / sqrt(total_var)) psi = psi / total_var mu = mu / sqrt(total_var) # rotate loadings if (k > 1){ if (rotation == "oblimin"){lambda = matrix(GPArotation::oblimin(lambda)$loadings, ncol = k)}