meteoR: Tools for Easier Analysis of Meteorological Fields

Documented in geom_arrow geom_vector

#' Arrows
#'
#' Parametrization of [ggplot2::geom_segment] either by location and displacement
#' or by magnitude and angle with default arrows. `geom_arrow()` is the same as
#' `geom_vector()` but defaults to preserving the direction under coordinate
#' transformation and different plot ratios.
#'
#' @param min.mag minimum magnitude for plotting vectors
#' @param skip,skip.x,skip.y numeric specifying number of gridpoints not to draw
#'  in the x and y direction
#' @param arrow.length,arrow.angle,arrow.ends,arrow.type parameters passed to
#' [grid::arrow]
#' @inheritParams ggplot2::geom_segment
#' @param direction direction of rotation (counter-clockwise or clockwise)
#' @param start starting angle for rotation in degrees
#' @param pivot numeric indicating where to pivot the arrow where 0 means at the
#' beginning and 1 means at the end.
#' @param preserve.dir logical indicating whether to preserve direction or not
#'
#' @details
#' Direction and start allows to work with different standards. For the
#' meteorological standard, for example, use `star = -90` and `direction = "cw"`.
#'
#' @section Aesthetics:
#' `geom_vector` understands the following aesthetics (required aesthetics are in bold)
#'
#' \itemize{
#' \item **x**
#' \item **y**
#' \item either **mag** and **angle**, or **dx** and **dy**
#' \item \code{alpha}
#' \item \code{colour}
#' \item \code{linetype}
#' \item \code{size}
#' \item \code{lineend}
#' }
#'
#' @examples
#' library(data.table)
#' library(ggplot2)
#'
#' data(seals)
#' # If the velocity components are in the same units as the axis,
#' # geom_vector() (or geom_arrow(preserve.dir = TRUE)) might be a better option
#' ggplot(seals, aes(long, lat)) +
#'     geom_arrow(aes(dx = delta_long, dy = delta_lat), skip = 1, color = "red") +
#'     geom_vector(aes(dx = delta_long, dy = delta_lat), skip = 1) +
#'     scale_mag()
#'
#' data(geopotential)
#' geopotential <- copy(geopotential)[date == date[1]]
#' geopotential[, gh.z := Anomaly(gh), by = .(lat)]
#' geopotential[, c("u", "v") := GeostrophicWind(gh.z, lon, lat)]
#'
#' (g <- ggplot(geopotential, aes(lon, lat)) +
#'     geom_arrow(aes(dx = dlon(u, lat), dy = dlat(v)), skip.x = 3, skip.y = 2,
#'                color = "red") +
#'     geom_vector(aes(dx = dlon(u, lat), dy = dlat(v)), skip.x = 3, skip.y = 2) +
#'     scale_mag( guide = "none"))
#'
#' # A dramatic illustration of the difference between arrow and vector
#' g + coord_polar()
#'
#' # When plotting winds in a lat-lon grid, a good way to have both
#' # the correct direction and an interpretable magnitude is to define
#' # the angle by the longitud and latitude displacement and the magnitude
#' # by the wind velocity. That way arrows are always parallel to streamlines
#' # and their magnitude are in the correct units.
#' ggplot(geopotential, aes(lon, lat)) +
#'     geom_contour(aes(z = gh.z)) +
#'     geom_vector(aes(angle = atan2(dlat(v), dlon(u, lat))*180/pi,
#'                    mag = Mag(v, u)), skip = 1, pivot = 0.5) +
#'     scale_mag()
#'
#' # Sverdrup transport
#' library(data.table)
#' b <- 10
#' d <- 10
#' grid <- as.data.table(expand.grid(x = seq(1, d, by = 0.5),
#'                                   y = seq(1, b, by = 0.5)))
#' grid[, My := -sin(pi*y/b)*pi/b]
#' grid[, Mx := -pi^2/b^2*cos(pi*y/b)*(d - x)]
#'
#' ggplot(grid, aes(x, y)) +
#'     geom_arrow(aes(dx = Mx, dy = My))
#'
#' # Due to limitations in ggplot2 (see: https://github.com/tidyverse/ggplot2/issues/4291),
#' # if you define the vector with the dx and dy aesthetics, you need
#' # to explicitly add scale_mag() in order to show the arrow legend.
#'
#' ggplot(grid, aes(x, y)) +
#'     geom_arrow(aes(dx = Mx, dy = My)) +
#'     scale_mag()
#'
#' # Alternative, use Mag and Angle.
#' ggplot(grid, aes(x, y)) +
#'     geom_arrow(aes(mag = Mag(Mx, My), angle = Angle(Mx, My)))
#'
#' @export
#' @family ggplot2 helpers
geom_arrow <- function(mapping = NULL, data = NULL,
                       stat = "arrow",
                       position = "identity",
                       ...,
                       start = 0,
                       direction = c("ccw", "cw"),
                       pivot = 0.5,
                       preserve.dir = TRUE,
                       min.mag = 0,
                       skip = 0,
                       skip.x = skip,
                       skip.y = skip,
                       arrow.angle = 15,
                       arrow.length = 0.5,
                       arrow.ends = "last",
                       arrow.type = "closed",
                       arrow = grid::arrow(arrow.angle, grid::unit(arrow.length, "lines"),
                                           ends = arrow.ends, type = arrow.type),
                       lineend = "butt",
                       na.rm = FALSE,
                       show.legend = NA,
                       inherit.aes = TRUE) {
  ggplot2::layer(geom = GeomArrow,
                 mapping = mapping,
                 data = data,
                 stat = stat,
                 position = position,
                 show.legend = show.legend,
                 inherit.aes = inherit.aes,
                 key_glyph = draw_key_vector,
                 params = list(
                   start = start,
                   direction = direction,
                   pivot = pivot,
                   preserve.dir = preserve.dir,
                   arrow = arrow,
                   lineend = lineend,
                   na.rm = na.rm,
                   skip.x = skip.x,
                   skip.y = skip.y,
                   min.mag = min.mag,
                   ...)
  )
}

