R/plotSeries.R

In luisdamiano/BayesHMM: Full Bayesian Inference for Hidden Markov Models

#' Plot the observation series along with many other customizable options.
#'
#' @param fit An object returned by either \code{\link{draw_samples}} or \code{\link{optimizing}}.
#' @param r An optional numeric vector with the numbers of the dimensions of the observation vector to be plotted. It defaults to plotting all variables.
#' @param features An optional vector of character strings with the names of the features to include in the plot. These may be \emph{stateShade} (shaded background), \emph{yColoredLine} (colored observation line), \emph{yColoredDots} (colored observation dots), \emph{bottomColoredMarks} (colored marks on the lower axis), \emph{topColoredMarks} (colored marks on the top axis). All colors are based on the hidden state assigned to each time step using the \emph{stateProbability} quantity. Note that more than one feature may be plotted at the same time. It allows for partial matching. It defaults to none.
#' @param stateProbability An optional character string with the name of the quantity used to classify the observations (either filtered, smoothed, or viterbi). It allows for partial matching. It defaults to the smoothed probability.
#' @param stateProbabilityFunction An optional function applied to the samples corresponding to one time step \emph{t}, one hidden state \emph{k}. The observation at time step \emph{t} is then assigned to the hidden state with largest value. Note that the user needs to supply a function as an argument, and not a character string with the name of the function. It defaults to the \code{\link{median}} function.
#' @param chain An optional integer number between 1 and the number of chains M. Only the samples generated by the selected chain are considered. It defaults to the first chain.
#' @param main An optional character string with the overall title for the plot.
#' @param xlab An optional character string with the title for the horizontal axis.
#' @param ylab An optional character string with the title for the vertical axis.
#' @param addLegend An optional boolean indicating whether a legend should be included.
#' @param legend.cex A numerical value giving the amount by which plotting legend text and symbols should be magnified relative to the default. It defaults to one.
#' @param ... Arguments to be passed to the \code{\link{plot}} function.
#' @family visualization functions
#' @examples
#' \dontrun{
#' plot_series(
#'   myFit,
#'   features = c("stateShade", "yColoredLine", "bottomColoredMarks")
#' )
#' }
plot_series <- function(fit, r = NULL,
                        features = NULL, stateProbability = "smoothed",
                        stateProbabilityFunction = median, chain = 1,
                        main = NULL, xlab = NULL, ylab = NULL,
                        addLegend = TRUE, legend.cex = 1, ...) {

  theme <- getOption("BayesHMM.theme")

  stateProbability <- match.arg(
    stateProbability,
    choices = c("filtered", "smoothed", "viterbi"),
    several.ok	= FALSE
  )

  features <- match.arg(
    features,
    choices = c(
      "stateShade", "yColoredLine", "yColoredDots",
      "bottomColoredMarks", "topColoredMarks"
    ),
    several.ok	= TRUE
  )

  y    <- extract_y(fit)
  tidx <- seq_len(get_dim(y)[1])
  zFun <- switch(
    stateProbability,
    filtered = classify_alpha,
    smoothed = classify_gamma,
    viterbi  = classify_zstar
  )
  z    <- zFun(fit, reduce = stateProbabilityFunction, chain = chain)
  K    <- extract_K(fit)
  R    <- extract_R(fit)
  rSeq <- if (is.null(r)) { seq_len(R) } else { r }

  opar <- par(
    oma   = c(2, 0, 0, 0),
    mfrow = rev(n2mfrow(length(rSeq)))
  )
  on.exit(par(opar))

  for (r in rSeq) {
    plot(
      NULL,
      main = if (is.null(main)) { "" } else { main },
      xlab = if (is.null(xlab)) { "Time" } else { xlab },
      ylab =
        if (is.null(ylab)) {
          colnames(y)[r]
        } else {
          ylab
        },
      xlim = quantile(tidx, probs = c(0, 1)),
      ylim = quantile(y, probs = c(0, 1)),
      ...
    )

    lines(x = tidx, y = y[, r], col = theme$seriesY)
    add_features(tidx, y[, r], z, features = features)
  }

  add_legend_overlay(
    x      = "bottom",
    legend = sprintf("Hidden state %d", 1:K),
    bty    = "n",
    horiz  = TRUE,
    fill   = theme$states[1:K],
    border = theme$states[1:K],
    cex    = legend.cex
  )
}

