R/generateHMM.R

In jiangrongo/HMM:

#' generateHMM
#'
#' \code{generateHMM} generates multiple trajectures of observed chain X according to the hidden Markov model.
#'
#' Assumes hidden Markov model, the emmision probability follows normal distribution.
#'     \itemize{
#'         \item  transition probability is specified by user via \code{trans}.
#'         Default value: [0, 3, 7; 5, 2, 3; 5, 5, 0]/10
#'         \item  emission prob: mean is specified by user via \code{u}.
#'         It must match the dimension of \code{trans}.
#'         Defualt vaule \code{u} = [0,5,-5].
#'         \item  emission prob: sd is specified by user via \code{sig}.
#'         It must match the dimension of \code{trans}.
#'         Defualt vaule \code{sig}=[1,1,1].
#'         \item  initial probability. Default value \code{pi0}=[1,1,1]/3
#'     }
#'
#' @export
#' @param num   number of Markov chains
#' @param n     length of Markov chain (assume all chains have equal length)
#' @param trans a matrix of transition probability
#' @param u     emission prob: mean. For the rest of arguments, the dimension must match with \code{trans}.
#' @param sig   emission prob: standard deviation.
#' @param pi0   initial probability.
#'
#' @return      A list containing:
#'                \itemize{
#'                  \item matrix of hidden state, Z. Each column is a Markov chain.
#'                  \item matrix of observed states, X. (each column corresponds to a column of Z)
#'                  \item pi0, the initial probability
#'                  \item \code{trans}
#'                  \item \code{u}, \code{sig}
#'                }
#'
#' @examples
#' set.seed(1221)
#' df <- generateHMM(num=2,n=10); df


generateHMM <- function(num, n=100,
                        trans=matrix(c(0,3,7,
                                       5,2,3,
                                       5,5,0),
                                     3,3, byrow = TRUE)/10,
                        u=c(-5,0,5),
                        sig=c(1,1,1),
                        pi0=c(1,1,1)/3 ){
  M <- dim(trans)[1]
  Z <- matrix(0, n, num)  # hidden variable
  X <- matrix(0, n, num)  # observed variable

  # simulate num trajectories
  for (j in 1:num){
    Z[1,j] <- sample(1:M, size=1, prob=pi0)
    X[1,j] <- rnorm(1, mean=u[Z[1,j]], sd=sig[Z[1,j]])
    for (k in 2:n){
      Z[k, j] <- sample(1:M, size=1, prob=trans[Z[k-1, j],] )
      X[k, j] <- rnorm(1, mean=u[Z[k, j]], sd=sig[Z[k, j]])
    }
  }

  return(list(Z=Z, X=X, pi0=pi0, trans=trans, u=u, sig=sig))
}