R/EM_msmm.R

In RBesT: R Bayesian Evidence Synthesis Tools

## EM for MSMM (Multi-variate Student t Mixture Model) with Nc components
## if init is not specified, then knn is used to initialize means,
## cluster weights and covariance matrices (taken from the knn
## determined clusters)

## this function follows the details given in ref. "Robust mixture
## modelling using the t distribution". D. Peel and G.J. McLachlan,
## Statistics and Computing (2000), 10, 339-348,

EM_msmm <- function(X, Nc, init, Ninit=50, verbose=TRUE, Niter.max=500, tol=1e-1) {
    ## in case X is no matrix, interpret it as uni-variate case
    if(!is.matrix(X))
        X <- matrix(X,ncol=1)

    N <- dim(X)[1]
    Nd <- dim(X)[2]

    ## initialize randomly
    if(missing(init)) {
        ## assume that the sample is ordered randomly
        ind <- seq(1,N-Nc,length=Ninit)
        knnInit <- list(mu=matrix(0,nrow=Nc,ncol=Nd), p=(1/seq(1,Nc))/sum(seq(1,Nc)))
        for(k in seq(Nc))
            knnInit$mu[k,] <- colMeans(X[ind+k-1,,drop=FALSE])
        ## use k means clustering with K=Nc as init; ignore warnings
        ## as we may hit the maximal number of iterations
        suppressWarnings(KNN <- knn(X, K=Nc, init=knnInit, verbose=verbose, Niter.max=50))
        pEst <- KNN$p
        cmin <- which.min(pEst)
        muEst <- KNN$center
        ##nuEst <- rlnorm(Nc, log(20), log(5)/1.96) ## avoid randomness during initialization
        nuEst <- rep(12, times=Nc)
        covEst <- array(0, dim=c(Nc, Nd, Nd))
        Xtau <- sqrt(colVars(X))
        for(i in seq(Nc)) {
            if(i == cmin) next
            ind <- KNN$cluster==i
            if(sum(ind) > 10) {
                covKNN <- as.matrix(cov(X[ind,,drop=FALSE]))
                R <- cov2cor(covKNN)
                tau <- sqrt(diag(covKNN))
                ## set variances below or equal to 0 to the sample variance
                tau[tau <= 0] <- Xtau[tau <= 0]
            } else {
                R <- diag(Nd)
                tau <- Xtau
            }
            ##tauR <- rlnorm(Nd, log(tau), log(3)/1.96)
            ##covEst[i,,] <- diag(tauR, Nd, Nd) %*% R %*% diag(tauR, Nd, Nd)
            ## ensure that the smallest variance is not less than the global
            ## variance divided by 100... which is to stabilize things
            tau <- pmax(tau, Xtau / 100)
            covEst[i,,] <- diag(tau, Nd, Nd) %*% R %*% diag(tau, Nd, Nd)
        }
        ##tauR <- rlnorm(Nd, log(Xtau), log(3)/1.96)
        ##covEst[cmin,,] <- diag(tauR, Nd, Nd) %*% diag(Nd) %*% diag(tauR, Nd, Nd)
        covEst[cmin,,] <- diag(Xtau, Nd, Nd) %*% diag(Nd) %*% diag(Xtau, Nd, Nd)
        muEst[cmin] <- sum(pEst * muEst)
    } else {
        pEst <- init$p
        muEst <- init$center
        nuEst <- init$nu
        covEst <- init$cov
    }

    iter <- 0
    logN <- log(N)
    traceLli <- c()
    Dlli <- Inf

    ## initialize component and element wise log-likelihood matrix
    lli <- array(-20, dim=c(N,Nc))

    ## and the log(U) matrix which is the inverse mahalanobis distances
    ## decorated with nuEst
    lU <- array(0, dim=c(N,Nc))

    if(verbose) {
        message("EM multi-variate student t with Nc =", Nc, ":\n")
    }

    nu_ml <- function(c1) {
        function(nu) {
            (log(nu/2) - digamma(nu/2) + c1 + digamma((nu+Nd)/2) - log((nu+Nd)/2))^2
        }
    }

    while(iter < Niter.max) {
        ## calculate responsabilities from the likelihood terms;
        ## calculations are done in log-space to avoid numerical
        ## difficulties if some points are far away from some
        ## component and hence recieve very low density
        for(i in seq(Nc)) {
            lli[,i] <- log(pEst[i]) + mvtnorm::dmvt(X, muEst[i,], as.matrix(covEst[i,,]), nuEst[i], log=TRUE)
        }
        ## ensure that the log-likelihood does not go out of numerical
        ## reasonable bounds
        lli <- apply(lli, 2, pmax, -30)
        ##lnresp <- apply(lli, 1, log_sum_exp)
        lnresp <- matrixStats::rowLogSumExps(lli)
        ## the log-likelihood is then given by the sum of lresp
        lliCur <- sum(lnresp)
        traceLli <- c(traceLli, lliCur)
        if(iter > 1) {
            ## Dlli is the slope of the log-likelihood evaulated with
            ## a second order method
            Dlli <- (traceLli[iter+1] - traceLli[iter - 1])/2
        }
        if(verbose)
            message("Iteration", iter, ": log-likelihood", lliCur, "; Dlli =", Dlli, "\n")
        if(Dlli < tol) {
            break
        }
        ## ... and the responisbility matrix follows from this by
        ## appropiate normalization.
        lresp <- sweep(lli, 1, lnresp, "-") ## Eq. 16
        resp <- exp(lresp)

        ## calculate additional weights of the U matrix aka latent
        ## tail mass of a point
        for(i in seq(Nc)) {
            Xc_i <- sweep(X, 2L, muEst[i,])
            Sigma_i <- as.matrix(covEst[i,,])
            ## in rare cases the covariance matrix becomes (almost)
            ## singular in which case the alternative cholesky
            ## factorization gives more stable results
            maha_dist <- tryCatch(
                mahalanobis(Xc_i, FALSE, Sigma_i),
                error = function(e) {
                    ## see https://stats.stackexchange.com/questions/147210/efficient-fast-mahalanobis-distance-computation
                    ## also adding eps to the diagonal to further stabilize the computation
                    L_i <- t(chol(Sigma_i + diag(5 * .Machine$double.eps, Nd, Nd)))
                    y_i <- forwardsolve(L_i, t(Xc_i))
                    colSums(y_i^2)
                }
            )
            lU[,i] <- log(nuEst[i] + Nd) - log(nuEst[i] + maha_dist) ## Eq. 20
        }

        ## mean probability to be in a specific mixture component ->
        ## updates pEst
        ##lzSum <- apply(lresp, 2, log_sum_exp)
        lzSum <- colLogSumExps(lresp)
        zSum <- exp(lzSum)
        ##zSum <- colSums(resp)
        ##pEst <- zSum/N ## Eq. 29
        pEst <-  exp(lzSum - logN)

        ## make sure it is scale to exactly 1 which may not happen due
        ## to small rounding issues
        pEst <- pEst/sum(pEst)

        ## product of u weight and responsabilities
        lW <- lresp + lU ## intermediate formed, i.e. tau_ij * u_ij
        ##wSum <- exp(apply(lW, 2, log_sum_exp))
        wSum <- exp(colLogSumExps(lW))
        ##wSum <- colSums(W)

        ## now obtain new estimates for each component of the mixtures
        ## of their mu vector and covariance matrices
        for(i in seq(Nc)) {
            ##xc <- sqrt(W[,i]) * sweep(X, 2, muEst[i,], check.margin = FALSE)
            ##covEst[i,,] <- crossprod(xc) / zSum[i]    ## Eq. 31
            ##muEst[i,] <- colSums(W[,i] * X) / wSum[i] ## Eq. 30
            ##xc <- exp(0.5 * lW[,i]) * sweep(X, 2, muEst[i,], check.margin = FALSE)
            ##covEst[i,,] <- crossprod(xc) / zSum[i]          ## Eq. 31
            xc <- exp(0.5 * (lW[,i] - lzSum[i])) * sweep(X, 2, muEst[i,], check.margin = FALSE)
            covEst[i,,] <- crossprod(xc)                    ## Eq. 31 (divisor moved to xc)
            muEst[i,] <- colSums(exp(lW[,i]) * X) / wSum[i] ## Eq. 30
            ## ensure that diagonal stays non-zero
            for(j in 1:Nd)
                covEst[i,j,j] <- max(covEst[i,j,j], .Machine$double.eps)
        }

        ## finally get the new nu estimates via numerical solution
        ## first calculate necessary constants which don't involve v_i
        ##c1 <- 1 + colSums(resp * (log(U) - U)) / zSum
        c1 <- 1 + colSums(resp * (lU - exp(lU))) / zSum
        for(i in seq(Nc)) {
            nuEstML <- optimize(nu_ml(c1[i]), c(0,150)) # Eq. 32
            if(is.na(nuEstML$objective)) {
                warning("Component ", i, " in iteration ", iter, " failed convergence.")
            } else if(nuEstML$objective > 1e-3 & nuEstML$minimum < 50) {
                ## only warn if we had trouble finding the minimum
                ## when below 50, larger values are anyway normals
                warning("Component ", i, " in iteration ", iter, " had convergence problems.\nObjective function = ",nuEstML$objective,"\n")
            }
            nuEst[i] <- nuEstML$minimum
        }

        iter <- iter + 1
    }
    if(iter == Niter.max)
        warning("Maximum number of iterations reached.")

    ## degrees of freedom
    ## covariance matrix df per component
    cov.df <- (Nd-1)*Nd/2 + Nd
    df <- Nc  * (Nd + cov.df) + Nc + Nc-1

    ## sort by largest weight
    o <- order(pEst, decreasing=TRUE)
    pEst <- pEst[o]
    muEst <- muEst[o,,drop=FALSE]
    covEst <- covEst[o,,,drop=FALSE]
    nuEst <- nuEst[o,drop=FALSE]

    if(Nd != 1) {
        rhoEst <- array(apply(covEst, 1, cov2cor), c(Nd,Nd,Nc))
        rhoEst <- apply(rhoEst, 3, function(x) x[lower.tri(x)])
        tauEst <- sqrt(t(apply(covEst, 1, diag)))
    } else {
        rhoEst <- NULL
        tauEst <- sqrt(as.vector(covEst))
    }

    invisible(list(cov=covEst, center=muEst, nu=nuEst, p=pEst, rho=rhoEst, tau=tauEst, lli=lliCur, df=df, Dlli=Dlli, niter=iter))
}