R/pretty_smooths.R

In edwardlavender/prettyGraphics: Pretty Graphics with R

######################################
######################################
#### pretty_smooth_1d()

#' @title Plot pretty one dimensional smooths
#' @description This function is designed to plot prettier versions of one dimensional smooths created by \code{\link[mgcv]{plot.gam}}. For each smooth, the necessary information for plotting is taken from a user-supplied list generated by \code{\link[mgcv]{plot.gam}}. Plots are created using \code{\link[prettyGraphics]{pretty_plot}}. Confidence intervals can be added via an internal call to \code{\link[prettyGraphics]{add_error_envelope}}. Partial residuals and a rug for the observed values can also be added and customised.
#' @param fit A list returned by \code{\link[mgcv]{plot.gam}}. Each element is a list which contains all the necessary information to recreate a plot for a specific model term.
#' @param shift A number which defines a value by which to shift model predictions/partial residuals vertically.
#' @param select An integer that specifies the smooth term(s) (i.e., the elements in \code{fit}) to be plotted.
#' @param add_fit (optional) A named list of arguments to customise the appearance of the fitted line. This is passed to \code{\link[graphics]{lines}} or the \code{add_fit} argument of \code{\link[prettyGraphics]{add_error_envelope}}.
#' @param add_se_type (optional) A character that defines the method by which CIs are added to the plot: as lines (\code{"lines"}) or as a shaded polygon (\code{"poly"}). This is passed to the \code{type} argument of \code{\link[prettyGraphics]{add_error_envelope}}.
#' @param add_se (optional) A named list of arguments to customise the appearance of the confidence intervals. This is passed to the \code{add_ci} argument of \code{\link[prettyGraphics]{add_error_envelope}}.
#' @param add_resid (optional) A named list of arguments to customise the appearance of partial residuals via \code{\link[graphics]{points}}. These are taken from the 'p.resid' elements in \code{fit}.
#' @param add_rug (optional) A named list of arguments to add a rug of observed values to the plot. Observed values are taken from the 'raw' elements in \code{fit}. This is passed to \code{\link[graphics]{rug}}.
#' @param xlim,ylim Axis limits for all plots. If \code{NULL}, axis limits are chosen according to the inputs to \code{pretty_axis_args}.
#' @param pretty_axis_args A named list of arguments, passed to \code{\link[prettyGraphics]{pretty_axis}} to customise axes.
#' @param add_xlab,add_ylab (optional) Named list of arguments to customise the x and y axis labels. Labels are taken from the 'xlab' and 'ylab' elements in \code{fit} respectively. Lists are passed to \code{\link[graphics]{mtext}}.
#' @param assign_main (optional) A logical input that defines whether or not to assign a title to the plot. If \code{TRUE}, each plot (i)'s title is given by \code{LETTERS[1:length(select)][i]}.
#' @param add_main (optional) A named list of arguments to customise plot titles. Labels are assigned (see \code{assign_main}) or taken from the 'main' elements in \code{fit}. This is passed to \code{\link[graphics]{mtext}}.
#' @param one_page A logical input that defines whether or not to plot all smooths on one page.
#' @param return_list (depreciated) A logical input which defines whether or not to return a list, with one element for each \code{select} value, each containing the named list of axis arguments from \code{\link[prettyGraphics]{pretty_axis}}.
#' @param ... Additional arguments (none implemented).
#'
#' @details For all \code{add_*} arguments, \code{add_* = NULL} suppresses the argument, \code{add_* = list()} implements the argument with default values and a named list customises the output.
#'
#' @return The function returns a pretty plot of one dimensional smooths(s) and, invisibly, the list of pretty axis parameters produced by \code{\link[prettyGraphics]{pretty_axis}}.
#'
#' @examples
#' #### Simulate some example data and fit an example model
#' n <- 100
#' x <- runif(n, 0, 100)
#' z <- runif(n, 0, 100)
#' y <- 0.1 * x^2 + 0.002 * z^3 + stats::rnorm(n, 0, 100)
#' mod <- mgcv::gam(y ~ s(x) + s(z))
#' fit <- mgcv::plot.gam(mod, residuals = TRUE, pages = 1)
#'
#' #### pretty_smooth_1d() implementation
#' # The default options
#' pretty_smooth_1d(fit)
#' # The number of smooth terms is controlled via select
#' # ... and can be plotted on one page via one_page
#' pretty_smooth_1d(fit, select = 1:2, one_page = TRUE)
#' # The fitted line can be controlled via add_fit
#' pretty_smooth_1d(fit, add_fit = list(col = "red"))
#' # Confidence intervals can be suppressed via add_se_type = NULL or
#' # ... controlled via add_se and
#' pretty_smooth_1d(fit, add_fit = list(col = "red"), add_se = NULL)
#' pretty_smooth_1d(fit,
#'                  add_fit = list(col = "red"),
#'                  add_se_type = "lines",
#'                  add_se = list(col = "blue"))
#' # Partial residuals can be added via add_resid
#' pretty_smooth_1d(fit, add_resid = list(cex = 0.25))
#' # A rug can be added/surpressed via add_rug
#' pretty_smooth_1d(fit, add_rug = NULL)
#' pretty_smooth_1d(fit, add_rug = list(ticksize = 0.01))
#' # Axis titles can be controlled via add_xlab, add_ylab and add_main
#' pretty_smooth_1d(fit, add_main = list(cex = 2))
#' # Axis titles are taken from the fitted object, so can be changed
#' # ... by changing the appropriate element in 'fit'
#' fit[[1]]$xlab <- "Updated x name"
#' pretty_smooth_1d(fit, add_main = list(cex = 2))
#'
#' @author Edward Lavender
#' @export
#'

