R/EM_gmm.R

In RBesT: R Bayesian Evidence Synthesis Tools

## EM for GMM with Nc components
EM_gmm <- function(x, Nc, mix_init, Ninit=50, verbose=FALSE, Niter.max=500, tol, Neps, eps=c(weight=0.005,alpha=0.005,beta=0.005)) {
    N <- length(x)
    assert_that(N+Nc >= Ninit)

    ## check data for 0  values which are problematic, but may be
    ## valid. Moving these to eps ensures proper handling during fit.
    x0 <- x==0
    if(any(x0)) {
        message("Detected ", sum(x0), " value(s) which are exactly 0.\nTo avoid numerical issues during EM such values are moved to smallest eps on machine.")
        x[x0] <- .Machine$double.eps
    }

    ## temporaries needed during EM
    Lx <- matrix(log(x), ncol=Nc, nrow=N)

    xRep <- matrix(x, ncol=Nc, nrow=N)

    ## initialize randomly using KNN
    if(missing(mix_init)) {
        ## assume that the sample is ordered randomly
        ind <- seq(1,N-Nc,length=Ninit)
        knnInit <- list(mu=matrix(0,nrow=Nc,ncol=1), p=rep(1/Nc, times=Nc))
        for(k in seq(Nc))
            knnInit$mu[k,1] <- mean(x[ind+k-1])
        KNN <- suppressWarnings(knn(x, K=Nc, init=knnInit, Niter.max=50))
        muInit <- rep(mean(x), times=Nc)
        varInit <- rep(1.5*var(x), times=Nc)
        for(k in 1:Nc) {
            kind <- KNN$cluster == k
            if(sum(kind) > 10) {
                muInit[k] <- KNN$center[k]
                varInit[k] <- var(x[kind])
            }
        }
        ## relocate the component with the least weight in the center
        ## and assign the sample variance to it; the idea is that we
        ## expect an informative component and a heavy tailed
        ## background which is best "absorbed" if a wide component is
        ## place initially at the center of the data
        cmin <- which.min(KNN$p)
        varInit[cmin] <- var(x)
        muInit[cmin] <- sum(KNN$center * KNN$p)
        ##varInit <- rlnorm(Nc, log(varInit), log(2.5)/1.96)
        bInit <- muInit / varInit
        aInit <- muInit * bInit
        mixEst <- rbind(KNN$p, aInit, bInit)
        dlink(mixEst) <- identity_dlink
        rownames(mixEst) <- c("w", "a", "b")
    } else {
        mixEst <- mix_init
    }

    if(verbose) {
        message("EM for gamma mixture model.\n")
        message("Initial estimates:\n")
        print(mixEst)
    }

    ## mixEst parametrization during fitting
    mixEstPar <- mixEst
    mixEstPar[1,] <- logit(mixEst[1,,drop=FALSE])
    mixEstPar[2,] <- log(mixEst[2,])
    mixEstPar[3,] <- log(mixEst[3,])
    rownames(mixEstPar) <-  c("w", "la", "lb")

    ## the optimizer needs a fixed range where search log-alpha
    MLrange <- c(min(mixEstPar[2,]) - log(1e4), max(mixEstPar[2,]) + log(1e4))

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

    iter <- 0
    logN <- log(N)
    traceMix <- list()
    traceLli <- c()
    Dlli <- Inf
    runMixPar <- array(-Inf, dim=c(Neps,3,Nc), dimnames=list(NULL, rownames(mixEstPar), NULL ))
    runOrder <- 0:(Neps-1)
    Npar <- Nc + 2*Nc
    if(Nc == 1) Npar <- Npar - 1

    ## find alpha and beta for a given component in log-space
    gmm_ml <- function(c1) {
        function(la) {
            (c1 - digamma(exp(la)) + la)^2
        }
    }

    gmm_ml_grad <- function(c1) {
        function(la) {
            a <- exp(la)
            val <- (c1 - digamma(a) + la)
            grad <- 2 * val * (1 - trigamma(a) * a)
            grad
        }
    }

    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
        ##li <- t(matrix(abmEst[,1] * dgamma(xRep, abmEst[,2], abmEst[,3]), nrow=Nc))

        ##lli <- t(matrix(log(mixEst[1,]) + dgamma(xRep, mixEst[2,], mixEst[3,], log=TRUE), nrow=Nc))
        w <- mixEst[1,]
        a <- mixEst[2,]
        b <- mixEst[3,]
        ## Gamma density: x^(a-1) * exp(-b * x) * b^a / Gamma(a)
        lli  <- sweep(sweep(Lx, 2, a-1, "*") - sweep(xRep, 2, b, "*"), 2, a * log(b) - lgamma(a) + log(w), "+")

        ## 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 norms
        lliCur <- sum(lnresp)
        ## record current state
        traceMix <- c(traceMix, list(mixEst))
        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 (log) 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
        ## mixEst first row
        ##lzSum <- apply(lresp, 2, log_sum_exp)
        lzSum <- colLogSumExps(lresp)
        zSum <- exp(lzSum)
        mixEst[1,] <- exp(lzSum - logN)

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

        ##lrx <- apply(Lx + lresp, 2, log_sum_exp)
        lrx <- colLogSumExps(Lx + lresp)
        resp_zscaled  <- exp(sweep(lresp, 2, lzSum, "-"))
        c1 <- colSums(Lx * resp_zscaled) + lzSum - lrx
        c2 <- lzSum - lrx

        ## now solve for new alpha and beta estimates
        for(i in 1:Nc) {
            Lest <- optimize(gmm_ml(c1[i]), MLrange)
            ##theta <- c(log(mixEst[2,i]))
            ##Lest <- optim(theta, gmm_ml(c1[i]), gr=gmm_ml_grad(c1[i]), method="BFGS", control=list(maxit=500))
            if(abs(Lest$objective) > 1E-4) {
                warning("Warning: Component", i, "in iteration", iter, "had convergence problems!")
            }
            mixEstPar[2,i] <- Lest$minimum
            mixEstPar[3,i] <- Lest$minimum + c2[i]
            mixEst[c(2,3),i] <- exp(mixEstPar[c(2,3),i])
        }

        mixEstPar[1,] <- logit(mixEst[1,,drop=FALSE])
        ind <- 1 + iter %% Neps
        runMixPar[ind,,] <- mixEstPar
        runOrder[ind] <- iter

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

    mixEst <- mixEst[,order(mixEst[1,], decreasing=TRUE),drop=FALSE]
    colnames(mixEst) <- paste("comp", seq(Nc), sep="")
    dlink(mixEst) <- identity_dlink
    class(mixEst) <- c("EM", "EMgmm", "gammaMix", "mix")

    ## give further details
    attr(mixEst, "df") <- Nc-1 + 2*Nc
    attr(mixEst, "nobs") <- N
    attr(mixEst, "lli") <- lliCur

    attr(mixEst, "Nc") <- Nc

    attr(mixEst, "tol") <- tol
    attr(mixEst, "traceLli") <- traceLli
    attr(mixEst, "traceMix") <- lapply(traceMix, function(x) {class(x) <- c("gammaMix", "mix"); x})
    attr(mixEst, "x") <- x

    mixEst
}


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