R/backward.R

In ricky-kotecha/rkHMM:

#' Backward Algorithm
#'
#' Calculate the backward probabilities given 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_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.
#' @param scale_factor Numeric vector of size T with the sclae factor for each point.
#'
#' @return List containing the forward probabilitys and the normalized forward probabilities.
#' @export
backward <- function(p_obs,
                     trn_mtx,
                     scale_factor = NULL) {
  n_states <- ncol(trn_mtx)
  N <- nrow(p_obs)
  p_obs_given_prev_state <- matrix(ncol = n_states, nrow = N)
  norm_p_obs_given_prev_state <- matrix(ncol = n_states, nrow = N)

  next_state <- c(1, 1)
  if (!is.null(scale_factor)) next_state <- next_state * scale_factor[N]

  p_obs_given_prev_state[N, ] <- next_state
  norm_p_obs_given_prev_state[N, ] <- next_state

  for (i in (N - 1):1) {
    p_obs_given_prev_state[i, ] <- trn_mtx %*% diag(p_obs[i + 1, ]) %*% next_state

    if (!is.null(scale_factor)) {
      norm_p_obs_given_prev_state[i, ] <- p_obs_given_prev_state[i, ] * scale_factor[i]
    } else {
      norm_p_obs_given_prev_state[i, ] <- p_obs_given_prev_state[i, ] / sum(p_obs_given_prev_state[i, ])
    }

    next_state <- norm_p_obs_given_prev_state[i, ]
  }
  list(
    backward_p = p_obs_given_prev_state,
    norm_backward_p = norm_p_obs_given_prev_state
  )
}