pretty_smooth_1d <- function(fit,
                             shift = 0,
                             select = 1,
                             add_fit = list(),
                             add_se_type = "poly",
                             add_se = list(col = scales::alpha("lightgrey", 0.8), border = FALSE),
                             add_resid = NULL,
                             add_rug = NULL,
                             xlim = NULL, ylim = NULL,
                             pretty_axis_args = list(),
                             add_xlab = list(line = 2),
                             add_ylab = list(line = 2),
                             assign_main = TRUE, add_main = list(adj = 0),
                             one_page = TRUE,
                             return_list = NULL,
                             ...){

  #### Set up
  if(one_page) pp <- graphics::par(mfrow = par_mf(length(select)))

  #### Create a plot for each variable
  if(requireNamespace("pbapply", quietly = TRUE)){
    loop <- function(X, FUN,...) pbapply::pblapply(X, FUN, ...)
  } else{
    loop <- function(X, FUN,...) lapply(X, FUN, ...)
  }
  axis_ls_by_term <- loop(1:length(select), function(i){

    #### Extract element of interest
    sel <- select[i]
    ifit <- fit[[sel]]

    #### Checks
    if(!is.null(add_resid)){
      if(is.null(ifit$p.resid)){
        message(paste0("add_resid = TRUE ignored: fit[[", sel, "]]$p.resid not provided."))
        add_resid <- NULL
      }
    }

    #### Define x and y limits
    if(shift != 0){
      ifit$fit <- ifit$fit + shift
      if(!is.null(add_resid)) ifit$p.resid <- ifit$p.resid + shift
    }
    x_rng <- range(ifit$x)
    cis <- list(fit = ifit$fit, lowerCI = ifit$fit - ifit$se, upperCI = ifit$fit + ifit$se)
    y_rng <- range(ifit$fit)
    if(!is.null(add_se_type)) y_rng <- range(c(y_rng, cis$lowerCI, cis$upperCI))
    if(!is.null(add_resid)) y_rng <- range(c(y_rng, ifit$p.resid))

    #### Make blank plot
    pretty_axis_args$x <- list(x =  x_rng, y = y_rng)
    if(is.null(pretty_axis_args$side)) pretty_axis_args$side <- 1:2
    axis_ls <- pretty_plot(ifit$x, ifit$fit,
                           xlim = xlim, ylim = ylim,
                           pretty_axis_args = pretty_axis_args,
                           xlab = "", ylab = "",
                           type = "n"
                           )

    #### Add CIs and fitted line
    if(is.null(add_se_type) & is.null(add_se)){
      if(!is.null(add_fit)) {
        add_fit$x <- ifit$x
        add_fit$y <- ifit$fit
        do.call(graphics::lines, add_fit)
      }
    } else{
      add_error_envelope_args <- list()
      add_error_envelope_args$x <- ifit$x
      add_error_envelope_args$ci <- cis
      add_error_envelope_args$type <- add_se_type
      add_error_envelope_args$add_ci <- add_se
      if(!is.null(add_fit)) {
        add_error_envelope_args$add_fit <- add_fit
        if(length(add_fit) == 0) add_fit$col <- "black"
        add_error_envelope_args$fitted_gp <- add_fit
      } else {
        add_error_envelope_args$add_fit <- NULL
      }
      do.call(add_error_envelope, add_error_envelope_args)
    }

    #### Add residuals
    if(!is.null(add_resid)){
      add_resid$x <- ifit$raw
      add_resid$y <- ifit$p.resid
      do.call(graphics::points, add_resid)
    }

    #### Add rug
    if(!is.null(add_rug)){
      if(is.null(add_rug$x)) add_rug$x <- ifit$raw
      if(is.null(add_rug$side)) add_rug$side <- 1
      if(is.null(add_rug$pos))  add_rug$pos  <- min(axis_ls[[2]]$lim)
      do.call(graphics::rug, add_rug)
    }

    #### Define/add labels
    ## x axis
    if(!is.null(add_xlab)) {
      add_xlab$side <- 1
      add_xlab$text <- ifit$xlab
      do.call(graphics::mtext, add_xlab)
    }
    ## y axis
    if(!is.null(add_ylab)) {
      add_ylab$side <- 2
      add_ylab$text <- ifit$ylab
      do.call(graphics::mtext, add_ylab)
    }
    ## main
    if(!is.null(add_main)) {
      if(assign_main) ifit$main <- LETTERS[i]
      if(!is.null(ifit$main)){
        if(is.null(add_main$side)) add_main$side <- 3
        add_main$text <- ifit$main
        do.call(graphics::mtext, add_main)
      }
    }

  })

  #### Reset graphics
  if(one_page) graphics::par(pp)

  #### Return outputs
  if(!is.null(return_list)) warning("The 'return_list' argument is depreciated.")
  return(invisible(axis_ls_by_term))

}


