R/plot.bcrosspred.R

In bdlnm: Bayesian Distributed Lag Non-Linear Models (B-DLNM)

#' Plot the predicted effects from a Bayesian distributed-lag model (B-DLNM).
#'
#' It plots the lag-specific, overall, slices, contour or 3D representations for predicted effects generated by [bcrosspred()].
#'
#' @param x An object of class `"bcrosspred"` returned by [bcrosspred()].
#' @param ptype Character. Plot type: one of `"slices"`, `"overall"`, `"3d"` or `"contour"`.
#' @param exp_at Numeric vector of exposure values (predictor) for the `"slices"` plot type (maximum length of 4).
#' @param lag_at Numeric vector of lags for the `"slices"` plot type (maximum length of 4).
#' @param ci Character. How to plot credible intervals: one of `"area"`, `"bars"`, `"lines"`, `"sampling"` (draw a line for each posterior sample) or `"none"`. Default: `"area"`.
#' @param ci.arg List of graphical arguments for the plotting of credible intervals passed to base plotting functions (see `Details`).
#' @param ci.level Numeric in `(0,1)`. Credible interval level. Default is taken from `x$ci.level`.
#' @param cumul Logical. If `TRUE`, plot incremental cumulative predictions (requires cumulative predictions in `x`).
#' @param exponentiate Logical. Set to `TRUE` to plot exponentiated effects (relative risks) and to `FALSE` to not plot them. The default is `NULL` in which the function uses `x$model.link` to automatically detect whether to use relative risks (when link is `log` or `logit`).
#' @param ... Additional graphical arguments passed to base plotting functions (see `Details`).
#'
#' @details
#'
#' The function supports different visualizations for posterior effects predicted by [bcrosspred()] specified in `ptype`:
#'
#' - `"slices"`: (optionally multi-panel) line plot of the predicted effect over lag(s) for selected exposure value(s) or over exposure value(s) for selected lag(s). Use `exp_at` and/or `lag_at` to specify which slices to draw (at least one must be supplied). At most 4 slices of each type are allowed to keep layouts readable. See [plot.default()], [lines()] and [points()] for information on additional graphical arguments.
#' - `"overall"`: plot the predicted overall cumulative effect over the whole lag period for each exposure value in the prediction grid. See [plot.default()], [lines()] and [points()] for information on additional graphical arguments.
#' - `"3d"`: 3D plot of the surface of the predicted effect for each exposure and lag. Uses the median of the posterior samples. Not meaningful for single lag models. Additional graphical arguments can be included, such as `theta` or `phi` (perspective), `border` or `shade` (surface), `xlab`, `ylab` or `zlab` (axis labelling) and `col.` See [persp()] for additional information.
#' - `"contour"`: filled contour of the surface of the predicted effect for each exposure and lag. Uses the median of the posterior samples. Not meaningful for single lag models. Additional graphical arguments can be included, such as `plot.title`, `plot.axes` or `key.title` for titles and axis and key labelling. See [filled.contour()] for additional information.
#'
#' The function will call the specified original base plotting functions for each different `ptype`. Via the argument `...` the user can change the graphical parameters that will be passed to these functions. See the original functions for a complete list of the arguments. Some arguments, if not specified, are set to different default values than the original functions.
#'
#' Credible intervals will only be drawn for `ptype` equal to `"slices"` or `"overall"`. The type of credible interval will be given by the `ci` argument, with options `"area"` (default), `"bars"`, `"lines"`, `"sampling"` (draw a line for each posterior sample) or `"none"` (no credible intervals). Their appearance may be modified through `ci.arg`, a list of arguments passed to to low-level plotting functions: [polygon()] for `"area"`, [segments()] for `"bars"` and [lines()] for `"lines"`. See the original functions for a complete list of the arguments. As above, some unspecified arguments are set to different default values. The credible interval level will be given by `ci.level` or inferred from `x$ci.level` by default.
#'
#' If exponentiated effects (relative risks) are specified (`exponentiate = TRUE`) or auto-detected if `x$model.link` is `log` or `logit`, the function will use the predicted relative risks stored in `x$matRRfit`, `x$allRRfit` or `x$cumRRfit` (if `cumul=TRUE`).
#'
#' In the presence of unlagged or single lag associations, only the `overall` plot can be produced. In this case, the `overall` plot represents the effect at each predictor value, rather than the overall cumulative effect across lags.
#'
#' @return No return value, called for side effects.
#' 
#' @author Pau Satorra, Marcos Quijal-Zamorano.
#'
#' @note This function is inspired by [dlnm::plot.crosspred()] developed by Gasparrini (2011) <doi:10.18637/jss.v043.i08>. It has been adapted to work in a Bayesian framework within the \pkg{bdlnm} package.
#'
#' @references
#'
#' Gasparrini A. (2011). Distributed lag linear and non-linear models in R: the package dlnm. _Journal of Statistical Software_, 43(8), 1-20. <doi:10.18637/jss.v043.i08>.
#'
#' Quijal-Zamorano M., Martinez-Beneito M.A., Ballester J., Marí-Dell'Olmo M. (2024). Spatial Bayesian distributed lag non-linear models (SB-DLNM) for small-area exposure-lag-response epidemiological modelling. _International Journal of Epidemiology_, 53(3), dyae061. <doi:10.1093/ije/dyae061>.
#'
#' @seealso [bcrosspred()] to predict exposure–lag–response associations for a `"bdlnm"` object,
#' @seealso [bdlnm()] to fit a Bayesian distributed lag non-linear model (`"bdlnm"`).
#'
#' @export
#'
#' @examples
#' # Set exposure-response and lag-response spline parameters
#'  dlnm_var <- list(
#'    var_prc = c(10, 75, 90),
#'    var_fun = "ns",
#'    lag_fun = "ns",
#'    max_lag = 21,
#'    lagnk = 3
#'  )
#'
#'
#' # Set cross-basis parameters
#'  argvar <- list(fun = dlnm_var$var_fun,
#'                 knots = stats::quantile(london$tmean,
#'                                  dlnm_var$var_prc/100, na.rm = TRUE),
#'                 Bound = range(london$tmean, na.rm = TRUE))
#'
#'  arglag <- list(fun = dlnm_var$lag_fun,
#'                 knots = dlnm::logknots(dlnm_var$max_lag, nk = dlnm_var$lagnk))
#'
#'  # Create crossbasis
#'  cb <- dlnm::crossbasis(london$tmean, lag = dlnm_var$max_lag, argvar, arglag)
#'
#'  # Seasonality of mortality time series
#'  seas <- splines::ns(london$date, df = round(8 * length(london$date) / 365.25))
#'
#'  # Prediction values (equidistant points)
#'  temp <- round(seq(min(london$tmean), max(london$tmean), by = 0.1), 1)
#'  # Ensure it falls inside the range of temperatures after rounding:
#'  temp <- temp[temp >= min(london$tmean) & temp <= max(london$tmean)]
#'
#' if (check_inla()) {
#'  # Fit the model
#'  mod <- bdlnm(mort_75plus ~ cb + factor(dow) + seas, data = london, family = "poisson",
#'              sample.arg = list(n = 1000, seed = 432))
#'
#'  # Prediction
#'  cpred <- bcrosspred(mod, exp_at = temp)
#'
#'  # Perform the plots:
#'
#'  #Overall
#'  plot(cpred, "overall", xlab = "Temperature (ºC)", ylab = "Relative Risk", col = 4,
#'  main="Overall", log = "y")
#'
#'  #3-D plot
#'  plot(cpred, "3d", zlab = "Relative risk", col = 4, lphi = 60, cex.axis = 0.6,
#'  xlab = "Temperature (ºC)", main = "3D graph of temperature effect")
#'
#'  #Contour
#'  plot(cpred, "contour", xlab = "Temperature (ºC)", ylab = "Lag", main="Contour plot")
#'
#'  #Slices (for a high temperature)
#'  htemp <- 23
#'  plot(cpred , "slices", exp_at = htemp, col=3, ylab="RR",
#'  main=paste0("Association for a high temperature (", htemp, "ºC)"))
#'
#'  #Slices (for lag 0)
#'  plot(cpred , "slices", lag_at = 0, col=4, ylab="RR", main=paste0("Association at Lag 0"))
#'
#' }
#'
plot.bcrosspred <- function(
  x,
  ptype,
  exp_at = NULL,
  lag_at = NULL,
  ci = "area",
  ci.arg,
  ci.level = x$ci.level,
  cumul = FALSE,
  exponentiate = NULL,
  ...
) {
  ## ---------------------------
  ## Basic checks
  ## ---------------------------

  # Check x object
  if (missing(x) || !is.list(x) || !inherits(x, "bcrosspred")) {
    cli::cli_abort(
      "{.arg x} must be a {.cls bcrosspred} object as returned by {.fn bcrosspred}."
    )
  }

  # Check ptype
  if (missing(ptype) || !ptype %in% c("slices", "3d", "contour", "overall")) {
    cli::cli_abort(
      "{.arg ptype} must be one of {.val 'slices'}, {.val '3d'}, {.val 'contour' } or {.val 'overall'}."
    )
  }

  # Check CI
  if (!ci %in% c("area", "bars", "lines", "sampling", "none")) {
    cli::cli_abort(
      "{.arg ptype} must be one of {.val 'area'}, {.val 'bars'}, {.val 'lines' } or {.val 'sampling'} or {.val 'none'}."
    )
  }

  if (missing(ci.arg)) {
    ci.arg <- list()
  } else if (!is.list(ci.arg)) {
    cli::cli_abort("{.arg ci.arg} must be a list.")
  }

  if (!is.numeric(ci.level) || ci.level >= 1 || ci.level <= 0) {
    cli::cli_abort(
      "{.arg ci.level} must be numeric and between {.val 0} and {.val 1}."
    )
  }

  # Validate exp_at & lag_at inputs for slices
  if (ptype == "slices") {
    if (is.null(exp_at) && is.null(lag_at)) {
      cli::cli_abort(
        "At least one of {.arg exp_at} or {.arg lag_at} must be supplied when {.arg ptype} = {.val 'slices'}."
      )
    }

    if (!is.null(exp_at)) {
      if (!is.numeric(exp_at)) {
        cli::cli_abort("{.arg exp_at} must be numeric.")
      }
      if (length(exp_at) > 4L) {
        cli::cli_abort("{.arg exp_at} length must be <= 4.")
      }
      # Check that requested exp_at values are present in predictions
      if (!all(exp_at %in% x$exp_at)) {
        cli::cli_abort(
          "All {.arg exp_at} values must be present in {.code x$exp_at}."
        )
      }
    }

    if (!is.null(lag_at)) {
      if (!is.numeric(lag_at)) {
        cli::cli_abort("{.arg lag_at} must be numeric.")
      }
      if (length(lag_at) > 4L) {
        cli::cli_abort("{.arg lag_at} length must be <= 4.")
      }
      # check lag_at validity against predicted lags
      if (!all(lag_at %in% x$lag_at)) {
        cli::cli_abort(
          "All {.arg lag_at} values must be present in {.val x$lag_at}."
        )
      }
    }
  }

  # decide whether to compute CI from matfit directly or use stored summaries
  ci_compute <- !identical(ci.level, x$ci.level)

  # exp
  if (!is.null(exponentiate) && !is.logical(exponentiate)) {
    cli::cli_abort("{.arg exponentiate} must be logical")
  }

  ## ---------------------------
  ## Prepare data
  ## ---------------------------

  # exponentiate if requested or if detected model.link suggests it
  noeff <- 0L
  if (
    (is.null(exponentiate) &&
      !is.null(x$model.link) &&
      x$model.link %in% c("log", "logit")) ||
      (!is.null(exponentiate) && exponentiate)
  ) {
    matfit <- if (!cumul) x$matRRfit else x$cumRRfit
    allfit <- x$allRRfit
    matfitsum <- if (!cumul) x$matRRfit.summary else x$cumRRfit.summary
    allfitsum <- x$allRRfit.summary
    noeff <- 1L
  } else {
    matfit <- if (!cumul) x$matfit else x$cumfit
    allfit <- x$allfit
    matfitsum <- if (!cumul) x$matfit.summary else x$cumfit.summary
    allfitsum <- x$allfit.summary
    noeff <- 0L
  }

  ## ---------------------------
  ## Plot: slices
  ## ---------------------------
  if (ptype == "slices") {
    if (is.null(x$lag_at)) {
      cli::cli_abort("slices plot is not meaningful for unlagged associations.")
    }

    if (min(x$lag_at) == max(x$lag_at)) {
      cli::cli_abort(
        "slices plot is not meaningful for single lag associations."
      )
    }

    # graphical layout: one plot per requested exp_at/lag_at value
    npanels <- length(exp_at) + length(lag_at)

    # set grey scale
    grey <- grey(0.9)

    if (npanels > 1) {
      #restore user graphical parameters
      oldpar <- graphics::par(no.readonly = TRUE)
      on.exit(graphics::par(oldpar))
      
      graphics::layout(matrix(
        1:npanels,
        ncol = sum(!is.null(exp_at), !is.null(lag_at))
      ))
      grey <- grey(0.8)
      graphics::par(mgp = c(2, 0.7, 0), mar = c(4.1, 4.1, 2.1, 1.1))
    }

    # medians: use stored summary
    matfitmed <- matfitsum[,, "0.5quant"]

    # low / high intervals
    if (!ci_compute) {
      #use stored summary
      quantiles <- grep("quant$", dimnames(matfitsum)[[3]])
      matfitlow <- matfitsum[,, quantiles[1]]
      matfithigh <- matfitsum[,, rev(quantiles)[1]]
    } else {
      #recompute quantiles again
      quantiles <- c((1 - ci.level) / 2, 1 - (1 - ci.level) / 2)
      matfitlow <- apply(matfit, c(1, 2), stats::quantile, probs = quantiles[1])
      matfithigh <- apply(
        matfit,
        c(1, 2),
        stats::quantile,
        probs = quantiles[2]
      )
    }

    # plot by requested lags
    if (!is.null(lag_at)) {
      for (i in lag_at) {
        # filter for i-th lag_at
        matfit_i <- matfit[, which(x$lag_at == i), ]

        # set default plot arguments
        plot.arg <- list(
          type = "l",
          xlab = "Exposure",
          ylab = "Outcome",
          ylim = c(min(matfit_i, na.rm = TRUE), max(matfit_i, na.rm = TRUE)),
          bty = "l"
        )

        if (npanels > 1L) {
          plot.arg$cex.axis <- 0.7
        }

        # merge user args
        plot.arg <- utils::modifyList(plot.arg, list(...))

        # set credible intervals arguments
        ci.list <- list(
          panel.first = call(
            "fci",
            ci = ci,
            x = x$exp_at,
            y = matfit_i,
            high = matfithigh[, which(x$lag_at == i)],
            low = matfitlow[, which(x$lag_at == i)],
            ci.arg,
            plot.arg,
            noeff = noeff
          )
        )

        plot.arg <- utils::modifyList(
          plot.arg,
          c(
            ci.list,
            list(
              x = x$exp_at,
              y = matfitmed[, which(x$lag_at == i)]
            )
          )
        )

        col <- plot.arg$col

        if (npanels > 1) {
          plot.arg$main <- ""
          plot.arg$xlab <- "Exposure"
        }

        # plot
        do.call("plot", plot.arg)

        if (length(lag_at) > 1) graphics::mtext(paste("Lag =", i), cex = 0.8)
      }
    }

    # plot by requested predictor values (exp_at)
    if (!is.null(exp_at)) {
      for (i in exp_at) {
        # filter for i-th exposure value
        matfit_i <- matfit[which(x$exp_at == i), , ]

        # set default plot arguments
        plot.arg <- list(
          type = "l",
          xlab = "Lag",
          ylab = "Outcome",
          ylim = c(min(matfit_i, na.rm = TRUE), max(matfit_i, na.rm = TRUE)),
          bty = "l"
        )

        if (npanels > 1L) {
          plot.arg$cex.axis <- 0.7
        }

        # merge user args
        plot.arg <- utils::modifyList(plot.arg, list(...))

        # set credible intervals arguments
        ci.list <- list(
          panel.first = call(
            "fci",
            ci = ci,
            x = x$lag_at,
            y = matfit_i,
            high = matfithigh[which(x$exp_at == i), ],
            low = matfitlow[which(x$exp_at == i), ],
            ci.arg,
            plot.arg,
            noeff = noeff
          )
        )

        plot.arg <- utils::modifyList(
          plot.arg,
          c(ci.list, list(x = x$lag_at, y = matfitmed[which(x$exp_at == i), ]))
        )

        if (npanels > 1L) {
          plot.arg$main <- ""
          plot.arg$xlab <- "Lag"
        }

        # plot
        do.call("plot", plot.arg)

        if (length(lag_at) > 1) {
          graphics::mtext(paste("Exposure =", i), cex = 0.8)
        }
      }
    }
  }

  ## ---------------------------
  ## Plot: overall
  ## ---------------------------

  if (ptype == "overall") {
    # set default values
    min_all <- min(allfit, na.rm = TRUE)
    max_all <- max(allfit, na.rm = TRUE)

    plot.arg <- list(
      type = "l",
      ylim = c(min_all, max_all),
      xlab = "Exposure",
      ylab = "Outcome",
      bty = "l"
    )

    # merge with user args
    plot.arg <- utils::modifyList(plot.arg, list(...))

    col <- plot.arg$col

    # medians: use stored summary
    allfitmed <- allfitsum[, "0.5quant"]

    # low / high intervals
    if (!ci_compute) {
      #use stored summary
      quantiles <- grep("quant$", colnames(allfitsum))
      allfitlow <- allfitsum[, utils::head(quantiles, 1)]
      allfithigh <- allfitsum[, utils::tail(quantiles, 1)]
    } else {
      #recompute quantiles again
      quantiles <- c((1 - ci.level) / 2, 1 - (1 - ci.level) / 2)
      allfitlow <- apply(allfit, 1, stats::quantile, probs = quantiles[1])
      allfithigh <- apply(allfit, 1, stats::quantile, probs = quantiles[2])
    }

    # set credible intervals arguments
    ci.list <- list(
      panel.first = call(
        "fci",
        ci = ci,
        x = x$exp_at,
        y = allfit,
        high = allfithigh,
        low = allfitlow,
        ci.arg,
        plot.arg,
        noeff = noeff
      )
    )
    plot.arg <- utils::modifyList(
      plot.arg,
      c(ci.list, list(x = x$exp_at, y = allfitmed))
    )

    # plot
    do.call("plot", plot.arg)
  }

  ## ---------------------------
  ## Plot: 3D
  ## ---------------------------

  if (ptype == "3d") {
    if (is.null(x$lag_at)) {
      cli::cli_abort("slices plot is not meaningful for unlagged associations.")
    }

    if (min(x$lag_at) == max(x$lag_at)) {
      cli::cli_abort(
        "slices plot is not meaningful for single lag associations."
      )
    }

    # get the median summary
    matfitmed <- matfitsum[,, "0.5quant"]

    # set plot defaults
    plot.arg <- list(
      ticktype = "detailed",
      theta = 210,
      phi = 30,
      xlab = "Exposure",
      ylab = "Lag",
      zlab = "Outcome",
      col = "lightskyblue",
      zlim = c(min(matfitmed, na.rm = TRUE), max(matfitmed, na.rm = TRUE)),
      ltheta = 290,
      shade = 1,
      r = sqrt(3),
      d = 5,
      border = NA
    )

    # merge with the user args
    plot.arg <- utils::modifyList(plot.arg, list(...))

    plot.arg <- utils::modifyList(
      plot.arg,
      list(
        x = x$exp_at,
        y = x$lag_at,
        z = matfitmed
      )
    )

    # plot
    do.call("persp", plot.arg)
  }

  ## ---------------------------
  ## Plot: Contour
  ## ---------------------------

  if (ptype == "contour") {
    if (is.null(x$lag_at)) {
      cli::cli_abort("slices plot is not meaningful for unlagged associations.")
    }

    if (min(x$lag_at) == max(x$lag_at)) {
      cli::cli_abort(
        "slices plot is not meaningful for single lag associations."
      )
    }

    # get the median summary
    matfitmed <- matfitsum[,, "0.5quant"]

    # set default values
    levels <- pretty(matfitmed, n = 20)
    col1 <- grDevices::colorRampPalette(c("blue", "white"))
    col2 <- grDevices::colorRampPalette(c("white", "red"))
    col <- c(col1(sum(levels < noeff)), col2(sum(levels > noeff)))
    graphics::filled.contour(
      x = x$exp_at,
      y = x$lag_at,
      z = matfitmed,
      col = col,
      levels = levels,
      ...
    )
  }
}