R/plt.distFacet.R

In lessR: Less Code with More Comprehensive Results

.plt.dist.facet <-
function(x, facet1, facet2, facet1.name, facet2.name,
         render = "density",
         bandwidth = "nrd0", kind = "both", histo = TRUE,
         fill_hist = "transparent", col_nrm = "transparent",
         fill_nrm = "transparent", col_gen, fill_gen, lwd = NULL,
         rug = FALSE, color_rug = "gray40", size_rug = 0.5,
         full_curve = TRUE, theme = getOption("theme"),
         x.min = NULL, x.max = NULL,
         stat = "count", breaks = NULL, bin_start = NULL, bin_width = NULL,
         bin_end = NULL, fill_area = TRUE,
         axis_fmt = "K", axis_x_pre = "", axis_y_pre = "",
         x.lab = NULL, y_axis = FALSE, n_row = NULL, n_col = NULL, quiet,
         by = NULL, by_name = NULL, horiz = TRUE) {

  freqpoly <- render == "freq_poly"
  if (freqpoly) y_axis <- TRUE  # counts / proportions are the point of the y

  # a  by  variable overlays one distribution per group within every panel,
  #   the same overlay as the grouped distribution of a single panel. With
  #   no  by  there is a single unnamed group, so one code path draws both
  grouped <- !is.null(by)
  grp  <- if (grouped) droplevels(factor(by)) else factor(rep("", length(x)))
  lv.g <- levels(grp)
  n.g  <- length(lv.g)

  # one color per group, and the histogram behind a density curve is
  #   dropped when grouped: several of them overlaid obscure the curves
  clr.g <- if (grouped) .color_range(fill_gen, n.g) else fill_gen
  if (grouped) histo <- FALSE

  # save graphics state ONCE, restore on exit
  old.par <- par(no.readonly = TRUE)
  on.exit({
    ro <- c("cin", "cra", "csi", "cxy", "din", "page")
    par(old.par[setdiff(names(old.par), ro)])
  }, add = TRUE)

  # panel-strip label:  VarName = "level"  (quote non-numeric levels only),
  #   matching the faceted contour/smooth/hexbin forms
  .lab.lv <- function(nm, lv) {
    q <- ifelse(is.na(suppressWarnings(as.numeric(lv))),
                paste0('"', lv, '"'), lv)
    paste0(nm, " = ", q)
  }

  # -------------------------------------------------------------------------
  # build the panel grid: one distribution per facet cell, or with no
  #   facet a single panel with no title
  if (!is.null(facet1)) {
    if (!is.factor(facet1)) facet1 <- factor(facet1)
    lv1 <- levels(droplevels(facet1))
  }

  if (is.null(facet1)) {
    nc <- 1; nr <- 1
    sel <- list(seq_along(x))
    labs <- character(0)
  }
  else if (!is.null(facet2)) {  # two facets: rows = facet2, cols = facet1
    if (!is.factor(facet2)) facet2 <- factor(facet2)
    lv2 <- levels(droplevels(facet2))
    nc <- length(lv1); nr <- length(lv2)
    cells <- expand.grid(c1 = lv1, c2 = lv2, stringsAsFactors = FALSE)
    cells <- cells[order(match(cells$c2, lv2), match(cells$c1, lv1)), ]
    sel <- Map(function(a, b) which(facet1 == a & facet2 == b),
               cells$c1, cells$c2)
    labs <- paste0(.lab.lv(facet1.name, cells$c1), ", ",
                   .lab.lv(facet2.name, cells$c2))
  }
  else {  # one facet: near-square grid unless n_row / n_col supplied
    k <- length(lv1)
    if (!is.null(n_col))      { nc <- n_col; nr <- ceiling(k / nc) }
    else if (!is.null(n_row)) { nr <- n_row; nc <- ceiling(k / nr) }
    else { nc <- ceiling(sqrt(k)); nr <- ceiling(k / nc) }
    sel <- lapply(lv1, function(a) which(facet1 == a))
    labs <- .lab.lv(facet1.name, lv1)
  }

  # -------------------------------------------------------------------------
  # first pass: per-panel estimates, and the shared x-range / y-max so panels
  #   are directly comparable
  est <- vector("list", length(sel))
  y.max <- 0
  xr <- NULL

  # One set of bins for every panel. Binning each panel separately makes the
  #   bin width a function of that panel's sample size, so the bars of a
  #   small panel are both narrower and, on the shared density scale,
  #   taller than those of a large panel holding the same distribution
  .shared.breaks <- function() {
    xa <- x[is.finite(x)]
    if (!is.null(bin_width) || !is.null(bin_start)) {
      start.dflt <- is.null(bin_start)
      if (start.dflt) bin_start <- pretty(range(xa))[1]
      if (is.null(bin_width))
        bin_width <- diff(graphics::hist(xa, plot = FALSE)$breaks[1:2])
      # a pretty() start can sit more than a bin below the minimum,
      #   leaving empty leading bins, so move it up to the first data bin
      if (start.dflt)
        while (bin_start + bin_width < min(xa))
          bin_start <- bin_start + bin_width
      end <- if (!is.null(bin_end)) bin_end else max(xa)
      brk <- seq(bin_start, end, bin_width)
      while (max(brk) < max(xa)) brk <- seq(bin_start, max(brk) + bin_width,
                                            bin_width)
      brk
    }
    else if (!is.null(breaks)) breaks
    else graphics::hist(xa, plot = FALSE, breaks = "Sturges")$breaks
  }

  # est[[panel]][[group]] is the estimate drawn by one group in one panel,
  #   NULL where that combination holds too little data
  brk <- .shared.breaks()
  for (i in seq_along(sel)) {
    e.i <- vector("list", n.g)
    for (j in seq_len(n.g)) {
      xi <- x[sel[[i]]][grp[sel[[i]]] == lv.g[j]]
      xi <- xi[is.finite(xi)]

      if (freqpoly) {
        if (length(xi) < 1) next
        h <- graphics::hist(xi, breaks = brk, plot = FALSE)
        val <- if (stat == "proportion") h$counts / sum(h$counts)
               else h$counts
        e.i[[j]] <- list(mids = h$mids, val = val, bw = diff(brk)[1])
        y.max <- max(y.max, max(val))
      }
      else {  # density
        if (length(xi) < 2 || stats::sd(xi) == 0) next  # too few, no spread
        h     <- graphics::hist(xi, plot = FALSE, breaks = brk)
        d.gen <- suppressWarnings(stats::density(xi, bw = bandwidth))
        e.i[[j]] <- list(xi = xi, h = h, d.gen = d.gen,
                         mean = mean(xi), sd = stats::sd(xi))
        y.max <- max(y.max, max(d.gen$y), if (histo) max(h$density) else 0)
        xr <- range(c(xr, if (full_curve) range(d.gen$x) else range(xi)))
      }
    }
    est[[i]] <- e.i
  }
  if (freqpoly) xr <- range(brk)
  if (is.null(xr)) {
    cat("\n"); stop(call. = FALSE, "\n------\n",
      "No facet cell has enough data to plot.\n\n")
  }
  if (is.null(x.min)) x.min <- xr[1]
  if (is.null(x.max)) x.max <- xr[2]

  # normal curve also contributes to y-max (evaluated on the shared grid)
  if (!freqpoly && kind %in% c("normal", "both")) {
    xx <- seq(x.min, x.max, length = 200)
    for (i in seq_along(est)) for (e in est[[i]]) if (!is.null(e)) {
      dn <- stats::dnorm(xx, e$mean, e$sd)
      y.max <- max(y.max, max(dn))
    }
  }

  axT1 <- pretty(c(x.min, x.max))
  lbl.x <- .axis.format(axT1, axis_fmt, axis_x_pre, axis_y_pre = "no")
  ayT2 <- pretty(c(0, y.max))
  lbl.y <- .axis.format(ayT2, axis_fmt, axis_x_pre = "no", axis_y_pre)

  # -------------------------------------------------------------------------
  # draw the grid
  ax <- .axes_dim()
  col.bg  <- getOption("panel_fill")
  col.box <- getOption("panel_color")
  par(bg = getOption("window_fill"))
  # the value axis is at the bottom, or with a vertical value axis at the
  #   left, the other axis shown only with  y_axis
  oma.v <- 5.0;  oma.d <- if (y_axis) 3.4 else 1.0
  par(oma = c(if (horiz) 3.4 else oma.d,
              if (horiz) (if (y_axis) 4.6 else 2.4) else oma.v,
              if (grouped) 2.8 else 1.0, 0.8))  # top: room for the legend
  par(mfrow = c(nr, nc))
  par(mgp = c(3.5, 0.4, 0), tcl = -0.25)

  clr.g.tr <- .maketrans(clr.g, 100)

  # draw with f(), such as polygon() or lines(), the values  v  and their
  #   densities or counts  dn, each on its own axis
  .vd <- function(f, v, dn, ...) if (horiz) f(v, dn, ...) else f(dn, v, ...)
  # the border of a group is its own color, so that overlaid curves stay
  #   distinguishable; ungrouped keeps the single supplied line color
  brd.g <- if (grouped) clr.g else rep(col_gen, n.g)

  for (i in seq_len(nr * nc)) {
    row.i <- ((i - 1) %/% nc) + 1
    col.i <- ((i - 1) %%  nc) + 1
    last.row  <- row.i == nr
    first.col <- col.i == 1

    par(mar = c(0.6, 0.6, 1.5, 0.5))
    plot.new()
    if (horiz)
      plot.window(xlim = c(x.min, x.max), ylim = c(0, y.max))
    else
      plot.window(xlim = c(0, y.max), ylim = c(x.min, x.max))
    rect(par("usr")[1], par("usr")[3], par("usr")[2], par("usr")[4],
         col = col.bg, border = col.box)

    if (i <= length(est)) for (j in seq_len(n.g)) {
      e <- est[[i]][[j]]
      if (is.null(e)) next
      if (freqpoly) {
        # frequency polygon: bin values at midpoints, dropped to zero one bin
        #   beyond each end; optional fill under the line
        px <- c(e$mids[1] - e$bw, e$mids, e$mids[length(e$mids)] + e$bw)
        py <- c(0, e$val, 0)
        lwd.fp <- if (!is.null(lwd)) lwd else 1.35
        if (fill_area)
          .vd(polygon, px, py, col = clr.g.tr[j], border = brd.g[j],
              lwd = lwd.fp)
        else
          .vd(lines, px, py, col = brd.g[j], lwd = lwd.fp)
      }
      else {
        if (histo) {
          if (horiz)
            plot(e$h, add = TRUE, freq = FALSE, col = fill_hist,
                 border = "transparent")
          else
            rect(0, e$h$breaks[-length(e$h$breaks)], e$h$density,
                 e$h$breaks[-1], col = fill_hist, border = "transparent")
        }
        if (kind %in% c("normal", "both")) {
          xx <- seq(x.min, x.max, length = 200)
          dn <- stats::dnorm(xx, e$mean, e$sd)
          lwd.nrm <- if (!is.null(lwd)) lwd
                     else if (fill_nrm == "transparent") 1.35 else 1
          .vd(polygon, c(x.min, xx, x.max), c(0, dn, 0), col = fill_nrm,
              border = col_nrm, lwd = lwd.nrm)
        }
        if (kind %in% c("general", "both")) {
          lwd.gen <- if (!is.null(lwd)) lwd
                     else if (clr.g[j] == "transparent") 1.35 else 1
          .vd(polygon, e$d.gen$x, e$d.gen$y, col = clr.g.tr[j],
              border = brd.g[j], lwd = lwd.gen)
        }
        if (rug) rug(e$xi, col = color_rug, lwd = size_rug, ticksize = 0.05,
                     side = if (horiz) 1 else 2)
      }
    }

    # the value axis along the bottom row, or vertical along the first
    #   column, the axis of the densities or counts on the other side
    if (if (horiz) last.row else first.col)
      axis(if (horiz) 1 else 2, at = axT1, labels = lbl.x,
           cex.axis = if (horiz) ax$axis_x_size else ax$axis_y_size,
           col.axis = if (horiz) ax$axis_x_text_color
                      else ax$axis_y_text_color, las = 1)
    if ((if (horiz) first.col else last.row) && y_axis)
      axis(if (horiz) 2 else 1, at = ayT2, labels = lbl.y,
           cex.axis = if (horiz) ax$axis_y_size else ax$axis_x_size,
           col.axis = if (horiz) ax$axis_y_text_color
                      else ax$axis_x_text_color)
    if (i <= length(labs))
      title(main = labs[i], cex.main = getOption("lab_size") * 1.05,
            col.main = getOption("lab_color"), font.main = 1, line = 0.3)
    box(col = col.box)
  }

  # shared axis titles in the outer margin
  lab.cex <- getOption("lab_size") * 1.15
  mtext(x.lab, side = if (horiz) 1 else 2, outer = TRUE,
        line = if (horiz) 1.6 else 3.4, cex = lab.cex,
        col = getOption("lab_color"))
  if (y_axis) {
    y.lab <- if (!freqpoly) "Density"
             else if (stat == "proportion") "Proportion" else "Count"
    mtext(y.lab, side = if (horiz) 2 else 1, outer = TRUE,
          line = if (horiz) 2.8 else 1.6, cex = lab.cex,
          col = getOption("lab_color"))
  }

  # legend for the overlaid groups, across the top of the whole grid
  if (grouped) {
    par(fig = c(0, 1, 0, 1), oma = c(0, 0, 0, 0), mar = c(0, 0, 0, 0),
        new = TRUE)
    plot.new()
    legend("top", legend = lv.g, title = by_name, horiz = TRUE, bty = "n",
           fill = clr.g.tr, border = brd.g, cex = lab.cex * 0.92,
           text.col = getOption("lab_color"))
  }

  # on.exit() restores par
}