R/plot_dllm.R

In dlmwwbe: Dynamic Linear Model for Wastewater-Based Epidemiology

#' Plot a Fitted Dynamic Local Level Model
#'
#' Produces a plot for an object of class \code{dllm} (typically created by \code{dllm}).
#' The function displays the observed data along with the fitted curves computed using filtered and/or
#' smoothed state estimates.
#'
#' @param x An object of class \code{dllm}, as returned by \code{\link[dlm]{dlm}}.
#' @param type Character; one of \code{"smoother"} or \code{"filter"} (default: \code{"smoother"}).
#'   Specifies which fitted curves to display.
#' @param plot_data Logical; if \code{TRUE} (default) the observed data points are plotted.
#' @param obs_cols Character; an optional argument specifying the variables to be plotted. If \code{NULL},
#'    plot all variables.
#' @param obs_colors Optional character vector specifying custom colors for observed data.
#'   If not supplied, a default palette is used.
#' @param filter_colors Optional character vector specifying custom colors for filtered curves.
#'   If not supplied, a default palette is used.
#' @param smoother_colors Optional character vector specifying custom colors for smoothed curves.
#'   If not supplied, a default palette is used.
#' @param conf.int Logical; if \code{TRUE}, plot confidence intervals with the given sig.level.
#' @param sig.level Numeric; significance level for confidence intervals (default: 0.95).
#' @param ... Additional graphical parameters to pass to the underlying plotting functions.
#'
#' @return This function produces a plot of the fitted DLM and returns \code{NULL} invisibly.
#'
#'
#' @importFrom graphics legend matlines matplot par
#' @importFrom utils modifyList
#' @importFrom stats qnorm
#' @importFrom graphics polygon
#' @importFrom grDevices adjustcolor
#' @export

