Make facet formulas compatible with density plots

DESCRIPTION

@@ -1,7 +1,7 @@
 Package: plot2
 Type: Package
 Title: Lightweight extension of base R plot
-Version: 0.0.3.9018
+Version: 0.0.3.9019
 Authors@R: 
   c(
     person(

NAMESPACE

@@ -29,11 +29,15 @@ importFrom(graphics,plot.new)
 importFrom(graphics,plot.window)
 importFrom(graphics,points)
 importFrom(graphics,polygon)
+importFrom(graphics,rect)
 importFrom(graphics,segments)
+importFrom(graphics,text)
 importFrom(graphics,title)
 importFrom(methods,as)
 importFrom(stats,as.formula)
 importFrom(stats,model.frame)
+importFrom(stats,na.omit)
 importFrom(stats,update)
+importFrom(utils,head)
 importFrom(utils,modifyList)
 importFrom(utils,tail)

NEWS.md

@@ -1,6 +1,6 @@
 # News
 
-##  0.0.3.918 (development version)
+##  0.0.3.919 (development version)
 
 Website:
 
@@ -26,8 +26,8 @@ existing plot window. (#60 @grantmcdermott)
 - `plot2` gains a new `facet` argument for drawing faceted plots. Users can
 override the default square arrangement by passing the desired number of facet
 rows or columns to the companion `facet.args` helper function. Facets can be
-combined with `by` grouping, or used on their own. (#83, #91, #94, #96, #101
-@grantmcdermott)
+combined with `by` grouping, or used on their own.
+(#83, #91, #94, #96, #101, #103 @grantmcdermott)
 - Users can now control `plot2`-specific graphical parameters globally via
 the new `par2()` function (which is modeled on the base `par()` function). At
 the moment only a subset of global parameters, mostly related to legend and

R/par2.R

@@ -40,6 +40,8 @@
 #' @export
 par2 = function(...) {
 
+  facet.col = facet.bg = facet.border = NULL
+
   opts = list(...)
   par2_old = as.list(.par2)
   nam = names(opts)

R/plot2.R

@@ -198,8 +198,9 @@
 #'   the "Details" section of \code{\link[graphics]{plot}}.
 #'   
 #' @importFrom grDevices adjustcolor extendrange palette palette.colors palette.pals hcl.colors hcl.pals xy.coords
-#' @importFrom graphics abline arrows axis Axis box grconvertX grconvertY lines par plot.default plot.new plot.window points polygon segments title mtext
-#' @importFrom utils modifyList tail
+#' @importFrom graphics abline arrows axis Axis box grconvertX grconvertY lines par plot.default plot.new plot.window points polygon segments title mtext text rect
+#' @importFrom utils modifyList head tail
+#' @importFrom stats na.omit
 #' 
 #' @examples
 #' 
@@ -1141,6 +1142,8 @@ plot2.default = function(
 plot2.formula = function(
     x = NULL,
     data = parent.frame(),
+    facet = NULL,
+    facet.args = NULL,
     type = "p",
     xlim = NULL,
     ylim = NULL,
@@ -1199,8 +1202,60 @@ plot2.formula = function(
   m$formula = formula
   ## need stats:: for non-standard evaluation
   m[[1L]] = quote(stats::model.frame)
+  # catch for facets (need to ensure that na.omit, na.action, etc. are at done
+  # at the same level to avoid mismatches if there any missing variables)
+  has_facet_fml = !is.null(facet) && inherits(facet, "formula")
+  if (has_facet_fml) {
+    facet_fml = facet
+    facet_fml[[1]] = as.name("+")
+    if (no_y) {
+      m$formula = eval(substitute(
+        update(formula,  ~  . + facet_fml), list(facet_fml = facet_fml)
+      ))
+    } else {
+      m$formula = eval(substitute(
+        update(formula, .  ~ . + facet_fml), list(facet_fml = facet_fml)
+      ))
+    }
+  }
   mf = eval.parent(m)
 
+  ## We need to do some extra work if we included facet variables in the model
+  ## frame above
+  if (has_facet_fml) {
+    # separate the facet columns from the rest of the model frame
+    facet_n_cols = length(all.vars(facet_fml))
+    fmf = mf[, tail(seq_along(mf), facet_n_cols), drop = FALSE]
+    mf = mf[, head(seq_along(mf), -facet_n_cols), drop = FALSE]
+    # now do some prep work on the facets themselves for nicer plotting (e.g,
+    # grid arrangement for two-sided facet formula)
+    no_yfacet = length(facet) == 2L
+    facet_fml_rhs = if (no_yfacet) 2L else 3L
+    if (no_yfacet) {
+      yfacet_loc = NULL
+      xfacet_loc = 1L
+    } else {
+      yfacet_loc = 1L
+      xfacet_loc = 2L
+    }
+    if (NCOL(fmf) < xfacet_loc) stop("formula should specify at least one variable on the right-hand side")
+    yfacet = if (no_yfacet) NULL else fmf[, yfacet_loc]
+    xfacet = fmf[, xfacet_loc:NCOL(fmf)]
+
+    ## return object
+    xfacet = interaction(xfacet, sep = ":")
+    if (no_yfacet) {
+      facet = xfacet
+    } else {
+      # yfacet = interaction(yfacet, sep = ":")
+      ## NOTE: We "swap" the formula LHS and RHS since mfrow plots rowwise
+      facet = interaction(xfacet, yfacet, sep = "~")
+      attr(facet, "facet_grid") = TRUE
+      attr(facet, "facet_nrow") = length(unique(yfacet))
+    }
+
+  }
+
   ## extract variables: x, y (if any), by (if any)
   if (no_y) {
     y_loc = NULL
@@ -1237,6 +1292,7 @@ plot2.formula = function(
 
   plot2.default(
     x = x, y = y, by = by,
+    facet = facet, facet.args = facet.args,
     data = data,
     type = type,
     xlim = xlim,
@@ -1269,6 +1325,7 @@ plot2.density = function(
     x = NULL,
     by = NULL,
     facet = NULL,
+    facet.args = NULL,
     type = c("l", "area"),
     xlim = NULL,
     ylim = NULL,
@@ -1294,6 +1351,12 @@ plot2.density = function(
   type = match.arg(type)
   ## override if bg = "by"
   if (!is.null(bg) || !is.null(fill)) type = "area"
+
+  # catch for facet_grid
+  if (!is.null(facet)) {
+    facet_attributes = attributes(facet)
+    # facet_grid = attr(facet, "facet_grid")
+  }
 
   if (inherits(x, "density")) {
     object = x
@@ -1303,6 +1366,12 @@ plot2.density = function(
   } else {
     ## An internal catch for non-density objects that were forcibly
     ## passed to plot2.density (e.g., via a one-side formula)
+    if (anyNA(x)) {
+      x = na.omit(x)
+      if (!is.null(by)) by = by[-attr(x, "na.action")]
+      if (!is.null(facet)) facet = facet[-attr(x, "na.action")]
+      x = as.numeric(x)
+    }
     object = stats::density(x)
     legend.args = list(...)[["legend.args"]]
   }
@@ -1313,6 +1382,12 @@ plot2.density = function(
     y = object$y
   } else {
     x = eval(str2lang(object$data.name), envir = parent.frame())
+    if (anyNA(x)) {
+      x = na.omit(x)
+      if (!is.null(by) && length(by) != length(x)) by = by[-attr(x, "na.action")]
+      if (!is.null(facet) && length(facet) != length(x)) facet = facet[-attr(x, "na.action")]
+      x = as.numeric(x)
+    }
     if (is.null(facet) || identical(by, facet)) {
       split_x = split(x, f = by) 
     } else if (is.null(by)) {
@@ -1341,8 +1416,20 @@ plot2.density = function(
     }
     by_names = tryCatch(as(by_names, class(by)), error = function(e) if (inherits(by, "factor")) as.factor(by_names) else by_names)
     # need to coerce facet variables to factors for faceting to work properly later on
-    facet_names = tryCatch(as(facet_names, class(facet)), error = function(e) facet_names)
-    facet_names = tryCatch(as.factor(facet_names), error = function(e) facet_names)
+    # if we originally passed a factor, try to preserve this order for grid arrangement
+    if (inherits(facet, "factor")) {
+      orig_len = length(levels(facet))
+      new_len = length(facet_names)
+      if (orig_len == new_len) {
+        facet_names = levels(facet)
+      } else {
+        ## need to recycle names if nested in multiple by splits
+        facet_names = rep(levels(facet), each = new_len/orig_len)
+      }
+    } else {
+      facet_names = tryCatch(as(facet_names, class(facet)), error = function(e) facet_names)
+      facet_names = tryCatch(as.factor(facet_names), error = function(e) facet_names)
+    }
 
     split_object = lapply(seq_along(split_object), function(ii) {
       lst = list(
@@ -1397,8 +1484,18 @@ plot2.density = function(
   if (is.null(xlab)) xlab = paste0("N = ", object$n, "   Bandwidth = ", sprintf("%.4g", object$bw))
   if (is.null(main)) main = paste0(paste(object$call, collapse = "(x = "), ")")
 
+  # if (!is.null(facet)) attr(facet, "facet_grid") = facet_grid
+  if (!is.null(facet)) {
+    if (!is.null(facet_attributes[["levels"]])) {
+      facet = factor(facet, levels = facet_attributes[["levels"]])
+    } else {
+      facet = factor(facet)
+    }
+    attr(facet, "facet_grid") = facet_attributes[["facet_grid"]]
+  }
+
   plot2.default(
-    x = x, y = y, by = by, facet = facet,
+    x = x, y = y, by = by, facet = facet, facet.args = facet.args,
     type = type,
     xlim = xlim,
     ylim = ylim,