######################################
######################################
#### pretty_smooth_2d()

#' @title Pretty two-dimensional smooths
#' @description This is function plots pretty two-dimensional smooths as a contour plot.
#' @param x An mgcv model (see \code{\link[mgcv]{vis.gam}}).
#' @param view A character vector of two variables (see \code{\link[mgcv]{vis.gam}}).
#' @param xlim,ylim,pretty_axis_args Axis control arguments. \code{xlim} and \code{ylim} control x and y axis limits via \code{\link[prettyGraphics]{pretty_axis}}. If not supplied, \code{xlim} and \code{ylim} are defined as the range across each corresponding variable. To use default 'pretty' limits instead, specify \code{xlim = NULL} and \code{ylim = NULL}. Additional axis customisation is implemented by passing a named list of arguments to \code{\link[prettyGraphics]{pretty_axis}} via \code{pretty_axis_args}.
#' @param add_xy A named list of arguments, passed to \code{\link[graphics]{points}}, to add observations to the plot. \code{add_xy = NULL} suppresses this option, \code{add_xy = list()} implements default arguments and a named list customises these.
#' @param add_rug_x,add_rug_y Named list of arguments, passed to \code{\link[graphics]{rug}}, to add observed values of the variables defined in \code{view} to the plot. \code{add_rug_* = NULL} suppresses this option, \code{add_rug_*} implements default arguments and a named list customises these.
#' @param return_list (depreciated) A logical input which defines whether or not to return the list produced by \code{\link[prettyGraphics]{pretty_axis}}.
#' @param ... Additional arguments passed to \code{\link[mgcv]{vis.gam}}, excluding \code{"plot.type"} (since only contour plots are supported by this function) and \code{"axes"}.
#'
#' @details At present, the function is simply a wrapper for \code{\link[mgcv]{vis.gam}} with the additional flexibility provided by the \code{\link[prettyGraphics]{pretty_axis}} function and the \code{add_xy}, \code{add_rug_x} and \code{add_rug_y} arguments.
#'
#' @return The function returns a contour plot of the predictions of a generalised addition model for the two variables defined in \code{view} and, invisibly, the list of pretty axis parameters produced by \code{\link[prettyGraphics]{pretty_axis}}.
#'
#' @examples
#' #### Simulate example data and fit model (following ?mgcv::vis.gam examples)
#' set.seed(0)
#' n    <- 200
#' sig2 <- 4
#' x0   <- runif(n, 0, 1)
#' x1   <- runif(n, 0, 1)
#' x2   <- runif(n, 0, 1)
#' y    <- x0^2 + x1 * x2 + runif(n, -0.3, 0.3)
#' g    <- mgcv::gam(y ~ s(x0, x1, x2))
#'
#' #### Example (1): Contour plot using default options
#' pretty_smooth_2d(g, view = c("x1", "x2"))
#'
#' #### Example (2): Customise axes via xlim, ylim and pretty_axis_args
#' # Use xlim and ylim
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  xlim = c(0, 1),
#'                  ylim = c(0, 1))
#' # Use pretty_axis_args
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  pretty_axis_args = list(side = 1:4))
#'
#' #### Example (3): Add observed data
#' # Specify list() to use default options
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  add_xy = list())
#' # Customise addition of observed data
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  add_xy = list(pch = ".'", cex = 5))
#'
#' #### Example (4): Add rugs for the x and y variables
#' # Use default options
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  add_rug_x = list(),
#'                  add_rug_y = list())
#' # Customise options
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  add_rug_x = list(col = "grey"),
#'                  add_rug_y = list(col = "grey"))
#'
#' #### Example (5): Pass additional options to mgcv::vis.gam() via ...
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  main = "", xlab = "x1[...]", ylab = "x2[...]",
#'                  color = "gray",
#'                  contour.col = "red")
#'
#' #### Example (5): Integrate with add_colour_bar()
#' # Define plotting window with space for legend
#' pp <- graphics::par(oma = c(2, 2, 2, 10))
#' # Define z-limits for plot and legend and associated colours
#' zlim  <- c(0, 1.5)
#' n_col <- 100
#' col_param <- pretty_cols_brewer(zlim = zlim,
#'                                 pal = grDevices::heat.colors,
#'                                 n_breaks = n_col + 1)
#' # Make plot with colour scheme
#' pretty_smooth_2d(g, view = c("x1", "x2"),
#'                  zlim = col_param$zlim,
#'                  col = "heat",
#'                  nCol = n_col,
#'                  nlevels = n_col)
#' # Define legend param required for add_colour_bar()
#' legend_at <-
#'   data.frame(x = col_param$breaks[1:(length(col_param$breaks) - 1)],
#'              col = col_param$col)
#' legend_axis <- pretty_axis(side = 4, x = list(legend_at$x))
#' # Add legend as subplot
#' TeachingDemos::subplot(add_colour_bar(legend_at,
#'                                       pretty_axis_args = legend_axis),
#'                                       x = 1, y = 0.05,
#'                                       vadj = 0, hadj = 0,
#'                                       size = c(0.2, 2))
#'
#' @author Edward Lavender
#' @export
#'