#' Plot the estimated hidden path along with many other customizable options.
#'
#' @inherit plot_series
#' @param stateProbability An optional character string with the name of the quantity used to classify the observations (either filtered or smoothed). It allows for partial matching. It defaults to the smoothed probability.
#' @param features An optional vector of character strings with the names of the features to include in the plot. These may be \emph{stateShade} (shaded background), \emph{probabilityColoredLine} (colored probability line), \emph{probabilityColoredDots} (colored probability dots), \emph{probabilityFan} (shaded probability fan for posterior intervals), \emph{bottomColoredMarks} (colored marks on the lower axis), \emph{topColoredMarks} (colored marks on the top axis). All colors are based on the hidden state assigned to each time step using the \emph{stateProbability} quantity. Note that more than one feature may be plotted at the same time. It allows for partial matching. It defaults to none.
#' @param stateProbabilityInterval An optional function returning the upper and lower bound of the posterior intervals, such as \code{\link{posterior_intervals}}. It defaults to `posterior_intervals(c(0.1, 0.9))`.
#' @family visualization functions
#' @examples
#' \dontrun{
#' plot_state_probability(
#'   myFit, stateProbability = "filtered",
#'   features = c("stateShade", "probabilityFan", "bottomColoredMarks"),
#'   chain = 2
#' )
#' }
plot_state_probability <- function(fit, stateProbability = "smoothed",
                                   features = NULL, stateProbabilityFunction = median,
                                   stateProbabilityInterval = posterior_intervals(c(0.1, 0.9)),
                                   chain = 1, main = NULL, xlab = NULL,
                                   ylab = NULL, addLegend = TRUE, legend.cex = 1, ...) {

  theme <- getOption("BayesHMM.theme")

  stateProbability <- match.arg(
    stateProbability,
    choices = c("filtered", "smoothed"),
    several.ok	= FALSE
  )

  features <- match.arg(
    features,
    choices = c(
      "probabilityColoredLine", "probabilityColoredDots", "probabilityFan",
      "stateShade", "bottomColoredMarks", "topColoredMarks"
    ),
    several.ok	= TRUE
  )

  zFun <- switch(
    stateProbability,
    filtered = classify_alpha,
    smoothed = classify_gamma,
    viterbi  = classify_zstar
  )
  z    <- zFun(fit, reduce = stateProbabilityFunction, chain = chain)

  zFun <- switch(
    stateProbability,
    filtered = extract_alpha,
    smoothed = extract_gamma,
    viterbi  = extract_zstar
  )
  p    <- zFun(fit, reduce = stateProbabilityFunction, chain = chain)
  pInt <- zFun(fit, reduce = stateProbabilityInterval, chain = chain)
  tidx <- seq_len(get_dim(z)[1])
  K    <- extract_K(fit)

  opar <- par(no.readonly = TRUE)
  on.exit(par(opar))

  par(
    oma   = c(6, 0, 0, 0),
    mfrow = c(K, 1)
  )
  for (k in 1:K) {
    par(
      mar = par_edit(
        par()$mar,
        if (k != K) { 0 },
        NULL,
        if (k != 1) { 0 },
        NULL
      )
    )

    plot(
      x    = tidx,
      y    = p[k, ],
      main = if (k == 1) { main },
      xlab = if (k == K) { xlab } else { "" },
      ylab = sprintf("Hidden state %d", k),
      type = "l",
      col  = theme$states[k],
      xaxt = "n",
      yaxt = "n",
      ...
    )

    add_features(
      x = tidx, z = z, p = p[k, ],
      pInt = pInt[, k, ], k = k, features = features
    )

    axis(if (k %% 2) { 4 } else { 2 })

    if (k == K) { axis(1) }
  }

  if (addLegend) {
    add_legend_overlay(
      x      = "bottom",
      legend = sprintf("Hidden state %d", 1:K),
      bty    = "n",
      horiz  = TRUE,
      fill   = theme$states[1:K],
      border = theme$states[1:K],
      cex    = legend.cex
    )
  }
}