R/traj_plot.R

In dtkaplan/mosaicUSAFA: Functions, data, and tutorials for the Math 141Z course

#' Plots to correspond with integrateODE()
#'
#' Utilities to work with solutions to ODEs as generated by integrateODE
#'
#' @param one a formula describing the plot, e.g. B(t) ~ R(t) or B(t) ~ t or a ggplot
#' object
#' @param two placeholder for handling when `one` is a ggplot object
#' @param soln the solution as produced by `integrateODE()`
#' @param npts number of plotted points (default: 500)
#' @param nt number of tick marks to use in a trajectory plot
#' @param domain Optional list like `domain=domain(t=c(0,100))`. By default, this
#' will be inferred from `soln`
#'
#' @export
traj_plot <- function(one, two, soln, npts=500, nt=5, domain=NULL,  ...) {
  if (inherits(one, "formula")) {form <- one; soln <- two; one <- NULL}
  else if (inherits(one, "ggplot")) {form <- two}
  else stop("First argument should be either a formula or a piped-in ggplot")
  # process the formula
  # should have length 3
  as_traj <- TRUE
  if (length(form) != 3) stop("formula should be a two-sided formula")
  if (is.name(form[[3]])) {
    # just a simple variable name
    right_var_name <- right_in_name <- as.character(form[[3]])
    as_traj <- FALSE
  } else {
    if (length(form[[3]]) != 2) stop("use var name as function on right side of formula, e.g. x(t)")
    right_in_name <- as.character(form[[3]][[2]])
    right_var_name <- as.character(form[[3]][[1]])
  }
  if (length(form[[2]]) != 2)
    stop("use var name as function on left side of formula, e.g. y(t)")

  left_in_name <- as.character(form[[2]][[2]])
  left_var_name <- as.character(form[[2]][[1]])
  if (left_in_name != right_in_name) stop("Use same parametric name on both sides of formula")

  DF <- soln_to_df(soln, domain=domain, npts=npts, xname=left_in_name)
  tick_times <- pretty(range(DF[[left_in_name]]), n = nt)
  Ticks <- soln_to_df(soln, invals=tick_times)
  names(Ticks)[1] <- "label"

  plot_formula <- as.formula(paste(left_var_name, "~", right_var_name))
  P <- gf_path(one, plot_formula, data = DF, ...)
  if (as_traj && nt > 0) { # add tick marks
    # the_aes <- do.call(aes, list(x=as.name(right_var_name),
    #                              y=as.name(left_var_name),
    #                              label = as.name("label")))
    P <- P %>%
      gf_label(plot_formula, data = Ticks, label=~label, vjust=1, hjust=1,...) %>%
      gf_point(plot_formula, data=Ticks, ...)

  }

  P
}
#'
#' @export
soln_to_df <- function(soln, domain=NULL, npts=300, invals=NULL, xname="t") {
  vars <- names(soln)
  for (v in vars) {
    if( !is.function(soln[[v]]))
      stop("`soln` must be a list of interpolating functions")
  }

  first <- soln[[vars[[1]]]]
  if (is.null(invals)) { # figure out input values, if not explicitly given as arg.
    input_range <- if (is.list(domain)) domain[[1]]
                   else range(environment(first)$x)
    invals <- seq(input_range[1], input_range[2], length = npts)
  }
  Res <- tibble::tibble(x = invals)
  names(Res) <- xname
  for (v in vars) Res[[v]] <- soln[[v]](invals)

  return(Res)
}