R/viterbi.R

In ricky-kotecha/rkHMM:

viterbi <- function(states,
                    A,
                    prior,
                    factor_density_array,
                    na.allow = TRUE) {
  # nt <- sum(object@ntimes)
  # lt <- length(object@ntimes)
  # et <- cumsum(object@ntimes)
  # bt <- c(1,et[-lt]+1)
  # ns <- object@nstates
  nt <- NROW(factor_density_array[ , 1, ])
  lt <- 1
  et <- NROW(factor_density_array[ , 1, ])
  bt <- 1

  ns <- states

  delta <- psi <- matrix(nrow=nt,ncol=ns)
  state <- vector(length=nt)

  B <- factor_density_array
  # B <- object@dens
  if(na.allow) B <- replace(B,is.na(B),1)
  B <- apply(B,c(1,3),prod)

  # initialization
  case <- 1
  delta[bt[case],] <- prior[case,]*B[bt[case],]
  delta[bt[case],] <- delta[bt[case],]/(sum(delta[bt[case],]))
  psi[bt[case],] <- 0
  # recursion
  if(nt[case]>1) {
    for(tt in ((bt[case]+1):et[case])) {
      for(j in 1:ns) {
        delta[tt, j] <- max(delta[tt-1,] * (A[1,j,])) * B[tt,j]
        k <- which.max(delta[tt-1,]*A[1,j,])
        if(length(k) == 0) k <- 0 # what's this doing here??? can this ever occur? FIX ME
        psi[tt,j] <- k
      }
      delta[tt,] <- delta[tt,]/(sum(delta[tt,]))
    }
  }

  # trace maximum likely state
  state[et[case]] <- which.max(delta[et[case],])

  # this doesn't need a for loop does it???? FIX ME
  if (nt > 1) {
    for (i in (et[case]-1):bt[case]) {
      state[i] <- psi[i+1,state[i+1]]
    }
  }

  colnames(delta) <- paste("S",1:ns,sep="")

  data.frame(state, delta)
}