R/panel_fix.R

In thisplot: Utility Functions for Plotting

#' @title Set the panel width/height of a plot to a fixed value
#'
#' @description
#' The ggplot object, when stored, can only specify the height and width of the entire plot, not the panel.
#' The latter is obviously more important to control the final result of a plot.
#' This function can set the panel width/height of plot to a fixed value and rasterize it.
#'
#' @md
#' @inheritParams thisutils::log_message
#' @param x A ggplot object, a grob object, or a combined plot made by patchwork or cowplot package.
#' @param panel_index Specify the panel to be fixed.
#' If `NULL`, will fix all panels.
#' @param respect Whether row heights and column widths should respect each other.
#' @param width The desired width of the fixed panels.
#' @param height The desired height of the fixed panels.
#' @param margin The margin to add around each panel, in inches.
#' Default is `1`.
#' @param padding The padding to add around each panel, in inches.
#' Default is `0`.
#' @param units The units in which `height`, `width` and `margin` are given.
#' Can be `"mm"`, `"cm"`, `"in"`, etc. See [grid::unit].
#' @param raster Whether to rasterize the panel.
#' @param dpi Plot resolution.
#' @param return_grob Whether to return a grob object instead of a wrapped `patchwork` object.
#' Default is `FALSE`.
#' @param save `NULL` or the file name used to save the plot.
#' @param bg_color The background color of the plot.
#' @param ... Additional arguments passed to other functions.
#'
#' @return If `return_grob` is `TRUE`, returns a gtable object.
#' Otherwise, returns a patchwork object with fixed panel sizes.
#' The returned object has a `size` attribute containing width, height, and units.
#'
#' @export
#'
#' @examples
#' library(ggplot2)
#' p <- ggplot(
#'   data = mtcars, aes(x = mpg, y = wt, colour = cyl)
#' ) +
#'   geom_point() +
#'   facet_wrap(~gear, nrow = 2)
#' # fix the size of panel
#' panel_fix(
#'   p,
#'   width = 5,
#'   height = 3,
#'   units = "cm"
#' )
#' # rasterize the panel
#' panel_fix(
#'   p,
#'   width = 5,
#'   height = 3,
#'   units = "cm",
#'   raster = TRUE,
#'   dpi = 90
#' )
#'
#' # `panel_fix` will build and render the plot when input a ggplot object.
#' # so after `panel_fix`, the size of the object will be changed.
#' object.size(p)
#' object.size(
#'   panel_fix(
#'     p,
#'     width = 5,
#'     height = 3,
#'     units = "cm"
#'   )
#' )
panel_fix <- function(
  x = NULL,
  panel_index = NULL,
  respect = NULL,
  width = NULL,
  height = NULL,
  margin = 1,
  padding = 0,
  units = "in",
  raster = FALSE,
  dpi = 300,
  return_grob = FALSE,
  bg_color = "white",
  save = NULL,
  verbose = FALSE,
  ...
) {
  if (!inherits(x, "gtable")) {
    tryCatch(
      {
        gtable <- as_gtable(x)
      },
      error = function(error) {
        log_message(
          error, "\nCannot convert the x to a gtable object",
          message_type = "error"
        )
      }
    )
  } else {
    gtable <- x
  }
  args <- as.list(match.call())[-1]
  depth <- args[["depth"]]
  if (is.null(depth)) {
    depth <- 1
  }

  if (is.null(panel_index)) {
    non_zero <- grep(
      pattern = "zeroGrob",
      vapply(gtable$grobs, as.character, character(1)),
      invert = TRUE
    )
    panel_index <- grep(
      pattern = "panel|full",
      gtable[["layout"]][["name"]]
    )
    panel_index <- intersect(panel_index, non_zero)
  }
  if (length(panel_index) == 0 && length(gtable$grobs) == 1) {
    panel_index <- 1
  }
  add_margin <- TRUE
  for (i in panel_index) {
    geom_index <- grep(
      pattern = "GeomDrawGrob",
      names(gtable$grobs[[i]][["children"]])
    )
    if (length(geom_index) > 0) {
      log_message(
        "panel {.val {i}} is detected as generated by plot_grid",
        verbose = verbose
      )
      for (j in geom_index) {
        subgrob <- gtable$grobs[[i]][["children"]][[j]][["children"]][[1]][["children"]][[1]]

        if (length(subgrob$grobs[[1]][["children"]]) > 0 &&
          all(sapply(subgrob$grobs[[1]][["children"]], function(x) inherits(x, "recordedGrob")))) {
          subgrob <- panel_fix_overall(
            x = subgrob$grobs[[1]][["children"]],
            width = width,
            height = height,
            margin = padding,
            units = units,
            raster = raster,
            dpi = dpi,
            return_grob = TRUE
          )
        } else {
          subgrob <- panel_fix(
            x = subgrob,
            width = width,
            height = height,
            margin = padding,
            units = units,
            raster = raster,
            dpi = dpi,
            return_grob = TRUE,
            verbose = verbose,
            depth = depth + 1
          )
        }
        gtable$grobs[[i]][["children"]][[j]][["children"]][[1]][["children"]][[1]] <- subgrob
      }
      sum_width <- grid::convertWidth(
        sum(subgrob[["widths"]]),
        unitTo = units,
        valueOnly = TRUE
      ) / as.numeric(gtable$grobs[[i]][["children"]][[j]]$vp$width)
      sum_height <- grid::convertHeight(
        sum(subgrob[["heights"]]),
        unitTo = units,
        valueOnly = TRUE
      ) / as.numeric(gtable$grobs[[i]][["children"]][[j]]$vp$height)
      gtable <- panel_fix_overall(
        gtable,
        panel_index = i,
        width = sum_width,
        height = sum_height,
        margin = ifelse(depth == 1, margin, 0),
        units = units,
        raster = FALSE,
        return_grob = TRUE
      )
    } else if (gtable$grobs[[i]]$name == "layout" || inherits(x, "patchwork")) {
      log_message(
        "panel {.val {i}} is detected as generated by patchwork",
        verbose = verbose
      )
      # if (i == panel_index[1] && length(panel_index) > 1 && isTRUE(verbose)) {
      #   log_message("More than 2 panels detected. panel_fix may not work as expected.")
      # }
      subgrob <- gtable$grobs[[i]]
      if (length(subgrob[["children"]]) > 0 &&
        all(sapply(subgrob[["children"]], function(x) inherits(x, "recordedGrob")))) {
        subgrob <- panel_fix_overall(
          subgrob[["children"]],
          width = width,
          height = height,
          margin = 0,
          units = units,
          raster = raster,
          dpi = dpi,
          return_grob = TRUE
        )
      } else {
        subgrob <- panel_fix(
          subgrob,
          width = width,
          height = height,
          margin = 0,
          units = units,
          raster = raster,
          dpi = dpi,
          return_grob = TRUE,
          verbose = verbose,
          depth = depth + 1
        )
      }
      gtable$grobs[[i]] <- subgrob
      layout <- gtable$layout
      layout[["rowranges"]] <- lapply(
        seq_len(nrow(layout)),
        function(n) layout$t[n]:layout$b[n]
      )
      layout[["colranges"]] <- lapply(
        seq_len(nrow(layout)),
        function(n) layout$l[n]:layout$r[n]
      )
      p_row <- c(layout$t[i], layout$b[i])
      p_col <- c(layout$l[i], layout$r[i])
      background_index <- grep(
        pattern = "background", layout$name
      )
      background_index <- background_index[order(layout$z[background_index], decreasing = TRUE)]
      for (bgi in background_index) {
        if (all(p_row %in% layout[["rowranges"]][[bgi]]) && all(p_col %in% layout[["colranges"]][[bgi]])) {
          p_background_index <- bgi
          break
        }
      }
      gtable <- gtable::gtable_add_rows(
        gtable,
        heights = grid::unit(padding, units),
        pos = layout$t[p_background_index] - 1
      )
      gtable <- gtable::gtable_add_rows(
        gtable,
        heights = grid::unit(padding, units),
        pos = layout$b[p_background_index]
      )
      gtable <- gtable::gtable_add_cols(
        gtable,
        widths = grid::unit(padding, units),
        pos = layout$l[p_background_index] - 1
      )
      gtable <- gtable::gtable_add_cols(
        gtable,
        widths = grid::unit(padding, units),
        pos = layout$r[p_background_index]
      )
      sum_width <- grid::convertWidth(
        sum(subgrob[["widths"]]),
        unitTo = units,
        valueOnly = TRUE
      )
      sum_height <- grid::convertHeight(
        sum(subgrob[["heights"]]),
        unitTo = units,
        valueOnly = TRUE
      )

      gtable <- panel_fix_overall(
        gtable,
        panel_index = i,
        width = sum_width,
        height = sum_height,
        margin = ifelse(depth == 1 & add_margin, margin, 0),
        units = units,
        raster = FALSE,
        respect = TRUE,
        return_grob = TRUE
      )
      if (depth == 1 & add_margin) {
        add_margin <- FALSE
      }
    } else {
      gtable <- panel_fix_overall(
        gtable,
        panel_index = i,
        width = width,
        height = height,
        margin = margin,
        units = units,
        raster = raster,
        dpi = dpi,
        return_grob = TRUE
      )
    }
  }

  if (!is.null(respect)) {
    gtable$respect <- respect
  }

  if (isTRUE(return_grob)) {
    return(gtable)
  } else {
    p <- patchwork::wrap_plots(gtable) +
      theme(
        plot.background = element_rect(
          fill = bg_color, color = bg_color
        )
      )
    if (units != "null") {
      plot_width <- grid::convertWidth(
        sum(gtable[["widths"]]),
        unitTo = units,
        valueOnly = TRUE
      )
      plot_height <- grid::convertHeight(
        sum(gtable[["heights"]]),
        unitTo = units,
        valueOnly = TRUE
      )
      attr(p, "size") <- list(
        width = plot_width,
        height = plot_height,
        units = units
      )
    }

    if (!is.null(save) && is.character(save) && nchar(save) > 0) {
      if (units == "null") {
        log_message(
          "{.arg units} can not be 'null' if want to save the plot",
          message_type = "error"
        )
      }
      filename <- normalizePath(save)
      log_message(
        "Save the plot to the file: {.file {filename}}",
        verbose = verbose
      )
      if (!dir.exists(dirname(filename))) {
        dir.create(dirname(filename), recursive = TRUE, showWarnings = FALSE)
      }
      ggplot2::ggsave(
        plot = p,
        filename = filename,
        width = plot_width,
        height = plot_height,
        units = units,
        dpi = dpi,
        limitsize = FALSE
      )
    }
    return(p)
  }
}

