R/EM_mnmm.R

In RBesT: R Bayesian Evidence Synthesis Tools

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

EM_mnmm <- function(X, Nc, mix_init, Ninit=50, verbose=FALSE, Niter.max=500, tol, Neps, eps=c(w=0.005,m=0.005,s=0.005)) {
    ## 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]
    assert_that(N+Nc >= Ninit)

    ## initialize normal EM using a Student-t EM which is very robust
    ## against outliers
    if(missing(mix_init)) {
        mix_init <- EM_msmm(X, Nc, Ninit=Ninit, verbose=verbose, Niter.max=round(Niter.max/2), tol=0.1)
    }
    pEst <- mix_init$p
    muEst <- mix_init$center
    covEst <- mix_init$cov

    ## take current estimates and transform to scale on which
    ## convergence is assessed
    est2par <- function(p, mu, cov) {
        est <- rbind(logit(p), matrix(sapply(1:Nc, function(i) mv2vec(mu[i,], cov[i,,])), ncol=Nc))
        est[(1+Nd+1):(1+2*Nd),] <- log(est[(1+Nd+1):(1+2*Nd),])
        est
    }
    ## in case tolerance is not specified, then this criteria is
    ## ignored
    if(missing(tol)) {
        checkTol <- FALSE
        tol <- -1
    } else
        checkTol <- TRUE

    if(missing(Neps)) {
        ## in case tolerance has been declared, but Neps not, we flag
        ## to disable checking of running mean convergence check
        checkEps <- FALSE
        Neps <- 5
    } else
        checkEps <- TRUE

    ## if nothing is specified, we declare convergence based on a
    ## running mean of differences in parameter estimates
    if(!checkTol & !checkEps) {
        checkEps <- TRUE
    }

    assert_that(Neps > 1)
    assert_that(ceiling(Neps) == floor(Neps))

    ## eps can also be given as a single integer which is interpreted
    ## as number of digits
    if(length(eps) == 1) eps <- rep(10^(-eps), 3)

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

    ## expand eps according to the dimensionality
    eps <- c(eps[1], rep(eps[2], Nd), rep(eps[3], Nd), rep(eps[2], (Nd-1)*Nd/2))

    iter <- 0
    logN <- log(N)
    traceMix <- list()
    traceLli <- c()
    Dlli <- Inf
    runMixPar <- array(-Inf, dim=c(Neps,df.comp,Nc))
    runOrder <- 0:(Neps-1)
    if(Nc == 1) Npar <- df else Npar <- df + 1

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

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

    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]) + dmvnorm(X, muEst[i,], as.matrix(covEst[i,,]), log=TRUE, checkSymmetry=FALSE)
        }
        ## 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)
        ## record current state
        traceMix <- c(traceMix, list(list(p=pEst, mean=muEst, sigma=covEst)))
        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(Nc > 1) {
            smean <- apply(runMixPar[order(runOrder),,,drop=FALSE], c(2,3), function(x) mean(abs(diff(x))))
            eps.converged <- sum(sweep(smean, 1, eps, "-") < 0)
        } else {
            smean <- apply(runMixPar[order(runOrder),-1,,drop=FALSE], c(2,3), function(x) mean(abs(diff(x))))
            eps.converged <- sum(sweep(smean, 1, eps[-1], "-") < 0)
        }
        if(is.na(eps.converged)) eps.converged <- 0
        if(verbose)
            message("Iteration ", iter, ": log-likelihood = ", lliCur, "; Dlli = ", Dlli, "; converged = ", eps.converged, " / ", Npar, "\n", sep="")
        if(checkTol & Dlli < tol) {
            break
        }
        if(iter >= Neps & checkEps & eps.converged == Npar) {
            break
        }
        ## ... and the responseability matrix follows from this by
        ## appropiate normalization.
        lresp <- sweep(lli, 1, lnresp, "-")
        ##resp <- exp(lresp)

        ## mean probability to be in a specific mixture component -> updates
        ## pEst
        ##lzSum <- apply(lresp, 2, log_sum_exp)
        lzSum <- colLogSumExps(lresp)
        ##zSum <- exp(lzSum)
        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)

        ## now obtain new estimates for each component of the mixtures
        ## of their mu vector and covariance matrices
        for(i in seq(Nc)) {
            upd <- cov.wt(X, exp(lresp[,i] - lzSum[i]), method="ML")
            ##upd <- cov.wt(X, resp[,i]/zSum[i], method="ML")
            muEst[i,] <- upd$center
            covEst[i,,] <- upd$cov
            ## ensure that diagonal stays non-zero
            for(j in 1:Nd)
                covEst[i,j,j] <- max(covEst[i,j,j], .Machine$double.eps)
        }

        ind <- 1 + iter %% Neps
        runMixPar[ind,,] <- est2par(pEst, muEst, covEst)
        runOrder[ind] <- iter

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

    ## sort by largest weight
    o <- order(pEst, decreasing=TRUE)
    pEst <- pEst[o]
    muEst <- muEst[o,,drop=FALSE]
    covEst <- covEst[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))
    ##}

    ##mixEst <- list(p=pEst, mean=muEst, sigma=covEst)
    
    mixEst <- do.call(mixmvnorm, lapply(1:Nc, function(i) c(pEst[i], muEst[i,,drop=TRUE], matrix(covEst[i,,], Nd, Nd))))
    
    ## give further details
    attr(mixEst, "df") <- df
    attr(mixEst, "nobs") <- N
    attr(mixEst, "lli") <- lliCur

    attr(mixEst, "Nc") <- Nc

    convert <- function(est) suppressWarnings(do.call(mixmvnorm, lapply(1:Nc, function(i) c(est$p[i], est$mean[i,,drop=FALSE], matrix(est$sigma[i,,], Nd, Nd)))))

    attr(mixEst, "tol") <- tol
    attr(mixEst, "traceLli") <- traceLli
    attr(mixEst, "traceMix") <- lapply(traceMix, convert)
    attr(mixEst, "x") <- X

    class(mixEst) <- c("EM", "EMmvnmm", "mvnormMix", "mix")

    mixEst
}