pretty_smooth_2d <- function(x,
                             view = NULL,
                             xlim, ylim,
                             pretty_axis_args = list(side = 1:4,
                                                     axis = list(list(),
                                                                 list(),
                                                                 list(labels = FALSE),
                                                                 list(labels = FALSE))),
                             add_xy = NULL,
                             add_rug_x = NULL, add_rug_y = NULL,
                             return_list = NULL,
                             ...) {

  #### Define data and fix limits across range of data
  check...(c("plot.type", "axes"), ...)
  dat <- stats::model.frame(x)
  if(is.null(view)) view <- colnames(dat)[1:2]
  if(length(view) != 2L) stop("'view' should contain two variable (column) names.")
  if(missing(xlim)) xlim <- range(dat[, view[1]])
  if(missing(ylim)) ylim <- range(dat[, view[2]])

  #### Implement pretty_axis_args
  if(!is.null(pretty_axis_args$side)) {
    if(!all(c(1, 2) %in% pretty_axis_args$side))
      stop("pretty_axis_args$side is expected to contain sides 1 and 2 (at a minimum) for a two-dimensional plot.")
  }
  axis_ls <- implement_pretty_axis_args(x = list(dat[, view[1]], dat[, view[2]]),
                                        pretty_axis_args = pretty_axis_args,
                                        xlim = xlim,
                                        ylim = ylim,...)
  xlim <- axis_ls[[1]]$lim
  ylim <- axis_ls[[2]]$lim

  #### Use mgcv::vis.gam() to plot smooth (contour)
  # Make plot
  mgcv::vis.gam(x = x, view = view,
                plot.type = "contour",
                axes = FALSE, xlim = xlim, ylim = ylim,...)

  #### Add observed data
  if(!is.null(add_xy)) {
    add_xy$x <- dat[, view[1]]
    add_xy$y <- dat[, view[2]]
    do.call(graphics::points, add_xy)
  }

  #### Add rugs
  # x variable
  if(!is.null(add_rug_x)) {
    if(is.null(add_rug_x$x)) add_rug_x$x <- dat[, view[1]]
    if(is.null(add_rug_x$side)) add_rug_x$side <- 1
    if(is.null(add_rug_x$pos)) {
      if(add_rug_x$side == 1) {
        add_rug_x$pos <- ylim[1]
      } else if(add_rug_x$side == 3) {
        add_rug_x$pos <- ylim[2]
      } else {
        stop("add_rug_x$side should be 1 or 3.")
      }
    }
    do.call(graphics::rug, add_rug_x)
  }
  # y variable
  if(!is.null(add_rug_y)) {
    if(is.null(add_rug_y$x)) add_rug_y$x <- dat[, view[2]]
    if(is.null(add_rug_y$side)) add_rug_y$side <- 2
    if(is.null(add_rug_y$pos)) {
      if(add_rug_y$side == 2) {
        add_rug_y$pos <- xlim[1]
      } else if(add_rug_y$side == 3) {
        add_rug_y$pos <- xlim[2]
      } else {
        stop("add_rug_y$side should be 2 or 4.")
      }
    }
    do.call(graphics::rug, add_rug_y)
  }

  #### Add pretty axes at the end
  pretty_axis(axis_ls = axis_ls, add = TRUE)

  #### Return blank
  if(!is.null(return_list)) warning("The 'return_list' argument is depreciated.")
  return(invisible(axis_ls))
}