R/plt.contourFacet.R

In lessR: Less Code with More Comprehensive Results

.plt.contour.facet <-
function(x, y, facet1, facet2, facet1.name, facet2.name,
         contour_n, contour_nbins, contour_points, size.pt,
         theme, scale_x, scale_y, pad_x, pad_y, x.lab, y.lab,
         ellipse, ellipse_color, ellipse_lwd, fit.line, fit_color, fit_lwd,
         axis_x_pre, axis_y_pre, axis_fmt, n_row, n_col, quiet,
         render = "contour", smooth_power = 0.25, smooth_points = 100,
         smooth_bins = 128, main = NULL) {

  smooth <- render == "smooth"

  # 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)

  # -------------------------------------------------------------------------
  # colors: contour uses the single-plot filled-band ramp; smooth uses the
  #   single-plot density ramp (window_fill -> theme hue) painted as a raster
  if (smooth) {
    hue <- .get.h(theme)
    satur <- ifelse(theme %in% c("gray", "white"), 0, 80)
    col.s <- colorRampPalette(c(getOption("window_fill"),
                                hcl(hue, satur, 16)))(256)
  } else {
    if (theme %in% c("white", "gray", "light")) {
      clrs <- colorRampPalette(c("white",
        gray.colors(contour_n, rev = TRUE, start = 0.20, end = 0.92)))
    } else {
      sc <- .corcolors()
      clrs <- colorRampPalette(c("white", getColors(sc$fill_hi)))
    }
    col.n <- clrs(contour_n)
  }
  n.grid <- if (smooth) smooth_bins else contour_nbins

  # -------------------------------------------------------------------------
  # build the panel grid: one contour per facet cell
  # build a panel label:  VarName = "level"  (quote non-numeric levels only,
  #   so an integer level reads as  Year = 2020,  not  Year = "2020")
  .lab.lv <- function(nm, lv) {
    q <- ifelse(is.na(suppressWarnings(as.numeric(lv))),
                paste0('"', lv, '"'), lv)
    paste0(nm, " = ", q)
  }

  if (!is.factor(facet1)) facet1 <- factor(facet1)
  lv1 <- levels(droplevels(facet1))

  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)
    # panels filled row-major: for each facet2 level (row), each facet1 (col)
    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)
  }

  # -------------------------------------------------------------------------
  # shared scale across all panels
  # common KDE grid + bandwidth so density (z) is comparable panel-to-panel
  good.all <- complete.cases(x, y)
  x.all <- x[good.all]; y.all <- y[good.all]

  x_ext <- diff(range(x.all)) * 0.12
  y_ext <- diff(range(y.all)) * 0.12
  kde.lims <- c(range(x.all) + c(-x_ext, x_ext),
                range(y.all) + c(-y_ext, y_ext))
  # smooth matches the single-panel smoothScatter() so the two are consistent:
  #   its kernel SD is (5-95% quantile spread)/25, and kde2d's h is 4x the
  #   kernel SD.  contour keeps the wider normal-reference bandwidth, the
  #   established look for density contours.
  if (smooth) {
    ss.bw <- function(v) diff(stats::quantile(v, c(0.05, 0.95),
                              na.rm = TRUE, names = FALSE)) / 25
    h.common <- 4 * c(ss.bw(x.all), ss.bw(y.all))
  } else {
    h.common <- c(MASS::bandwidth.nrd(x.all), MASS::bandwidth.nrd(y.all))
  }
  h.common[h.common <= 0] <- NA  # let kde2d pick if degenerate

  # per-panel KDE on the common grid; collect densities and global z range
  dens.list <- vector("list", length(sel))
  z.rng <- NULL  # grows to the global density range across drawn panels
  for (i in seq_along(sel)) {
    ok <- intersect(sel[[i]], which(good.all))
    if (length(ok) < 3) next  # too few points for a density
    xi <- x[ok]; yi <- y[ok]
    h.i <- h.common
    h.i[is.na(h.i)] <- c(MASS::bandwidth.nrd(xi),
                         MASS::bandwidth.nrd(yi))[is.na(h.i)]
    if (any(h.i <= 0)) next  # still degenerate (e.g., all points collinear)
    d <- MASS::kde2d(xi, yi, n = n.grid, lims = kde.lims, h = h.i)
    dens.list[[i]] <- d
    z.rng <- range(c(z.rng, d$z))
  }
  if (is.null(z.rng)) {
    cat("\n"); stop(call. = FALSE, "\n------\n",
      "No facet cell has enough data (>= 3 points) for a ",
      if (smooth) "smooth density." else "contour.", "\n\n")
  }
  # contour: shared band breaks; smooth: shared image() zlim (transformed,
  #   since the raster maps the smooth_power-transformed density to color)
  lv.z <- seq(z.rng[1], z.rng[2], length.out = contour_n + 1)
  zlim.s <- pmax(z.rng, 0) ^ smooth_power

  # shared plot limits across panels
  if (is.null(scale_x)) {
    x.lim <- range(x.all)
    x.lim <- x.lim + c(-diff(x.lim) * pad_x[1], diff(x.lim) * pad_x[2])
    axT1 <- pretty(x.lim, eps.correct = 0)
  } else {
    x.lim <- scale_x[1:2]; axT1 <- axTicks(1, axp = scale_x)
  }
  if (is.null(scale_y)) {
    y.lim <- range(y.all)
    y.lim <- y.lim + c(-diff(y.lim) * pad_y[1], diff(y.lim) * pad_y[2])
    ayT2 <- pretty(y.lim, eps.correct = 0)
  } else {
    y.lim <- scale_y[1:2]; ayT2 <- axTicks(2, axp = scale_y)
  }
  lbl.x <- .axis.format(axT1, axis_fmt, axis_x_pre, axis_y_pre = "no")
  lbl.y <- .axis.format(ayT2, axis_fmt, axis_x_pre = "no", axis_y_pre)

  # -------------------------------------------------------------------------
  # draw the grid
  ax <- .axes_dim()
  par(bg = getOption("window_fill"))
  main.ok <- !is.null(main) && nzchar(main[1])
  # outer margin for shared axis titles; top expands for a title if given
  par(oma = c(3.4, 4.6, if (main.ok) 2.4 else 0.7, 0.8))
  par(mfrow = c(nr, nc))
  par(mgp = c(3.5, 0.4, 0), tcl = -0.25)

  n.pan <- nr * nc
  for (i in seq_len(n.pan)) {
    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()
    plot.window(xlim = x.lim, ylim = y.lim)

    if (i <= length(sel) && !is.null(dens.list[[i]])) {
      d <- dens.list[[i]]
      if (smooth)
        image(d$x, d$y, d$z ^ smooth_power, col = col.s, zlim = zlim.s,
              add = TRUE, useRaster = TRUE)
      else
        .filled.contour(d$x, d$y, d$z, levels = lv.z, col = col.n)

      ok <- intersect(sel[[i]], which(good.all))
      xi <- x[ok]; yi <- y[ok]
      if (smooth) {
        # overplot the smooth_points least-dense points (smoothScatter's
        #   nrpoints), so sparse tails show individual observations
        if (smooth_points > 0) {
          ix <- pmin(pmax(findInterval(xi, d$x), 1L), length(d$x))
          iy <- pmin(pmax(findInterval(yi, d$y), 1L), length(d$y))
          take <- head(order(d$z[cbind(ix, iy)]), smooth_points)
          points(xi[take], yi[take], pch = 21, cex = size.pt,
                 bg = "#00000055", col = "white", lwd = 0.2)
        }
      }
      else if (contour_points)
        points(xi, yi, pch = 21, cex = size.pt,
               bg = "#00000044", col = "white", lwd = 0.2)
      if (ellipse[1] > 0 && length(ok) >= 3) {
        mu <- c(mean(xi), mean(yi)); Sig <- cov(cbind(xi, yi))
        for (j in seq_along(ellipse)) {
          e <- ellipse::ellipse(Sig, centre = mu, level = ellipse[j],
                                npoints = 100)
          lines(e, col = ellipse_color, lwd = ellipse_lwd)
        }
      }
      if (fit.line != "off" && length(ok) >= 3) {
        pf <- .plt.fit(xi, yi, fit.line = fit.line, fit_power = 0,
                       fit_new = NULL, sz.lv = length(xi))
        lines(pf$x.lv, pf$f.ln, col = fit_color, lwd = fit_lwd)
      }
    }

    if (last.row)
      axis(1, at = axT1, labels = lbl.x, cex.axis = ax$axis_x_size,
           col.axis = ax$axis_x_text_color)
    if (first.col)
      axis(2, at = ayT2, labels = lbl.y, las = 1,
           cex.axis = ax$axis_y_size,
           col.axis = ax$axis_y_text_color)
    if (i <= length(labs))
      title(main = labs[i], cex.main = getOption("lab_size") * 0.95,
            col.main = getOption("lab_color"), font.main = 1, line = 0.3)
    box(col = getOption("panel_color"))
  }

  # shared axis titles in the outer margin
  lab.cex <- getOption("lab_size") * 1.15
  mtext(x.lab, side = 1, outer = TRUE, line = 1.6, cex = lab.cex,
        col = getOption("lab_color"))
  mtext(y.lab, side = 2, outer = TRUE, line = 2.8, cex = lab.cex,
        col = getOption("lab_color"))
  if (main.ok)
    mtext(main, side = 3, outer = TRUE, line = 0.4, font = 2,
          cex = getOption("main_size", 1.1),
          col = getOption("main_color"))

  # on.exit() restores par
}