## uni-variate case
EM_nmm <- function(X, Nc, mix_init, verbose=FALSE, Niter.max=500, tol, Neps, eps=c(w=0.005,m=0.005,s=0.005)) {
    if(is.matrix(X)) {
        assert_matrix(X, any.missing=FALSE, ncols=1)
    }
    mixEst <- EM_mnmm(X=X, Nc=Nc, mix_init=mix_init, verbose=verbose, Niter.max=Niter.max, tol=tol, Neps=Neps, eps=eps)
    rownames(mixEst) <- c("w", "m", "s")
    class(mixEst) <- c("EM", "EMnmm", "normMix", "mix")
    attr(mixEst, "traceMix") <- lapply(attr(mixEst, "traceMix"),
                                       function(x) {
                                           class(x) <- class(mixEst)
                                           rownames(x) <- rownames(mixEst)
                                           x
                                       })
    likelihood(mixEst) <- "normal"
    mixEst
}

#' @export
print.EMnmm <- function(x, ...) {
    cat("EM for Normal Mixture Model\nLog-Likelihood = ", logLik(x), "\n\n",sep="")
    NextMethod()
}

#' @export
print.EMmvnmm <- function(x, ...) {
    cat("EM for Multivariate Normal Mixture Model\nLog-Likelihood = ", logLik(x), "\n\n",sep="")
    NextMethod()
}

## given a vector of means and a covariance matrix for a multi-variate
## normal (and optionally a vector of df), return a vector of
## coefficients with a deterministic mapping
mv2vec <- function(mean, sigma, df, label=TRUE) {
    Nd <- length(mean)
    if(Nd != 1) {
        rho <- cov2cor(sigma)[lower.tri(sigma)]
        tau <- sqrt(diag(sigma))
    } else {
        rho <- NULL
        tau <- sqrt(sigma)
    }
    if(missing(df)) df <- NULL
    res <- c(mean, tau, rho, df)
    if(label) {
        tauN <- paste("sd", 1:Nd, sep="")
        cols <- names(mean)
        if(is.null(cols))
            cols <- paste("mu", 1:Nd, sep="")
        if(Nd == 1)
            corNames <- NULL
        else if(length(cols) == 2)
            corNames <- "cor"
        else {
            corNames <- outer(cols, cols, paste, sep="_")
            corNames <- paste("cor", corNames[lower.tri(corNames)], sep=".")
        }
        if(!is.null(df))
            dfNames <- paste("df", 1:Nd, sep="")
        else
            dfNames <- NULL
        names(res) <- c(cols, tauN, corNames, dfNames)
    }
    return(res)
}

