R/fasthsmmfit.R

In ziphsmm: Zero-Inflated Poisson Hidden (Semi-)Markov Models

#######################################################
#' Fast gradient descent / stochastic gradient descent algorithm to learn the parameters
#' in a specialized zero-inflated hidden semi-Markov model, where zero-inflation only happens
#' in State 1. And if there were covariates, they could only be the same ones for the
#' state-dependent log Poisson means and the logit structural zero proportion. In addition,
#' the dwell time distributions are nonparametric for all hidden states.
#'
#' @param y observed time series values
#' @param x matrix of covariates for the log poisson means and logit zero proportion.
#' Default to NULL.
#' @param ntimes a vector specifying the lengths of individual, 
#' i.e. independent, time series. If not specified, the responses are assumed to 
#' form a single time series, i.e. ntimes=length(y).
#' @param M number of latent states
#' @param trunc a vector specifying truncation at the maximum number of dwelling time in 
#' each state.
#' @param dt_init a matrix whose i,j th element is the probability of staying in 
#' state i for duration j, which is the nonparametric state duration distributions.
#' @param prior_init a vector of initial values for prior probability for each state
#' @param tpm_init a matrix of initial values for transition probability matrix
#' @param emit_init a vector of initial values for the means for each poisson distribution
#' @param zero_init a scalar initial value for the structural zero proportion 
#' @param yceil a scalar defining the ceiling of y, above which the values will be
#' truncated. Default to NULL.
#' @param stochastic Logical. Should the stochastic gradient descent methods be used.
#' @param nmin a scalar for the minimum number of observations before the first
#' iteration of stochastic gradient descent. Default to 1000.
#' @param nupdate a scalar specifying the total number of updates for stochastic
#' gradient descent. Default to 100.
#' @param power a scalar representing the power of the learning rate, which should lie
#' between (0.5,1]. Default to 0.7
#' @param rate a vector of learning rate in stochastic gradient descent for the logit
#' parameters and log parameters. Default to c(1,0.05).
#' @param method method to be used for direct numeric optimization. See details in
#' the help page for optim() function. Default to Nelder-Mead.
#' @param hessian Logical. Should a numerically differentiated Hessian matrix be returned?
#' Note that the hessian is for the working parameters, which are the generalized logit of prior 
#' probabilities (except for state 1), the generalized logit of the transition probability 
#' matrix(except 1st column), the logit of non-zero zero proportions, and the log of 
#' each state-dependent poisson means
#' @param ... Further arguments passed on to the optimization methods
#' @return the maximum likelihood estimates of the zero-inflated hidden Markov model
#' @references Walter Zucchini, Iain L. MacDonald, Roland Langrock. Hidden Markov Models for 
#' Time Series: An Introduction Using R, Second Edition. Chapman & Hall/CRC
#' @examples
#' 
#' #1. no covariates
#' set.seed(123)
#' prior_init <- c(0.5,0.2,0.3)
#' dt_init <- matrix(c(0.4,0.3,0.2,0.1,0.5,0.2,0.2,0.1,
#'                    0.25,0.25,0.25,0.25),3,4,byrow=TRUE)
#' emit_init <- c(10,50,100)
#' zeroprop <- c(0.6,0,0)
#' omega <- matrix(c(0,0.3,0.7,0.4,0,0.6,0.5,0.5,0),3,3,byrow=TRUE)
#' sim1 <- hsmmsim(n=1000,M=3,prior=prior_init,dt_dist="nonparametric",
#'                dt_parm=dt_init, tpm_parm=omega,
#'                emit_parm=emit_init,zeroprop=zeroprop)
#' str(sim1)
#' y <- sim1$series
#' 
#' 
#' fitnost <- fasthsmmfit(y,NULL,NULL,3,trunc=c(4,4,4),
#'            dt_init=matrix(c(0.3,0.3,0.2,0.2,
#'            0.4,0.2,0.2,0.2,0.25,0.25,0.25,0.25),3,4,byrow=TRUE),
#'            prior_init=c(0.3,0.3,0.4),
#'            tpm_init=matrix(c(0,0.7,0.3,0.5,0,0.5,0.3,0.7,0),3,3,byrow=TRUE), 
#'            emit_init=c(8,40,80),zero_init=0.4,method="BFGS",control=list(trace=1))
#'            
#' decode1 <- hsmmviterbi(y, ntimes=NULL, 3, trunc=c(4,4,4), fitnost$prior, 
#'                       dt_dist="nonparametric", 
#'                       fitnost$dt, fitnost$tpm, 
#'                       fitnost$emit, c(fitnost$zeroprop,0,0))
#' 
#' #2. with covariates
#' dtparm <-  matrix(c(0.4,0.3,0.2,0.1,0.7,0.2,0.1,0,0.25,0.25,0.25,0.25),3,4,byrow=TRUE)
#' priorparm <- c(0,0)
#' zeroindex <- c(1,0,0)
#' zeroparm <- c(0,-1,1)
#' emitparm <- c(2,0.5,-0.5,3,0.3,-0.2,4,0.4,-0.4)
#' tpmparm <- c(-1,0,1)
#' workparm <- c(priorparm,zeroparm,emitparm,tpmparm) 
#' trunc <- c(4,3,4)
#' 
#' designx <- matrix(rnorm(4000),nrow=2000,ncol=2)
#' result <- hsmmsim2(workparm,3,2000,zeroindex,"nonparametric",
#'                   emit_x=designx,zeroinfl_x=designx,dt_x=dtparm)
#' y <- result$series
#' 
#' fitnost <- fasthsmmfit(y,designx,NULL,3,trunc=c(4,3,4),
#'            dt_init=matrix(c(0.3,0.3,0.2,0.2,0.4,0.2,0.2,0.2,
#'            0.25,0.25,0.25,0.25),3,4,byrow=TRUE),
#'            prior_init=c(0.3,0.3,0.4),
#'            tpm_init=matrix(c(0,0.8,0.2,0.4,0,0.6,0.2,0.8,0),3,3,byrow=TRUE), 
#'            emit_init=c(8,40,80),zero_init=0.4,method="BFGS",control=list(trace=1))
#'
#' decode2 <- hsmmviterbi2(y,NULL,3,trunc=c(4,3,4),
#'                        fitnost$working_parm[-(1:8)],dt_x=fitnost$dt,
#'                        dt_dist="nonparametric", zero_init=c(1,0,0),
#'                        emit_x=designx,zeroinfl_x=designx)
#'
#' #3. stochastic gradient descent without covariates
#' prior_init <- c(0.5,0.2,0.3)
#' dt_init <- matrix(c(0.4,0.3,0.2,0.1,0.5,0.2,0.2,0.1,
#'                    0.25,0.25,0.25,0.25),3,4,byrow=TRUE)
#' emit_init <- c(10,50,100)
#' zeroprop <- c(0.6,0,0)
#' omega <- matrix(c(0,0.3,0.7,0.4,0,0.6,0.5,0.5,0),3,3,byrow=TRUE)
#' sim1 <- hsmmsim(n=50000,M=3,prior=prior_init,dt_dist="nonparametric",
#'                dt_parm=dt_init, tpm_parm=omega,
#'                emit_parm=emit_init,zeroprop=zeroprop)
#' y <- sim1$series
#'
#' fitst <- fasthsmmfit(y,NULL,NULL,3,trunc=c(4,4,4),
#'          dt_init=matrix(c(0.4,0.3,0.2,0.1,0.4,0.2,0.2,0.2,
#'          0.25,0.25,0.25,0.25),3,4,byrow=TRUE),
#'          prior_init=c(0.3,0.3,0.4),
#'          tpm_init=matrix(c(0,0.8,0.2,0.5,0,0.5,0.2,0.8,0),3,3,byrow=TRUE), 
#'          emit_init=c(15,40,90),zero_init=0.4,stochastic=TRUE,
#'          nmin=500,nupdate=500,power=0.6,rate=c(0.5,0.08))
#' str(fitst)
#'
#' #4. stochastic descent without covariates
#' dtparm <-  matrix(c(0.4,0.3,0.2,0.1,0.7,0.2,0.1,0,
#'            0.25,0.25,0.25,0.25),3,4,byrow=TRUE)
#' priorparm <- c(0,0)
#' zeroindex <- c(1,0,0)
#' zeroparm <- c(0,-0.5,0.5)
#' emitparm <- c(2,0.1,-0.1,3,0.3,-0.2,4,0.4,-0.4)
#' tpmparm <- c(-1,0,1)
#' workparm <- c(priorparm,zeroparm,emitparm,tpmparm) 
#' trunc <- c(4,3,4)
#' 
#' designx <- matrix(rnorm(100000),nrow=50000,ncol=2)
#' result <- hsmmsim2(workparm,3,50000,zeroindex,"nonparametric",
#'                   emit_x=designx,zeroinfl_x=designx,dt_x=dtparm)
#' 
#' y <- result$series
#'
#' fitst <- fasthsmmfit(y,designx,NULL,3,trunc=c(4,3,4),
#'          dt_init=matrix(c(0.4,0.3,0.2,0.1,0.6,0.3,0.1,0,
#'          0.25,0.25,0.25,0.25),3,4,byrow=TRUE),
#'          prior_init=c(0.3,0.3,0.4),
#'          tpm_init=matrix(c(0,0.8,0.2,0.5,0,0.5,0.2,0.8,0),3,3,byrow=TRUE), 
#'          emit_init=c(15,40,90),zero_init=0.6,stochastic=TRUE,
#'          nmin=500,nupdate=500,power=0.6,rate=c(0.3,0.05))
#' str(fitst)
#' @useDynLib ziphsmm
#' @importFrom Rcpp evalCpp
#' @export
#main function to fit a homogeneous hidden semi-markov model
fasthsmmfit <- function(y, x=NULL, ntimes=NULL, M, trunc,
                        dt_init, prior_init, tpm_init, emit_init, 
                       zero_init,  yceil=NULL,
                       stochastic=FALSE, nmin=1000, nupdate=100, power=0.7, 
                       rate=c(1,0.05),
                       method="Nelder-Mead", hessian=FALSE, ...){
  if(is.null(ntimes)) ntimes <- length(y)
  if(!is.null(yceil)) y <- ifelse(y>yceil, yceil, y)
  
  if(is.null(x)){
    allparm <- rep(NA, sum(trunc)-M + M-1 + M+1 + M*(M-2))
    lastindex <- 0
    for(i in 1:M) {
      allparm[(lastindex+1):(lastindex+trunc[i]-1)] <- glogit(dt_init[i,1:trunc[i]])
      lastindex <- lastindex + trunc[i] - 1
    }
    
    allparm[(lastindex+1):(lastindex+M-1)] <- glogit(prior_init)
    lastindex <- lastindex + M - 1
    
    allparm[lastindex+1] <- log(zero_init) - log(1-zero_init)
    lastindex <- lastindex + 1
    allparm[(lastindex+1):(lastindex+M)] <- log(emit_init)
    lastindex <- lastindex + M
    
    if(M>2){
      for(i in 1:M) {allparm[(lastindex+1):(lastindex+M-2)] <- glogit(tpm_init[i,-i])
      lastindex <- lastindex + M - 2}
    }
    
    
    if(stochastic==FALSE){
      result <- optim(allparm,hsmmnegloglik_nocov,gradhsmmnegloglik_nocov,
                      M=M,y=y,ntimes=ntimes,trunc=trunc,
                      method=method,hessian=hessian,...)
      
      fitparm <- result$par
      negloglik <- result$value
      natparm <- hsmmretrieve_nocov(fitparm,M,trunc)
      return(list(dt=natparm$dt,prior=natparm$delta,tpm=natparm$gamma,
                  zeroprop=natparm$theta,emit=natparm$lambda,
                  working_parm=fitparm,negloglik=negloglik))
    }else{
      
      ratevec <- c(rep(rate[1],sum(trunc)-M+M-1+1),
                   rep(rate[2],M),rep(rate[1],M*(M-2)))
      fitparm <- sgd_hsmm_nocov(allparm,y,M,ntimes,nmin,nupdate,power,
                               ratevec,trunc)
      natparm <- hsmmretrieve_nocov(fitparm[nupdate,],M,trunc)
      return(list(dt=natparm$dt,prior=natparm$delta,tpm=natparm$gamma,
                  zeroprop=natparm$theta,emit=natparm$lambda,
                  working_parm=fitparm))
    }
  }else{
    x <- cbind(1,x)
    ncolx <- ncol(x)
    allparm <- rep(NA, sum(trunc)-M+M-1+ncolx*(1+M)+M*(M-2))
    
    lastindex <- 0
    for(i in 1:M) {
      allparm[(lastindex+1):(lastindex+trunc[i]-1)] <- glogit(dt_init[i,1:trunc[i]])
      lastindex <- lastindex + trunc[i] - 1
    }
    
    allparm[(lastindex+1):(lastindex+M-1)] <- glogit(prior_init)
    lastindex <- lastindex + M - 1
    
    allparm[(lastindex+1):(lastindex+ncolx)] <- 
      c(log(zero_init)-log(1-zero_init),rep(0,ncolx-1))
    lastindex <- lastindex + ncolx
    
    for(i in 1:M){
      allparm[(lastindex+1):(lastindex+ncolx)] <- 
        c(log(emit_init[i]),rep(0,ncolx-1))
      lastindex <- lastindex + ncolx
    }
    
    if(M>2){
      for(i in 1:M) {allparm[(lastindex+1):(lastindex+M-2)] <- glogit(tpm_init[i,-i])
      lastindex <- lastindex + M - 2}
    }
    
    if(stochastic==FALSE){
      result <- optim(allparm,hsmmnegloglik_cov,gradhsmmnegloglik_cov,
                      M=M,y=y,covariates=x,ntimes=ntimes,trunc=trunc,
                      method=method,hessian=hessian,...)
      
      fitparm <- result$par
      negloglik <- result$value
      natparm <- hsmmretrieve_cov(fitparm,M,ncolx,trunc)
      return(list(dt=natparm$dt,prior=natparm$delta,tpm=natparm$gamma,
                  zeroparm=natparm$thetaparm,emitparm=natparm$lambdaparm,
                  working_parm=fitparm,negloglik=negloglik))
    }else{
      ratevec <- c(rep(rate[1],sum(trunc)-M+M-1+ncolx),
                   rep(rate[2],M*ncolx),rep(rate[1],M*(M-2)))
      fitparm <- sgd_hsmm_cov(allparm,y,x,M,ntimes,nmin,nupdate,power,
                             ratevec,trunc)
      natparm <- hsmmretrieve_cov(fitparm[nupdate,],M,ncolx,trunc)
      return(list(dt=natparm$dt,prior=natparm$delta,tpm=natparm$gamma,
                  zeroparm=natparm$thetaparm,emitparm=natparm$lambdaparm,
                  working_parm=fitparm))
    }
  }
}