#' @export
#' @rdname geom_arrow
geom_vector <- function(mapping = NULL, data = NULL,
                        stat = "arrow",
                        position = "identity",
                        ...,
                        start = 0,
                        direction = c("ccw", "cw"),
                        pivot = 0.5,
                        preserve.dir = FALSE,
                        min.mag = 0,
                        skip = 0,
                        skip.x = skip,
                        skip.y = skip,
                        arrow.angle = 15,
                        arrow.length = 0.5,
                        arrow.ends = "last",
                        arrow.type = "closed",
                        arrow = grid::arrow(arrow.angle, grid::unit(arrow.length, "lines"),
                                            ends = arrow.ends, type = arrow.type),
                        lineend = "butt",
                        na.rm = FALSE,
                        show.legend = NA,
                        inherit.aes = TRUE) {
  ggplot2::layer(geom = GeomArrow,
                 mapping = mapping,
                 data = data,
                 stat = stat,
                 position = position,
                 show.legend = show.legend,
                 inherit.aes = inherit.aes,
                 key_glyph = draw_key_vector,
                 params = list(
                   start = start,
                   direction = direction,
                   pivot = pivot,
                   preserve.dir = preserve.dir,
                   arrow = arrow,
                   lineend = lineend,
                   na.rm = na.rm,
                   skip.x = skip.x,
                   skip.y = skip.y,
                   min.mag = min.mag,
                   ...)
  )
}


draw_key_vector <- function(data, params, size) {
    if (is.null(data$linetype)) {
        data$linetype <- 0
    } else {
        data$linetype[is.na(data$linetype)] <- 0
    }

    `%||%` <- function(a, b) if (is.null(a)) b else a
    params$arrow$length <- data$mag*params$arrow$length

    grob <- grid::segmentsGrob(
        x0 = 0.1, y0 = 0.5, x1 = grid::unit(0.1, "npc") + grid::unit(data$mag, "cm"), y1 = 0.5,
        gp = grid::gpar(
            col = scales::alpha(data$colour %||% data$fill %||% "black", data$alpha),
            fill = scales::alpha(params$arrow.fill %||% data$colour %||% data$fill %||% "black", data$alpha),
            lwd = (data$linewidth %||% 0.5) * .pt,
            lty = data$linetype %||% 1,
            lineend = params$lineend %||% "butt"
        ),
        arrow = params$arrow
    )

    # Magick number is 1.25 because we like to span 80% of the width with the segment, so the
    # total width is 1 / 0.8 * size == 1.25 * size
    attr(grob, "width") <- 1.25 * data$mag # assumes cm
    grob

}


