R/plot_methods.R

In fitlandr: Fit Vector Fields and Potential Landscapes from Intensive Longitudinal Data

#' Plot a 2D vector field
#'
#' Plot a 2D vector field estimated by [fit_2d_vf()]. Powered by [ggplot2::ggplot()].
#'
#' @param x A `vectorfield` object estimated by [fit_2d_vf()].
#' @param arrow The description of the arrow heads of the vectors on the plot (representing the vector field). Generated by [grid::arrow()]. Also see the `arrow` parameter of [ggplot2::geom_segment()].
#' @param show_estimated_vector Show the vectors from the estimated model? `TRUE` by default.
#' @param estimated_vector_enlarge A number. How many times should the vectors (representing the estimated vector field) be enlarged on the plot? This can be useful when the estimated vector field is too strong or too weak.
#' @param estimated_vector_options A list passing other customized parameters to [ggplot2::geom_segment()] to control the vectors representing the estimated vector field.
#' @param show_point Show the original data points? `TRUE` by default.
#' @param point_options A list passing other customized parameters to [ggplot2::geom_point()] to control the points representing the original data point.
#' @param show_original_vector Show the original vectors (i.e., the vectors between data points)? `FALSE` by default.
#' @param original_vector_enlarge A number. How many times should the original vectors be enlarged on the plot?
#' @param original_vector_options A list passing other customized parameters to [ggplot2::geom_segment()] to control the vectors representing the original data.
#' @param show_used_vector Only for vector fields estimated by the "VFC" method. Should the vectors from the original data that are considered inliers be specially marked? `FALSE` by default.
#' @param used_vector_options Only for vector fields estimated by the "VFC" method. A list passing other customized parameters to [ggplot2::geom_segment()] to control the vectors representing the inliers. Red by default.
#' @param show_v_norm Show the norm of the estimated vectors (the strength of the vector field)? `FALSE` by default.
#' @param v_norm_options A list passing other customized parameters to [ggplot2::geom_raster()] to control the layer representing the norm of the estimated vectors.
#' @param ... Not in use.
#'
#' @return A `ggplot2` plot.
#'
#' @export
plot.vectorfield <- function(x, arrow = grid::arrow(length = grid::unit(0.1, "cm")),
                             show_estimated_vector = TRUE,
                             estimated_vector_enlarge = 1,
                             estimated_vector_options = list(),
                             show_point = TRUE,
                             point_options = list(size = 0.5),
                             show_original_vector = FALSE,
                             original_vector_enlarge = 1,
                             original_vector_options = list(),
                             show_used_vector = FALSE,
                             used_vector_options = list(color = "red"),
                             show_v_norm = FALSE,
                             v_norm_options = list(),
                             ...) {
  p <-
    ggplot2::ggplot(x$vec_grid, ggplot2::aes(x = x, y = y)) +
    ggplot2::theme_bw() +
    ggplot2::labs(x = x$x, y = x$y)

  if (show_estimated_vector) {
    p <- p +
      do.call(
        ggplot2::geom_segment,
        c(
          list(
            mapping = ggplot2::aes(
              xend = x + vx * estimated_vector_enlarge,
              yend = y + vy * estimated_vector_enlarge
            ),
            arrow = arrow
          ),
          estimated_vector_options
        )
      )
  }

  if (show_point) {
    p <- p +
      do.call(
        ggplot2::geom_point,
        c(
          list(
            data = x$data %>% as.data.frame(),
            mapping = ggplot2::aes(x = .data[[x$x]], y = .data[[x$y]])
          ),
          point_options
        )
      )
  }

  if (show_original_vector) {
    p <- p +
      do.call(
        ggplot2::geom_segment,
        c(
          list(
            data = x$original_vectors %>% as.data.frame(),
            mapping = ggplot2::aes(
              x = x,
              y = y,
              xend = x + vx * original_vector_enlarge,
              yend = y + vy * original_vector_enlarge
            ),
            arrow = arrow
          ),
          original_vector_options
        )
      )
  }

  if (show_used_vector) {
    p <- p +
      do.call(
        ggplot2::geom_segment,
        c(
          list(
            data = x$original_vectors[x$VFCresult$VFCIndex, ] %>% as.data.frame(),
            mapping = ggplot2::aes(
              x = x,
              y = y,
              xend = x + vx * original_vector_enlarge,
              yend = y + vy * original_vector_enlarge
            ),
            arrow = arrow
          ),
          used_vector_options
        )
      )
  }

  if (show_v_norm) {
    p <- p -
      do.call(
        ggplot2::geom_raster,
        c(
          list(
            mapping = ggplot2::aes(
              fill = v_norm
            )
          ),
          v_norm_options
        )
      ) + ggplot2::scale_fill_viridis_c()
  }

  return(p)
}


#' @references krassowski's answer at https://stackoverflow.com/questions/20249653/insert-layer-underneath-existing-layers-in-ggplot2-object
`-.gg` <- function(plot, layer) {
  if (missing(layer)) {
    stop("Cannot use `-.gg()` with a single argument. Did you accidentally put - on a new line?")
  }
  if (!ggplot2::is.ggplot(plot)) {
    stop("Need a plot on the left side")
  }
  plot$layers <- c(layer, plot$layers)
  plot
}