R/plot_pdlm.R

In dlmwwbe: Dynamic Linear Model for Wastewater-Based Epidemiology

#' Plot a Fitted Predictive Dynamic Linear Model
#'
#' Produces a plot for an object of class \code{pdlm} (typically created by \code{pdlm}).
#' 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{pdlm}, as returned by \code{\link[dlm]{dlm}}.
#' @param plot_data Logical; if \code{TRUE} (default) the observed data points are plotted.
#' @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.pdlm <- function(x,
                      plot_data=TRUE,
                      conf.int=FALSE,
                      sig.level=0.95,
                      ...) {
  if (!inherits(x, "pdlm"))
    stop("The object 'x' must be of class 'pdlm'. Ensure it was created using pdlm().")

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

  obs_col <- all.vars(x$formula)[1]
  obs <- x$data[, obs_col, drop = FALSE]
  pred <- x$ypred[, paste0(obs_col, '_pred'), drop=FALSE]

  time_index <- if (!is.null(x$date) && x$date %in% names(x$data)) {
    x$data[[x$date]]
  } else {
    seq_len(nrow(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 (x$log10){
    pred.temp <- log(pred+1, base=10)
  }
  var.pred <- x$var.pred

  upper_pred <- pred.temp + qnorm(p=sig.level, sd=sqrt(var.pred))
  lower_pred <- pred.temp - qnorm(p=sig.level, sd=sqrt(var.pred))

  if (x$log10){
    upper_pred <- 10**upper_pred - 1
    lower_pred <- 10**lower_pred - 1
  }
  #print(obs)
  #print(pred)
  #print(upper_pred)
  #print(lower_pred)
  plot_range <- range(obs, pred, upper_pred, lower_pred, na.rm = TRUE)

  default_args <- list(
    xlab = '',
    ylab = "Value",
    ylim = plot_range
  )

  user_args <- list(...)
  plot_args <- modifyList(default_args, user_args)
  if (plot_data) {
    do.call(matplot, c(list(x = time_index, y = obs, type = "p", pch = 16,
                            cex = 0.75, col = 'black'), plot_args))
  } else {
    do.call(plot, c(list(x = time_index, y = obs[, 1], type = "n"), plot_args))
  }

  do.call(matlines, c(list(x = time_index, y = pred, lwd = 2,
                             cex = 0.75, col = 'red', lty='solid'), plot_args))
  if (conf.int){
    polygon(c(time_index, rev(time_index)),
            c(lower_pred, rev(upper_pred)),
            col = adjustcolor('red', alpha.f = 0.2),
            border = NA)
  }


  invisible(NULL)
}