#' @rdname geom_arrow
#' @usage NULL
#' @format NULL
#' @export
GeomArrow <- ggplot2::ggproto("GeomArrow", ggplot2::Geom,
  required_aes = c("x", "y"),
  default_aes = ggplot2::aes(color = "black", size = 0.5, min.mag = 0,
                             linetype = 1, alpha = NA,
                             angle = 0, mag = 0),
  draw_key = draw_key_vector,
  draw_panel = function(data, panel_scales, coord,
                        arrow = arrow, lineend = lineend,
                        start = start, direction = direction,
                        preserve.dir = FALSE, pivot = 0.5) {
      if (!is.finite(pivot)) {
          stopf("'pivot' must be a number between 0 and 1.", call. = FALSE)
      }
      if (pivot > 1) {
          pivot <- 1
          warningf("'pivot' greater than 1, setting it to 1.", call. = FALSE)
      }
      if (pivot < 0) {
          pivot <- 0
          warningf("'pivot' less than 0, setting it to 0.", call. = FALSE)
      }
      mag <- data$norm_mag

      if ("simpleUnit" %in% class(unit(1, "mm"))) {
          arrow$length <- mag*arrow$length
      } else {
          arrow$length <- unit(as.numeric(arrow$length)*mag,
                               attr(arrow$length, "unit"))
      }

      if (preserve.dir == FALSE) {
          # For non linear coords
          data$group <- seq(nrow(data))
          data$piece <- 1
          data2 <- data
          data2$piece <- 2

          # Approximation for non linear coords.
          data2$x <- with(data, x + dx/10000)
          data2$y <-  with(data, y + dy/10000)

          coords <- coord$transform(data, panel_scales)
          coords2 <- coord$transform(data2, panel_scales)

          coords$xend <- coords2$x
          coords$yend <- coords2$y
          coords$dx <- with(coords, xend - x)/100
          coords$dy <- with(coords, yend - y)/100

          pol <- vectorGrob(x = coords$x, y = coords$y,
                            dx = coords$dx, dy = coords$dy,
                            length = unit(coords$mag, "cm"),
                            pivot = pivot,
                            preserve.dir = preserve.dir,
                            default.units = "npc",
                            arrow = arrow,
                            gp = grid::gpar(col = coords$colour,
                                            fill = scales::alpha(coords$colour, coords$alpha),
                                            alpha = ifelse(is.na(coords$alpha), 1, coords$alpha),
                                            lwd = coords$size*.pt,
                                            lty = coords$linetype,
                                            lineend = lineend))

      } else {
          coords <- coord$transform(data, panel_scales)
          pol <- arrowGrob(x = coords$x, y = coords$y,
                           angle = coords$angle,
                           length = unit(coords$mag, "cm"),
                           pivot = pivot,
                           preserve.dir = preserve.dir,
                           default.units = "native",
                           arrow = arrow,
                           gp = grid::gpar(col = coords$colour,
                                           fill = scales::alpha(coords$colour, coords$alpha),
                                           alpha = ifelse(is.na(coords$alpha), 1, coords$alpha),
                                           lwd = coords$size*.pt,
                                           lty = coords$linetype,
                                           lineend = lineend))
      }
      pol
  })

#' @rdname geom_arrow
#' @usage NULL
#' @format NULL
#' @export
StatArrow <- ggplot2::ggproto("StatArrow", ggplot2::Stat,
  required_aes = c("x", "y"),
  default_aes = ggplot2::aes(min.mag = 0, dx = NULL, dy = NULL,
                             mag = NULL, angle = NULL),
  compute_group = function(self, data, scales,
                           skip.x = skip.x, skip.y = skip.y,
                           min.mag = min.mag, start = 0, direction = -1,
                           preserve.dir = TRUE, ...) {
      data
  },
  setup_data = function(data, params) {


      params$direction <- switch(params$direction[1],
                                 ccw = -1,
                                 cw = 1,
                                 stopf("'direction' must be either \"ccw\" or \"cw\"", call. = FALSE)
      )
      if (is.null(data$mag) | is.null(data$angle)) {
          if (is.null(data$dx) | is.null(data$dy)) {
              stopf("stat_arrow needs wither 'mag' and 'angle' or 'dx' and 'dy'.", call. = FALSE)
          }
          data$mag <- with(data, Mag(dx, dy))
          data$angle <- with(data, atan2(dy, dx)*180/pi)
      } else {
          # Turn into mathematical angle
          data$angle <-  params$start - data$angle*params$direction
          data$dx <- with(data, mag*cos(angle*pi/180))
          data$dy <- with(data, mag*sin(angle*pi/180))
      }

      data <- subset(data, x %in% JumpBy(sort(unique(x)), params$skip.x + 1) &
                         y %in% JumpBy(sort(unique(y)), params$skip.y + 1) &
                         mag >= params$min.mag)

      data$norm_mag <- with(data, mag/max(mag, na.rm = TRUE))
      data
  },
  compute_panel = function(self, data, scales,
                           skip.x = 0, skip.y = 0,
                           min.mag = 0, start = 0, direction = -1,
                           preserve.dir = TRUE, ...) {
      if (plyr::empty(data)) return(data.frame())

      groups <- split(data, data$group)
      stats <- lapply(groups, function(group) {
          self$compute_group(data = group, scales = scales, ...)
      })

      stats <- mapply(function(new, old) {
          if (plyr::empty(new)) return(data.frame())
          unique <- ggplot2:::uniquecols(old)
          missing <- !(names(unique) %in% names(new))
          cbind(
              new,
              unique[rep(1, nrow(new)), missing,drop = FALSE]
          )
      }, stats, groups, SIMPLIFY = FALSE)

      data <- do.call(plyr::rbind.fill, stats)

      min.mag <- data$min.mag %||% min.mag

      # Warnings for good usage
      if (preserve.dir == FALSE) {
          if (scales$x$is_discrete() | scales$y$is_discrete()) {
              warningf("The use of preserve.dir = FALSE with discrete scales is not recommended.", call. = FALSE)
          }

          trans_name <- scales$x$trans$name
          if (trans_name == "date") {
              warningf("The use of preserve.dir = FALSE with date scales is not recommended.", call. = FALSE)
          }
      }

      data
  }
)