R/EM_bmm_mun.R

In RBesT: R Bayesian Evidence Synthesis Tools

## EM for Beta Mixture Models (BMM) with Nc components

EM_bmm_mun <- function(
  x,
  Nc,
  mix_init,
  Ninit = 50,
  verbose = FALSE,
  Niter.max = 500,
  tol,
  Neps,
  eps = c(w = 0.005, m = 0.005, N = 0.005)
) {
  N <- length(x)
  assert_that(N + Nc >= Ninit)

  ## check data for 0 and 1 values which are problematic, but may be
  ## valid, depending on a and b. Moving these to eps or 1-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
  }
  x1 <- x == 1
  if (any(x1)) {
    message(
      "Detected ",
      sum(x1),
      " value(s) which are exactly 1.\nTo avoid numerical issues during EM such values are moved to one minus smallest eps on machine."
    )
    x[x1] <- 1 - .Machine$double.eps
  }

  ## temporaries needed during EM
  LxO <- matrix(logit(x), ncol = Nc, nrow = N)
  LxC <- matrix(log1p(-x), ncol = Nc, nrow = N)
  sLxO <- colSums(LxO)
  sLxC <- colSums(LxC)

  xRep <- rep(x, each = Nc)

  ## 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])
      }
    }
    nInit <- pmax(muInit * (1 - muInit) / varInit - 1, 1, na.rm = TRUE)
    ## place the component which recieved the least weight at the
    ## data center with roughly the variance of the sample
    cmin <- which.min(KNN$p)
    muInit[cmin] <- sum(KNN$p * KNN$center)
    ## muInit[cmin] <- mean(x) ## could be considered here
    nInit[cmin] <- pmax(
      muInit[cmin] * (1 - muInit[cmin]) / var(x) - 1,
      1,
      na.rm = TRUE
    )
    ## Nmax <- max(2, max(nInit))
    ## ensure n is positive for each cluster; if this is not the
    ## case, sample uniformly from the range of n we have
    ## Nneg <- nInit <= .Machine$double.eps
    ## Nsmall <- nInit <= 0.5
    ## if(any(Nsmall))
    ##    nInit[Nsmall] <- runif(sum(Nsmall), 0.5, Nmax)
    ## nInitR <- 0.5 + rlnorm(Nc, log(nInit), log(5)/1.96)
    ## mixEst <- rbind(KNN$p, nInitR*muInit, nInitR*(1-muInit))
    mixEst <- rbind(KNN$p, nInit * muInit, nInit * (1 - muInit))
    dlink(mixEst) <- identity_dlink
    rownames(mixEst) <- c("w", "a", "b")
  } else {
    mixEst <- mix_init
  }

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

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

  ## 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
  bmm_mun_ml <- function(c1, c2) {
    function(par) {
      mu <- inv_logit(par[1])
      nmu <- inv_logit(-par[1]) ## 1-mu
      n <- exp(par[2])
      ab <- c(mu * n, nmu * n)
      di <- digamma(ab[2])
      eq1 <- digamma(ab[1]) - di
      eq2 <- di - digamma(n)
      (eq1 - c1)^2 + (eq2 - c2)^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
    lli <- t(matrix(
      log(mixEst[1, ]) + dbeta(xRep, mixEst[2, ], mixEst[3, ], log = TRUE),
      nrow = Nc
    ))
    lnresp <- matrixStats::rowLogSumExps(lli)
    ## ensure that the log-likelihood does not go out of numerical
    ## reasonable bounds
    lli <- apply(lli, 2, pmax, -30)
    ## 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 <- matrixStats::colLogSumExps(lresp)
    zSum <- exp(lzSum)
    mixEst[1, ] <- exp(lzSum - logN)

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

    c1 <- colSums(LxO * resp) / zSum
    c2 <- colSums(LxC * resp) / zSum

    ## now solve for new alpha and beta estimates jointly for each
    ## component
    for (i in 1:Nc) {
      Lest <- optim(mixEstPar[c(2, 3), i], bmm_mun_ml(c1[i], c2[i]))
      if (Lest$convergence != 0) {
        warning(
          "Warning: Component",
          i,
          "in iteration",
          iter,
          "had convergence problems!"
        )
      }
      mixEstPar[c(2, 3), i] <- Lest$par
      mui <- inv_logit(Lest$par[1])
      nmui <- inv_logit(-Lest$par[1])
      ni <- exp(Lest$par[2])
      mixEst[2, i] <- mui * ni
      mixEst[3, i] <- nmui * ni
    }

    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", "EMbmm", "betaMix", "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("betaMix", "mix")
    x
  })
  attr(mixEst, "x") <- x

  mixEst
}