R/forward.R

In ricky-kotecha/rkHMM:

#' Forward Algorithm
#'
#' Calculate the forward probabilities given an initial state, probability of observations and a transition matrix. I
#' looked through various lecture notes and books but I found wikipedia to be the most helpful:
#'
#' https://en.wikipedia.org/wiki/Forward%E2%80%93backward_algorithm#Forward_probabilities
#'
#'
#' @param p_state_0 Numeric vector containing the probability of each state at time 0.
#' @param p_obs Matrix of size T * N (length of time series * number of states) containing the probability of the
#' event/observation at each time step in every state.
#' @param trn_mtx Square matrix of size N states containing the probability of transitioning states. (i, j) would
#' contain the probability of moving from state i to state j.
#'
#' @return List containing the forward probabilitys and the normalized forward probabilities.
#' @export
forward <- function(p_state_0,
                    trn_mtx,
                    p_obs) {

  N <- ncol(trn_mtx)
  T <- nrow(p_obs)

  trellis <- matrix(nrow = T, ncol = N)
  scaled_trellis <- matrix(nrow = T, ncol = N)

  trellis[1, ] <- matrix(p_state_0, 1, N) %*% diag(p_obs[1, ])
  scaled_trellis[1, ] <- trellis[1, ] / sum(trellis[1, ])
  for (i in 2:T) {
    trellis[i, ] <- scaled_trellis[i - 1, , drop = FALSE] %*% trn_mtx %*% diag(p_obs[i, ])
    scaled_trellis[i, ] <- trellis[i, ] / sum(trellis[i, ])
  }

  list(
    forward_p = trellis,
    norm_forward_p = scaled_trellis
  )
}