R/plot_walk.R

In mbedward/walkies: Random walks

#' Plots walk paths using ggplot.
#'
#' Plots one or more \code{walk.data} object paths (generated by
#' \code{\link{make_walk}}).
#'
#' @param w Either a single \code{walk.data} object, or a list of
#'   such objects.
#' @param colour The variable to use for colouring within each walk path:
#'   \code{'direct'} for direct distance to start point;
#'   \code{'path'} for distance along path;
#'   \code{'none'} for no colouring
#' @param legend If \code{TRUE}, a legend is added to the plot when path
#'   colouring is used. Ignored if \code{colour} is \code{'none'}.
#' @param by.id If \code{TRUE}, plots are facetted by the \code{walk.data$id}
#'   field. If \code{FALSE} (default), a single plot is produced.
#' @param nrow If \code{by.id} is \code{TRUE}, this argument can be used to set the
#'   number of rows of facetted plots.
#' @param ncol If \code{by.id} is \code{TRUE}, this argument can be used to set the
#'   number of columns of facetted plots. See \code{link{ggplot2::facet_wrap}}.
#'
#' @return A plot object.
#'
#' @importFrom ggplot2 ggplot aes coord_equal facet_wrap
#'   geom_path scale_colour_continuous theme theme_bw
#'
#' @export
#'
plot_walk <- function(w, colour=c("none", "direct", "path"), legend=FALSE, by.id=FALSE, nrow=NULL, ncol=NULL) {

  colour <- match.arg(colour)

  walk2df <- function(wdat, index) {
    dat <- data.frame(x=wdat$x, y=wdat$y, step=0:wdat$nsteps, id=wdat$id, index=index)

    if (colour == "direct")
      dat$distance <- sqrt((wdat$x - wdat$xstart)^2 + (wdat$y - wdat$ystart)^2)
    else if (colour == "path")
      dat$distance <- cumsum(wdat$steplen)

    dat
  }

  if (inherits(w, "walk.data"))
    df <- walk2df(w, 1)
  else if (is.list(w)) {
    df <- lapply(1:length(w), function(i) {
      walk2df(w[[i]], i)
    })

    df <- do.call(rbind, df)
  }
  else
    stop("Argument w should be a walk.data object or a list of walk.data objects")

  g <- ggplot() +
    coord_equal() +
    theme_bw()

  if (colour == "direct" || colour == "path") {
    g <- g + geom_path(data=df, aes(x, y, colour=distance, group=index))

    if (colour == "direct")
      g <- g + scale_colour_continuous(name="Direct distance", low="darkred", high="blue")
    else
      g <- g + scale_colour_continuous(name="Path distance", low="darkred", high="blue")

    if (!legend)
      g <- g + theme(legend.position="none")
  }
  else { # no distance colouring
    g <- g + geom_path(data=df, aes(x, y, group=index), colour="darkgrey")
  }

  if (by.id)
    g <- g + facet_wrap(~ id, nrow = nrow, ncol = ncol)

  g
}