R/plot.str.R

In stR: Seasonal Trend Decomposition Using Regression

#' @importFrom graphics Axis
#' @importFrom graphics abline
#' @importFrom graphics layout
#' @importFrom graphics legend
#' @importFrom graphics lines
#' @importFrom graphics mtext
#' @importFrom graphics par
#' @importFrom graphics plot
#' @importFrom graphics plot.new
#' @importFrom graphics polygon

getLimits <- function(l) {
  d <- lapply(l, FUN = function(x) c(x$data, x$upper, x$lower))
  return(c(min(unlist(d), na.rm = TRUE), max(unlist(d), na.rm = TRUE)))
}

getYlab <- function(l) {
  names <- unlist(lapply(l, FUN = function(x) x$name))
  return(do.call("paste", c(as.list(names), sep = ", ")))
}

getLegend <- function(z) {
  x <- z$input
  l0 <- NULL
  if (!is.null(z$cvMSE)) {
    l0 <- paste0("LOO MSE = ", signif(z$cvMSE, 4), ", ")
  }
  if (!is.null(z$optim.CV.MSE)) {
    l0 <- paste0(l0, z$nFold, " fold ", z$gapCV, " gap MSE = ", signif(z$optim.CV.MSE, 4), ", ")
  }
  if (!is.null(z$optim.CV.MAE)) {
    l0 <- paste0(l0, z$nFold, " fold ", z$gapCV, " gap MAE = ", signif(z$optim.CV.MAE, 4), ", ")
  }

  if (length(x$lambdas) > 0) {
    l <- "Lambdas ="
    for (p in x$lambdas) {
      l <- paste(l, " (", paste(signif(p$lambdas, 3), collapse = ","), ")", sep = "")
    }
  } else {
    l <- NULL
  }

  return(paste0(l0, l))
}

createLayoutMatrix <- function(dataPanels, predictorPanels, randomPanels, forecastPanels, rHeader = 1, r = 3) {
  nPanels <- max(unlist(predictorPanels), dataPanels, randomPanels, forecastPanels)
  i <- 1
  v <- rep(i, rHeader)
  i <- i + 1
  for (k in seq_len(nPanels)) {
    v <- c(v, rep(i, r))
    i <- i + 1
  }
  return(as.matrix(v))
}

getDataToPlot <- function(scr, x, dataPanels, predictorPanels, randomPanels, forecastPanels, dataColor, predictorColors, randomColor, forecastColor) {
  toPlot <- list()
  j <- 1
  for (k in seq_along(predictorPanels)) {
    if (scr %in% predictorPanels[[k]] && !is.null(x$output$predictors[[k]]$upper)) {
      for (l in rev(seq_len(ncol(x$output$predictors[[k]]$upper)))) {
        toPlot[[j]] <- list(
          upper = x$output$predictors[[k]]$upper[, l],
          lower = x$output$predictors[[k]]$lower[, l],
          type = "polygon", col = "grey", border = "darkgrey", name = NULL
        )
        j <- j + 1
      }
    }
  }
  if (scr %in% forecastPanels && !is.null(x$output$forecast$upper)) {
    for (l in rev(seq_len(ncol(x$output$forecast$upper)))) {
      toPlot[[j]] <- list(
        upper = x$output$forecast$upper[, l],
        lower = x$output$forecast$lower[, l],
        type = "polygon", col = "grey", border = "darkgrey", name = NULL
      )
      j <- j + 1
    }
  }
  if (scr %in% dataPanels) {
    toPlot[[j]] <- list(data = x$input$data, type = "l", col = dataColor, name = "Observed")
    j <- j + 1
  }
  if (scr %in% randomPanels) {
    toPlot[[j]] <- list(data = x$output$random$data, type = "h", col = randomColor, name = "Random")
    j <- j + 1
  }
  for (k in seq_along(predictorPanels)) {
    if (scr %in% predictorPanels[[k]]) {
      toPlot[[j]] <- list(data = x$output$predictors[[k]]$data, type = "l", col = predictorColors[k], name = x$input$predictors[[k]]$name)
      j <- j + 1
    }
  }
  if (scr %in% forecastPanels) {
    toPlot[[j]] <- list(data = x$output$forecast$data, type = "l", col = forecastColor, name = "Fit/Forecast")
    j <- j + 1
  }
  return(toPlot)
}

