R/estimateHMM.R

In jiangrongo/HMM:

#' estimateHMM
#'
#' \code{estimateHMM} performs EM Algorithm on multiple trajectures of observed chain X
#'
#' Assumes hidden Markov model, the emmision probability follows normal distribution.
#' Detail model parameters are given by the function \code{generateHMM}.
#' Estimation would terminate in a given error tolerance.
#'
#' @export
#' @param X     a vector of observed states
#' @param M     the number of states of Markov chain
#' @param constr a constrain matrix for transition probability:
#'                 \itemize{
#'                    \item \code{constr[i,j]}=0 indicates p(i->j)=0.
#'                    \item \code{constr[i,j]}=1 indicates p(i->j)>0
#'                 }
#'
#' @param tol=0.001     error tolerance.
#' @return      A list containing:
#'                \itemize{
#'                  \item mean vector, \code{u}
#'                  \item standard deviation vector, \code{sig}
#'                  \item transition matrix, \code{trans}
#'                  \item number of iterations used, \code{iter}
#'                }
#'              Print the class: \code{EstConverge}
#'
#' @examples
#' set.seed(1221)
#' df <- generateHMM(num=6,n=100)
#' estimateHMM(df$X, M=3, constr=matrix(1,3,3), tol=0.001)
#' constr2 <- matrix(1,3,3)
#' constr2[1,1] <- 0; constr2[3,3] <- 0;
#' estimateHMM(df$X, M=3, constr=constr2, tol=0.001)

estimateHMM <- function(X, M, constr=matrix(1,M,M), tol=0.001){
  n <- dim(X)[1]     # length of Markov chains
  num <- dim(X)[2]   # number of Markov chains

  u <- matrix(0, M, 1)
  sig <- matrix(0, M, 1)
  # initial estimation of u, sig
  x <- sort(c(X))
  size <- round(n*num/M)
  for (k in 1:(M-1)){
    u[k,1] <- mean(x[(size*(k-1)+1):(size*k)])
    sig[k,1] <- sd(x[(size*(k-1)+1):(size*k)])
  }
  u[M,1] <- mean(x[(size*k+1) : (n*num)])
  sig[M,1] <- sd(x[(size*k+1) : (n*num)])
  # initial estimation transition prob
  trans <- list(diag(1/rowSums(constr)) %*% constr)

  # initialize list for iteration then start EM Alg
  alpha <- rep(list(matrix(0, M,n)), num)
  beta  <- rep(list(matrix(0, M,n)), num)
  L     <- rep(list(matrix(0, M,n)), num)
  sumH  <- rep(list(matrix(0, M,M)), num)

  iter <- 1  # iter : number of iteration of update
  while (1){
    for (i in 1:num){
      alpha[[i]] <- forwardAlg(X[,i], trans[[iter]], u[,iter], sig[,iter])
      beta[[i]]  <- backwardAlg(X[,i], trans[[iter]], u[,iter], sig[,iter])
      L[[i]]     <- computeL(alpha[[i]], beta[[i]])
      sumH[[i]]  <- computeH(alpha[[i]], beta[[i]], X[,i], trans[[iter]], u[,iter], sig[,iter])
    }

    # estimate mean
    numer_u <- matrix(0, M,1)
    denom   <- matrix(0, M,1)
    for (i in 1:num){
      numer_u <- numer_u + L[[i]] %*% as.matrix(X[,i])
      denom <- denom + rowSums(L[[i]])
    }
    u <- cbind(u, numer_u / denom)

    # estimate sigma one by one
    sig <- cbind(sig, 0)
    for (k in 1:M){
      numer_sig <- 0
      for (i in 1:num){
        numer_sig <- numer_sig + sum(L[[i]][k,] * (X[,i] - u[k,iter+1])^2)
      }
      sig[k,iter+1] <- sqrt(numer_sig / denom[k])
    }

    # estimate transition probability
    numer_trans <- matrix(0,M,M)
    denom_trans <- rep(0, M)
    for (i in 1:num){
      numer_trans <- numer_trans + sumH[[i]]
      denom_trans <- denom_trans + rowSums(L[[i]][,-n])
    }
    trans[[iter+1]] <- diag(1/denom_trans) %*% numer_trans

    # converge?
    err <- c(u[,iter+1]-u[,iter],  sig[,iter+1]-sig[,iter], trans[[iter+1]]-trans[[iter]])
    ifelse(prod(err < tol), break, iter <- iter+1)
  }

  # estimate initial probability
  pi0 <- matrix(0, M,num)
  for (i in 1:num){
    for (k in 1:M){
      pi0[k, i] <- L[[i]][[k,1]]
    }
  }
  # Estimate for each path and get the initial prob
  # ie: normalize first. To avoid underflow
  pi0 <- pi0/matrix(rep(colSums(pi0),M), M,num, byrow=TRUE)
  pi0 <- rowSums(pi0)/num

  res <- list(u=u, sig=sig, trans=trans, pi0=pi0, iter=iter)
  class(res) <- c('EstConverge', class(res))
  return(res)
}

#' print.EstConverge
#'
#' \code{print.EstConverge} prints the class EstConverge
#'
#' @export
#'
print.EstConverge <- function(res){
  h <- cbind(res$u[,res$iter+1], res$sig[,res$iter+1], res$trans[[res$iter+1]], res$pi0)
  colnames(h) <- c('u', 'sig', 'trans', rep('', dim(h)[1]-1), 'pi0')
  row.names(h) <- paste('state', 1:dim(h)[1], sep='')
  print(h)
  invisible(h)
  hiter <- res$iter
  names(hiter) <- 'Number of iterations: '
  print(hiter)
}