R/align_to_hm.R

In csgroen/ggheatmapper: Tile-able heatmaps using ggheatmap and patchwork

#' Align plot to ggheatmap
#'
#' Align a ggheatmap object to another ggplot that shares either its columns or
#' rows to form complex panels. This function is called internally by add_tracks,
#' and also exported by ggheatmap.
#'
#' @param gghm A object with class ggheatmap, generated by [ggheatmap()].
#' @param gplot Another object of class ggplot, that shares either columns
#' (if `pos='top'` or `pos='bottom'`) or rows (if `pos='left'`) with the heatmap.
#' Note: to create the gplot, make sure you use [gghmData()] to get the data, or
#' alternatively [get_rowLevels()] or [get_colLevels()].
#' @param pos One of 'top', 'bottom' or 'left'.
#' @param newplt_size_prop A float between 0 and 1, that indicates the proportion
#' of the height (if `pos='top'` or `pos='bottom'`) or width (if `pos='left'`) that
#' the new plot will occupy in the panel.
#' @param legend_action A string specifying how guides should be treated in the layout.
#' See: guides in [patchwork::plot_layout()].
#' @param tag_level A string ('keep' or 'new') to indicate how auto-tagging should behave.
#' See [patchwork::plot_annotation].
#'
#' @import tidyverse patchwork
#' @importFrom rlist list.prepend list.append
#' @importFrom stringr str_split str_sub
#' @export
#' @return A joined panel object of class `ggheatmap`.
align_to_hm <- function(gghm, gplot,
                        pos = "bottom",
                        newplt_size_prop = 0.5,
                        legend_action = NULL,
                        tag_level = 'new') {
    if(! "ggheatmap" %in% class(gghm)) {
        stop("`gghm` must be of class `ggheatmap`.")
    }
    gplot <- wrap_plots(gplot)
    params <- gghm$gghm$params
    plots <- gghm$gghm$plots
    design <- gghm$gghm$design

    plots <- list.append(plots, gplot + plot_layout(tag_level = tag_level))
    #-- Align plots
    if(pos == "bottom") {
        params$heights <- params$heights * (1-newplt_size_prop)
        params$heights <- c(params$heights, newplt_size_prop)
        new_col <- paste0(rep("#", length(params$widths)-1), LETTERS[length(plots)])
        new_design <- paste0(design,"\n",new_col)

    } else if (pos == "top") {
        params$heights <- params$heights * (1-newplt_size_prop)
        params$heights <- c(params$heights[1:2], newplt_size_prop,
                            params$heights[3:length(params$heights)])
        #-- Make design
        new_col <- paste0(rep("#", length(params$widths)-1), LETTERS[length(plots)])
        design_rows <- str_split(design, "\n")[[1]]
        new_design <- paste(c(design_rows[1:2], new_col,
                              design_rows[3:length(design_rows)]),
                            collapse = '\n')
        params$hm_row <- params$hm_row + 1
    } else if (pos == "left") {
        params$widths <- params$widths * (1-newplt_size_prop)
        params$widths <- c(params$widths[1], newplt_size_prop,
                           params$widths[2:length(params$widths)])
        #-- Make design
        design_rows <- str_split(design, "\n")[[1]]
        design_rows[params$hm_row] <- paste0("C", LETTERS[length(plots)], str_sub(design_rows[params$hm_row], 2))
        idx <- setdiff(1:length(design_rows), params$hm_row)
        design_rows[idx] <- paste0(str_sub(design_rows[idx],1,1), "#", str_sub(design_rows[idx],2))
        new_design <- paste(design_rows, collapse = "\n")

        params$hm_col <- params$hm_col + 1

    } else {
        params$widths <- params$widths * (1-newplt_size_prop)
        params$widths <- c(params$widths, newplt_size_prop)
        #-- Make design
        design_rows <- str_split(design, "\n")[[1]]
        design_rows[params$hm_row] <- paste0(design_rows[params$hm_row], LETTERS[length(plots)])
        idx <- setdiff(1:length(design_rows), params$hm_row)
        design_rows[idx] <- paste0(design_rows[idx], "#")
        new_design <- paste(design_rows, collapse = "\n")
    }
    #-- Re-assemble with new design
    new_gghm <- wrap_plots(plots, design = new_design,
                           widths = params$widths,
                           heights = params$heights,
                           guides = legend_action)

    #-- Update params
    class(new_gghm) <- append(class(new_gghm), "ggheatmap")
    new_gghm$data <- gghm$data
    new_gghm$gghm <- list(plots = plots,
                          params = params,
                          design = new_design,
                          row_levels = gghm$gghm$row_levels,
                          col_levels = gghm$gghm$col_levels,
                          line_geom = gghm$gghm$line_geom,
                          hclust = gghm$gghm$hclust)


    return(new_gghm)

}