## vec2mv <- function(vec) {
##     ## calculate dimension from the number of parameters
##     Nd <- -3/2 + sqrt( 9/4 + 2*length(vec) )
##     mean <- vec[1:Nd]
##     Tau <- diag(vec[(Nd+1):(2*Nd)])
##     sigma <- diag(Nd)/2
##     L <- lower.tri(sigma)
##     E <- sum(L)
##     sigma[L] <- vec[(2*Nd+1):(2*Nd+E)]
##     sigma <- sigma + t(sigma)
##     sigma <- Tau %*% sigma %*% Tau
##     list(mean=mean, sigma=sigma)
## }
##
## ## extracts results in a flattened form
## extractMVNmix <- function(emFit) {
##     pMix <- emFit$p
##     cols <- colnames(emFit$center)
##     if(is.null(cols))
##         cols <- paste("X", 1:ncol(emFit$center), sep="")
##     if(length(cols) == 1)
##         corNames <- NULL
##     else if(length(cols) == 2)
##         corNames <- "cor"
##     else {
##         corNames <- outer(cols, cols, paste, sep="_")
##         corNames <- paste("cor", corNames[lower.tri(corNames)], sep=".")
##     }
##     MAPmix <- do.call(cbind, list(emFit$center, emFit$tau, emFit$rho))
##     colnames(MAPmix) = c(paste("mean", cols, sep="."), paste("sd", cols, sep="."), corNames)
##     MAPmix <- as.data.frame(t(MAPmix))
##     colnames(MAPmix) <- paste("Mix", 1:length(pMix), sep="")
##     names(pMix) <- names(MAPmix)
##     list(mvn=MAPmix, pMix=pMix)
## }
##
## ## utility functions for multi-variate normal mixtures
## rmvnormMix <- function(n, p, m, sigma){
##   ind <- sample.int(length(p), n, replace = TRUE, prob = p)
##   d <- nrow(m)
##   samp <- matrix(0, nrow=n, ncol=d)
##   for(i in seq(n))
##       samp[i,] <- rmvnorm(1, m[ind[i],], as.matrix(sigma[ind[i],,]))
##   samp
## }
##
## dmvnormMix <- function(x, p, m, sigma) {
##     nc <- length(p)
##     ## logic is to replicate the original data vector such that each
##     ## item appears nc times which allows easy vectorized calls to
##     ## dnorm. Then we cast the result into a matrix with as many rows
##     ## as nc components which we sum together with a fast colSums call.
##     dens <- rep.int(0, nrow(x))
##     for(i in seq_along(p))
##         dens <- dens + p[i] * dmvnorm(x, m[i,], sigma[i,,])
##     dens
## }
##
## ## utility function to plot BVN mixtures
## bicontourNMM <- function(X, bvn, title="", ng=50) {
##     ## some pretty colors
##     ##library(colorspace)
##     k <- 15
##     ##my.cols <- diverge_hcl(k, c = c(100, 0), l = c(50, 90), power = 1.3)
##     my.cols <- c("#4A6FE3", "#6D84E1", "#8898E1", "#A0AAE2", "#B5BBE3", "#C7CBE3", "#D7D9E3",
##                  "#E2E2E2", "#E4D6D8", "#E6C4CA", "#E6AFB9", "#E498A7", "#E07E93", "#DB627F",
##                  "#D33F6A")
##
##     r <- apply(X, 2, range)
##     xg <- seq(r[1,1],r[2,1], length=ng)
##     yg <- seq(r[1,2],r[2,2], length=ng)
##     Z <- outer(xg,yg, function(x,y) dmvnormMix(cbind(x,y), bvn$p, bvn$center, bvn$cov))
##
##     Nc <- length(bvn$p)
##
##     plot(X, pch=19, cex=.2)
##     contour(xg, yg, Z, drawlabels=FALSE, nlevels=k, col=my.cols, add=TRUE, lwd=2)
##     abline(h=bvn$center[,2], v=bvn$center[,1], lwd=1, col=1:Nc)
##     title(title)
##     legend("bottomleft", legend=paste("Mix", 1:Nc, " ", round(100*bvn$p,1), "%", sep=""), text.col=1:Nc)
## }
##