plot.dllm <- function(x,
                      type=c("smoother", "filter"),
                      plot_data=TRUE,
                      obs_cols=NULL,
                      obs_colors=NULL,
                      filter_colors=NULL,
                      smoother_colors=NULL,
                      conf.int=FALSE,
                      sig.level=0.95,
                         ...) {
  if (!inherits(x, "dllm"))
    stop("The object 'x' must be of class 'dllm'. Ensure it was created using dllm().")

  valid_types <- c("smoother", "filter")
  if (!(type %in% valid_types)) {
    stop(sprintf("`type` must be one of: %s", paste(valid_types, collapse = ", ")))
  }


  if (!isTRUE(all(obs_cols %in% x$obs_cols))){
    invalid <- setdiff(obs_cols, x$obs_cols)
    stop("Invalid column names: ", paste(invalid, collapse=','))
  }

  if(is.null(obs_cols)){
    obs_cols <- x$obs_cols
  }

  if(sig.level<=0 | sig.level>=1)
    stop('Please provide a valid significant level!')

  obs <- x$data[, obs_cols, drop = FALSE]



  if(x$S=='univariate'){
    filters <- x$yf[, 1, drop=FALSE]
    var.filter <- x$var.filter[, 1, drop=FALSE]
    var.smoother <- x$var.smoother[, 1, drop=FALSE]
    colnames(filters) <- 'filter'
    smoothers <- x$ys[, 1, drop=FALSE]
    colnames(smoothers) <- 'smoother'
  }else{
    var.filter <- x$var.filter[, paste0(obs_cols, '_filter'), drop=FALSE]
    var.smoother <- x$var.smoother[, paste0(obs_cols, '_smoother'), drop=FALSE]
    filters <- x$yf[, paste0(obs_cols, '_filter'), drop=FALSE]
    smoothers <- x$ys[, paste0(obs_cols, '_smoother'), drop=FALSE]}

  if (x$log10){
    filters.temp <- log(filters, base=10)
    smoothers.temp <- log(smoothers, base=10)
    if(x$S=='kvariate'){
      filters <- 10**apply(filters.temp, MARGIN=1, FUN=mean)
      smoothers <- 10**apply(smoothers.temp, MARGIN=1, FUN=mean)
    }
  }

  if(x$S=='univariate'){
    upper_filter <- filters.temp + qnorm(p=sig.level, sd=sqrt(var.filter))
    lower_filter <- filters.temp - qnorm(p=sig.level, sd=sqrt(var.filter))
    upper_smoother <- smoothers.temp + qnorm(p=sig.level, sd=sqrt(var.smoother))
    lower_smoother <- smoothers.temp - qnorm(p=sig.level, sd=sqrt(var.smoother))
  }else{
    k <- ncol(filters.temp)
    upper_filter <- apply(X=filters.temp, MARGIN=1, FUN=mean) +
                    qnorm(p=sig.level, sd=sqrt(apply(var.filter, MARGIN=1, FUN=mean)/k))
    lower_filter <- apply(X=filters.temp, MARGIN=1, FUN=mean) -
      qnorm(p=sig.level, sd=sqrt(apply(var.filter, MARGIN=1, FUN=mean)/k))
    upper_smoother <- apply(X=smoothers.temp, MARGIN=1, FUN=mean) +
      qnorm(p=sig.level, sd=sqrt(apply(var.smoother, MARGIN=1, FUN=mean)/k))
    lower_smoother <- apply(X=smoothers.temp, MARGIN=1, FUN=mean) -
      qnorm(p=sig.level, sd=sqrt(apply(var.smoother, MARGIN=1, FUN=mean)/k))}

  if (x$log10){
    upper_filter <- 10**upper_filter
    lower_filter <- 10**lower_filter
    upper_smoother <- 10**upper_smoother
    lower_smoother <- 10**lower_smoother
  }


  time_index <- if (!is.null(x$date) && x$date %in% names(x$data)) {
    x$data[[x$date]]
  } else {
    seq_len(nrow(obs))
  }

  ncols <- ncol(obs)

  # Save and restore graphical parameters
  old_par <- par(no.readonly = TRUE)
  on.exit(par(old_par))
  par(mar = c(8, 4, 4, 2))  # Increase bottom margin for legends

  if (type=='filter'){
    plot_range <- range(obs, filters, na.rm = TRUE)
    if(conf.int){
      plot_range <- range(plot_range, upper_filter, lower_filter, na.rm = TRUE)
    }
  }else{
    plot_range <- range(obs, smoothers, na.rm = TRUE)
    if(conf.int){
      plot_range <- range(plot_range, upper_smoother, lower_smoother, na.rm = TRUE)
    }
  }
  default_args <- list(
    xlab = '',
    ylab = "Value",
    ylim = plot_range
  )
  user_args <- list(...)
  plot_args <- modifyList(default_args, user_args)

  default_obs_colors <- rep(x=c("#E69F00", "#56B4E9", "#009E73", "#0072B2", "#D55E00", "#CC79A7"),
                    length.out=ncols)

  if (is.null(obs_colors)){
    obs_colors <- default_obs_colors
  }else{
    if(length(obs_colors) != ncols)
      stop(sprintf("Required %d colors for the observed data, but received %d.",
                   ncols, length(obs_colors)))}

  if (is.null(filter_colors)){
    filter_colors <- 'red'
  }else{
    if(length(filter_colors) != ncols)
      stop(sprintf("Required %d colors for the filtered data, but received %d.",
                   ncols, length(filter_colors)))}

  if (is.null(smoother_colors)){
    smoother_colors <- 'red'
  }else{
    if(length(smoother_colors) != ncols)
      stop(sprintf("Required %d colors for the smoothed data, but received %d.",
                   ncols, length(smoother_colors)))}

  if (type=='filter'){
    main_text <- 'DLLM Filter'
  }else{
    main_text <- 'DLLM Smoother'
  }

  if (plot_data) {
    do.call(matplot, c(list(x = time_index, y = obs, type = "p", pch = 16, main=main_text,
                            cex = 0.75, col = obs_colors), plot_args))
  } else {
    do.call(plot, c(list(x = time_index, y = obs[, 1], type = "n", main=main_text), plot_args))
  }
  if (type=='filter'){
    do.call(matlines, c(list(x = time_index, y = filters, lwd = 2,
                            cex = 0.75, col = 'red', lty='solid'), plot_args))
    if (conf.int){
      polygon(c(time_index, rev(time_index)),
              c(lower_filter, rev(upper_filter)),
              col = adjustcolor('red', alpha.f = 0.2),
              border = NA)
    }
    #legend("topright",
    #       legend = c("filter"),
    #       lty = c('solid'),
    #       col = "black",
    #       bty = "n")
  }else if (type=='smoother'){


    do.call(matlines, c(list(x = time_index, y = smoothers, lwd = 2,
                           cex = 0.75, col = 'red', lty='solid'), plot_args))

    if (conf.int){
      polygon(c(time_index, rev(time_index)),
              c(lower_smoother, rev(upper_smoother)),
              col = adjustcolor('red', alpha.f = 0.2),
              border = NA)
    }
    #legend("topright",
    #       legend = c("smoother"),
    #       lty = c('solid'),
    #       col = "black",
    #       bty = "n")
  }
  #else{
  #  do.call(matlines, c(list(x = time_index, y = filters, lwd = 2,
  #                          cex = 0.75, col = filter_colors, lty='solid'), plot_args))
  #  do.call(matlines, c(list(x = time_index, y = smoothers, lwd = 2,
  #                          cex = 0.75, col = smoother_colors, lty='solid'), plot_args))
  #  legend("topright",
  #         legend = c("filter", "smoother"),
  #         lty = 'solid', #c('dotted', 'solid'),
  #         col = c('blue', 'red'),
  #         bty = "n")
  #}

  legend('topright',
         #x=mean(par("usr")[1:2]),
         #y=place_legend_below_xlab(-1),
         legend=obs_cols,
         col = obs_colors,
         cex=0.75,
         pch=16,
         xjust = 0.5,
         yjust = 0,
         xpd=TRUE,
         horiz = FALSE,
         bty='n',
         lty=NA,
         lwd=2)


  invisible(NULL)
}