#' @name plot.STR
#' @rdname plot.STR
#'
#' @title Plots the results of decomposition
#' @description \code{plot.STR} plots results of STR decomposition.
#' @seealso \code{\link{STRmodel}}, \code{\link{RSTRmodel}}, \code{\link{STR}}, \code{\link{AutoSTR}}
#' @param x Result of STR decomposition.
#' @param xTime Times for data to plot.
#' @param dataPanels Vector of panel numbers in which to plot the original data. Set to \code{NULL} to not show data.
#' @param predictorPanels A list of vectors of numbers where every such vector describes which panels should be used for plotting the corresponding predictor.
#' @param randomPanels Vector of panel numbers in which to plot the residuals.  Set to \code{NULL} to not show residuals.
#' @param forecastPanels Vector of panel numbers in which to plot the fit/forecast.  Set to \code{NULL} to not show forecasts.
#' @param dataColor Color to plot data.
#' @param predictorColors Vector of colors to plot components corresponding to the predictors.
#' @param randomColor Color to plot the residuals.
#' @param forecastColor Color to plot the fit/forecast.
#' @param vLines Vector of times where vertical lines will be plotted.
#' @param xlab Label for horizontal axis.
#' @param main Main heading for plot.
#' @param showLegend When \code{TRUE} (default) legend is shown at top of plot.
#' @param ... Other parameters to be passed directly to plot and lines functions in the implementation.
#' @author Alexander Dokumentov
#' @examples
#' \donttest{
#' fit <- AutoSTR(log(grocery))
#' plot(fit, forecastPanels = 0, randomColor = "DarkGreen", vLines = 2000:2010, lwd = 2)
#' }
#' @method plot STR
#' @export

plot.STR <- function(x, xTime = NULL, dataPanels = 1,
                     predictorPanels = as.list(seq_along(x$output$predictors)),
                     randomPanels = length(x$output$predictors) + 1,
                     forecastPanels = length(x$output$predictors) + 2,
                     dataColor = "black",
                     predictorColors = rep("red", length(x$output$predictors)),
                     randomColor = "red",
                     forecastColor = "blue",
                     vLines = NULL,
                     xlab = "Time",
                     main = ifelse(x$method %in% c("STR", "STRmodel"), "STR decomposition", "Robust STR decomposition"),
                     showLegend = TRUE, ...) {
  if (is.null(xTime)) {
    xTime <- as.vector(time(x$input$data))
  }
  if (length(x$input$data) != length(xTime)) {
    stop("Lengths of x and xTime should be same.")
  }
  op <- par(no.readonly = TRUE) # Resets parameters to the default state
  on.exit(par(op))

  lm <- createLayoutMatrix(dataPanels, predictorPanels, randomPanels, forecastPanels)
  layout(lm)
  par(mar = c(0, 4, 0, 0.5), oma = c(4.5, 0, 2, 0))
  plot.new()
  if (showLegend) {
    legendtext <- getLegend(x)
    if (!is.null(legendtext)) {
      legend("topleft", horiz = FALSE, bty = "n", legend = legendtext)
    }
  }

  nPanels <- max(unlist(predictorPanels), dataPanels, randomPanels, forecastPanels)
  for (scr in 1:nPanels) {
    toPlot <- getDataToPlot(scr, x, dataPanels, predictorPanels, randomPanels, forecastPanels, dataColor, predictorColors, randomColor, forecastColor)
    ylim <- getLimits(toPlot)
    ylab <- getYlab(toPlot)
    plot(xTime, x$input$data, ylab = ylab, type = "n", ylim = ylim, xaxt = "n", ...)
    Axis(x = xTime, side = 1, labels = scr == nPanels)
    abline(h = 0, col = "grey")
    if (!is.null(vLines)) {
      abline(v = vLines, col = "grey", ...)
    }
    for (p in toPlot) {
      if (p$type == "polygon") {
        polygon(c(xTime, rev(xTime)), c(p$upper, rev(p$lower)), col = p$col, border = p$border)
      } else {
        lines(xTime, p$data, col = p$col, type = p$type, ...)
      }
    }
  }
  mtext(xlab, side = 1, outer = TRUE, line = 2.5, cex = 0.9)
  mtext(main, side = 3, outer = TRUE)
}