R/forward_backward.R

In ricky-kotecha/rkHMM:

#' Forward Backward Algorithm
#'
#' Performs on iteration of the forward backward algorithm give an initial state, transition matrix and probaility of
#' each observation.
#'
#' @param p_state_0 Vector of length n_states containing the intial probability of being in each state.
#' @param trn_mtx Square matrix. Rows represent from states and columns represent to states.
#' @param p_obs Matrix n_time_steps by n_variables containing probability of each observation.
#'
#' @return List.
#' @export
forward_backward <- function(p_state_0, trn_mtx, p_obs) {
  fwd <- forward(
    p_state_0 = p_state_0,
    trn_mtx = trn_mtx,
    p_obs = p_obs
  )

  bwd <- backward(
    trn_mtx = trn_mtx,
    p_obs = p_obs,
    scale_factor = 1 / rowSums(fwd$forward_p)
  )

  gamma <- fwd$norm_forward_p * bwd$norm_backward_p / (1 / rowSums(fwd$forward_p))

  #  Again I looked in various places but found wikipedia to have the best explanation:
  #  https://en.wikipedia.org/wiki/Baum%E2%80%93Welch_algorithm#Update
  #
  #  Probability of being in state i and j at times t and t + 1 respectively
  #  given the observed sequence and parameters theta
  #
  #  In this case the observed sequence is joint_p_obs_per_state i.e. the probability of the observations.
  #
  #  My interpretation: Every t to t + 1 is a transition. At each of those points we have a sequence of forwards and
  #  backwards probabilities. Xi effective "connects" these with our current transition matrix and the p of the
  #  observation
  #
  #  Note we have already been normalising so we don't need to calculate the denominators
  N <- nrow(fwd$norm_forward_p)
  xi <- array(dim = c(N, 2, 2))
  for (t in 1:(N - 1)) {
    xi[t, , ] <- matrix(fwd$norm_forward_p[t, ], 2, 2, byrow = TRUE) *
      t(trn_mtx) *
      p_obs[t + 1, ] *
      bwd$norm_backward_p[t + 1, ]
  }

  list(
    alpha = fwd$norm_forward_p,
    beta = bwd$norm_backward_p,
    logLike = -sum(log(1 / rowSums(fwd$forward_p))),
    gamma = gamma,
    xi = xi
  )
}