#' @rdname panel_fix
#' @export
panel_fix_overall <- function(
  x,
  panel_index = NULL,
  respect = NULL,
  width = NULL,
  height = NULL,
  margin = 1,
  units = "in",
  raster = FALSE,
  dpi = 300,
  return_grob = FALSE,
  bg_color = "white",
  save = NULL,
  verbose = TRUE
) {
  if (!inherits(x, "gtable")) {
    if (inherits(x, "gTree")) {
      x <- x[["children"]]
    }
    tryCatch(
      {
        gtable <- as_gtable(x)
      },
      error = function(error) {
        log_message(
          error, "\nCannot convert the x to a gtable object",
          message_type = "error"
        )
      }
    )
  } else {
    gtable <- x
  }

  if (is.null(panel_index)) {
    non_zero <- grep(
      pattern = "zeroGrob",
      vapply(
        gtable$grobs, as.character, character(1)
      ), invert = TRUE
    )
    panel_index <- grep("panel|full", gtable[["layout"]][["name"]])
    panel_index <- intersect(panel_index, non_zero)
  }
  if (length(panel_index) == 0 && length(gtable$grobs) == 1) {
    panel_index <- 1
  }
  if (!length(width) %in% c(0, 1, length(panel_index)) || !length(height) %in% c(0, 1, length(panel_index))) {
    log_message(
      "The length of 'width' and 'height' must be 1 or the length of panels.",
      message_type = "error"
    )
  }

  if (inherits(x, "gList")) {
    panel_index <- 1
    panel_index_h <- panel_index_w <- list(1)
    w_comp <- h_comp <- list(grid::unit(1, "null"))
    w <- h <- list(grid::unit(1, "null"))
  } else if (length(panel_index) > 0) {
    panel_index_w <- panel_index_h <- list()
    w_comp <- h_comp <- list()
    w <- h <- list()
    for (i in seq_along(panel_index)) {
      index <- panel_index[i]
      panel_index_h[[i]] <- sort(
        unique(c(
          gtable[["layout"]][["t"]][index],
          gtable[["layout"]][["b"]][index]
        ))
      )
      panel_index_w[[i]] <- sort(
        unique(c(
          gtable[["layout"]][["l"]][index],
          gtable[["layout"]][["r"]][index]
        ))
      )
      w_comp[[i]] <- gtable[["widths"]][seq(min(panel_index_w[[i]]), max(panel_index_w[[i]]))]
      h_comp[[i]] <- gtable[["heights"]][seq(min(panel_index_h[[i]]), max(panel_index_h[[i]]))]

      if (length(w_comp[[i]]) == 1) {
        w[[i]] <- w_comp[[i]]
      } else if (length(w_comp[[i]]) > 1 && any(grid::unitType(w_comp[[i]]) == "null")) {
        w[[i]] <- grid::unit(1, units = "null")
      } else {
        w[[i]] <- sum(w_comp[[i]])
      }
      if (length(h_comp[[i]]) == 1) {
        h[[i]] <- h_comp[[i]]
      } else if (length(h_comp[[i]]) > 1 && any(grid::unitType(h_comp[[i]]) == "null")) {
        h[[i]] <- grid::unit(1, units = "null")
      } else {
        h[[i]] <- sum(h_comp[[i]])
      }
    }
  } else {
    log_message(
      "No panel detected",
      message_type = "error"
    )
  }

  if (units != "null") {
    raw_w <- sapply(
      w, function(x) {
        grid::convertWidth(x, unitTo = units, valueOnly = TRUE)
      }
    )
    raw_h <- sapply(
      h, function(x) {
        grid::convertHeight(x, unitTo = units, valueOnly = TRUE)
      }
    )
    for (i in seq_along(w)) {
      if (grid::unitType(w[[i]]) == "null" || grid::convertUnit(w[[i]], unitTo = "cm", valueOnly = TRUE) < 1e-10) {
        raw_w[i] <- 0
      }
    }
    for (i in seq_along(h)) {
      if (grid::unitType(h[[i]]) == "null" || grid::convertUnit(h[[i]], unitTo = "cm", valueOnly = TRUE) < 1e-10) {
        raw_h[i] <- 0
      }
    }
    if (isTRUE(gtable$respect)) {
      raw_aspect <- sapply(h, as.vector) / sapply(w, as.vector)
    } else {
      if (all(raw_w != 0) && all(raw_h != 0)) {
        raw_aspect <- raw_h / raw_w
      } else {
        raw_aspect <- grid::convertHeight(
          grid::unit(1, "npc"), "cm",
          valueOnly = TRUE
        ) / grid::convertWidth(grid::unit(1, "npc"), "cm", valueOnly = TRUE)
      }
    }

    if (is.null(width) && is.null(height)) {
      width <- raw_w
      height <- raw_h
      if (all(width == 0) && all(height == 0)) {
        width <- grid::convertWidth(
          grid::unit(1, "npc"), units,
          valueOnly = TRUE
        )
        height <- grid::convertHeight(
          grid::unit(1, "npc"), units,
          valueOnly = TRUE
        )
        if (isTRUE(gtable$respect)) {
          if (raw_aspect <= 1) {
            height <- width * raw_aspect
          } else {
            width <- height / raw_aspect
          }
        }
      }
    }

    for (i in seq_along(raw_aspect)) {
      if (is.finite(raw_aspect[i]) && raw_aspect[i] != 0) {
        if (is.null(width[i]) || is.na(width[i]) || width[i] == 0) {
          width[i] <- height[i] / raw_aspect[i]
        }
        if (is.null(height[i]) || is.na(height[i]) || height[i] == 0) {
          height[i] <- width[i] * raw_aspect[i]
        }
      }
    }

    for (i in seq_along(width)) {
      if (inherits(width[i], "unit")) {
        width[i] <- grid::convertWidth(
          width[i],
          unitTo = units,
          valueOnly = TRUE
        )
      }
    }
    for (i in seq_along(height)) {
      if (inherits(height[i], "unit")) {
        height[i] <- grid::convertHeight(
          height[i],
          unitTo = units,
          valueOnly = TRUE
        )
      }
    }
  }

  if (length(width) == 1) {
    width <- rep(width, length(panel_index))
  }
  if (length(height) == 1) {
    height <- rep(height, length(panel_index))
  }
  for (i in seq_along(panel_index)) {
    if (!is.null(width)) {
      width_unit <- width[i] / length(w_comp[[i]])
      gtable[["widths"]][seq(min(panel_index_w[[i]]), max(panel_index_w[[i]]))] <- rep(
        grid::unit(width_unit, units = units), length(w_comp[[i]])
      )
    }
    if (!is.null(height)) {
      height_unit <- height[i] / length(h_comp[[i]])
      gtable[["heights"]][seq(min(panel_index_h[[i]]), max(panel_index_h[[i]]))] <- rep(
        grid::unit(height_unit, units = units), length(h_comp[[i]])
      )
    }
  }
  gtable <- gtable::gtable_add_padding(
    gtable,
    padding = grid::unit(margin, units = units)
  )

  if (isTRUE(raster)) {
    for (i in seq_along(panel_index)) {
      index <- panel_index[i]
      g <- g_new <- gtable$grobs[[index]]
      vp <- g$vp
      children_order <- g$childrenOrder
      if (is.null(g$vp)) {
        g$vp <- grid::viewport()
      }

      for (j in seq_along(g[["children"]])) {
        child <- g[["children"]][[j]]
        child_nm <- names(g[["children"]])[j]
        if (!is.null(child$vp) ||
          any(grepl("(text)|(label)", child_nm)) ||
          any(grepl("(text)|(segments)|(legend)", class(child$list[[1]])))) {
          zero <- ggplot2::zeroGrob()
          zero$name <- g[["children"]][[j]]$name
          g[["children"]][[j]] <- zero
        } else if (inherits(child$list[[1]], "grob") || is.null(child$list[[1]])) {
          g_new[["children"]][[j]] <- ggplot2::zeroGrob()
        }
      }
      temp <- tempfile(fileext = "png")
      ragg::agg_png(
        temp,
        width = width[i],
        height = height[i],
        bg = "transparent",
        res = dpi,
        units = units
      )
      grid::grid.draw(g)
      grDevices::dev.off()
      g_ras <- grid::rasterGrob(
        png::readPNG(temp, native = TRUE)
      )
      unlink(temp)
      g <- grid::addGrob(g_new, g_ras)
      g$vp <- vp
      g$childrenOrder <- c(g_ras$name, children_order)
      gtable$grobs[[index]] <- g
    }
  }

  if (!is.null(respect)) {
    gtable$respect <- respect
  }

  if (isTRUE(return_grob)) {
    return(gtable)
  } else {
    p <- patchwork::wrap_plots(gtable) +
      theme(
        plot.background = element_rect(
          fill = bg_color, color = bg_color
        )
      )
    if (units != "null") {
      plot_width <- grid::convertWidth(
        sum(gtable[["widths"]]),
        unitTo = units,
        valueOnly = TRUE
      )
      plot_height <- grid::convertHeight(
        sum(gtable[["heights"]]),
        unitTo = units,
        valueOnly = TRUE
      )
      attr(p, "size") <- list(
        width = plot_width,
        height = plot_height,
        units = units
      )
    }

    if (!is.null(save) && is.character(save) && nchar(save) > 0) {
      if (units == "null") {
        log_message(
          "{.arg units} can not be 'null' if want to save the plot",
          message_type = "error"
        )
      }
      filename <- normalizePath(save)
      log_message(
        "Save plot to: {.file {filename}}",
        verbose = verbose
      )
      if (!dir.exists(dirname(filename))) {
        dir.create(
          dirname(filename),
          recursive = TRUE,
          showWarnings = FALSE
        )
      }
      ggplot2::ggsave(
        plot = p,
        filename = filename,
        width = plot_width,
        height = plot_height,
        units = units,
        dpi = dpi,
        limitsize = FALSE
      )
    }
    return(p)
  }
}