fix Align object plot will modify in place

Yunuuuu · Nov 7, 2024 · fd2f83c · fd2f83c
1 parent 9bf9e81
commit fd2f83c
Show file tree

Hide file tree

Showing 10 changed files with 97 additions and 80 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -104,9 +104,9 @@ S3method(ggalign_build,alignpatches)
 S3method(ggalign_build,ggplot)
 S3method(ggalign_gtable,alignpatches)
 S3method(ggalign_gtable,ggplot)
-S3method(ggalign_stat,Align)
 S3method(ggalign_stat,QuadLayout)
 S3method(ggalign_stat,StackLayout)
+S3method(ggalign_stat,align)
 S3method(ggalign_stat,default)
 S3method(ggoncoplot,"NULL")
 S3method(ggoncoplot,default)
@@ -152,7 +152,7 @@ S3method(patch,patchwork)
 S3method(patch,pheatmap)
 S3method(patch,recordedplot)
 S3method(patch,trellis)
-S3method(plot,Align)
+S3method(plot,align)
 S3method(plot,ggalign_area)
 S3method(plot_add,plot_data)
 S3method(plot_add,plot_theme)
@@ -174,7 +174,7 @@ S3method(quad_build,HeatmapLayout)
 S3method(quad_build,QuadLayout)
 S3method(quad_free,default)
 S3method(quad_free,uneval)
-S3method(quad_layout_add,Align)
+S3method(quad_layout_add,align)
 S3method(quad_layout_add,anno_init)
 S3method(quad_layout_add,data.frame)
 S3method(quad_layout_add,default)
@@ -206,8 +206,8 @@ S3method(quad_layout_subtract,quad_anno)
 S3method(stack_align,default)
 S3method(stack_free,default)
 S3method(stack_free,uneval)
-S3method(stack_layout_add,Align)
 S3method(stack_layout_add,QuadLayout)
+S3method(stack_layout_add,align)
 S3method(stack_layout_add,anno_init)
 S3method(stack_layout_add,default)
 S3method(stack_layout_add,free_gg)

diff --git a/R/align-.R b/R/align-.R
@@ -115,50 +115,60 @@ align <- function(align_class, params, data, size = NULL, controls = NULL,
         }
     }
 
-    # wrap all elements into this annotation ---------------------
-    ggproto(
-        NULL,
-        align_class,
-        isLock = FALSE,
-        # Following fields will be initialzed when added into the layout
-        # and will be saved and accessed across the plot rendering process
-        statistics = NULL,
-        direction = NULL,
-        position = NULL,
-        plot = NULL,
-        data = NULL,
-        params = NULL,
-        labels = NULL, # the original `vec_names()` of the `input_data`
+    new_align(
+        # wrap all elements into this annotation ---------------------
+        Object = ggproto(
+            NULL,
+            align_class,
+            isLock = FALSE,
+            # Following fields will be initialzed when added into the layout
+            # and will be saved and accessed across the plot rendering process
+            direction = NULL,
+            position = NULL,
+            params = NULL, # `$setup_params` method
+            data = NULL, # $setup_data method
+            statistics = NULL, # `$compute` method
+            labels = NULL, # the original `vec_names()` of the `input_data`
 
-        # user input -------------------------------
-        size = size,
-        # should we allow user switch between different plot with a string name?
-        # Should I remove "name" argument from the user input?
-        active = active,
-        # use `NULL` if this align don't require any data
-        # use `waiver()` to inherit from the layout data
-        input_data = data,
+            # use `NULL` if this align don't require any data
+            # use `waiver()` to inherit from the layout data
+            input_data = data,
+
+            # collect parameters
+            input_params = params[vec_set_intersect(names(params), all)],
+
+            # used to provide error message
+            call = call
+        ),
+        no_axes = no_axes,
         controls = controls,
         facet = facet,
         limits = limits,
-        no_axes = no_axes,
 
-        # collect parameters
-        input_params = params[intersect(names(params), all)],
+        # user input -------------------------------
+        size = size,
 
-        # used to provide error message
-        call = call
+        # should we allow user switch between different plot with a string name?
+        # Should I remove "name" argument from the user input?
+        active = active,
+        plot = NULL
     )
 }
 
+#' We create the align entity when initialize the Align object.
+#' @noRd
+new_align <- function(Object, ..., plot = NULL) {
+    structure(list(Object = Object, ..., plot = plot), class = "align")
+}
+
+is_align <- function(x) inherits(x, "align")
+
 #' @export
 #' @keywords internal
-plot.Align <- function(x, ...) {
+plot.align <- function(x, ...) {
     cli::cli_abort("You cannot plot {.obj_type_friendly {x}} object directly")
 }
 
-is_align <- function(x) inherits(x, "Align")
-
 #' @details
 #' Each of the `Align*` objects is just a [`ggproto()`][ggplot2::ggproto]
 #' object, descended from the top-level `Align`, and each implements various
@@ -185,7 +195,7 @@ Align <- ggproto("Align",
             align_method_params(self$ggplot, character()),
             align_method_params(
                 self$draw,
-                c("panel", "index", "extra_panel", "extra_index")
+                c("plot", "panel", "index", "extra_panel", "extra_index")
             ),
             self$extra_params
         )
@@ -261,7 +271,7 @@ Align <- ggproto("Align",
     # Following methods will be executed when building plot with the final
     # heatmap layout you shouldn't modify the `Align` object when drawing,
     # since all of above process will only run once.
-    draw = function(self, panel, index, extra_panel, extra_index) self$plot
+    draw = function(self, plot, panel, index, extra_panel, extra_index) plot
 )
 
 # Used to lock the `Align` object

diff --git a/R/align-add.R b/R/align-add.R
@@ -31,7 +31,7 @@ align_add.Coord <- function(object, align, object_name) {
 align_add.ggalign_option <- function(object, align, object_name) {
     name <- ggalign_option_name(object)
     align$controls[name] <- list(update_option(
-        object, .subset2(align$controls, name), object_name
+        object, .subset2(.subset2(align, "controls"), name), object_name
     ))
     align
 }
diff --git a/R/align-build.R b/R/align-build.R
@@ -1,22 +1,24 @@
 #' @importFrom ggplot2 theme element_blank
 #' @importFrom rlang inject
-align_build <- function(x, panel, index, controls, extra_layout) {
+align_build <- function(align, panel, index, controls, extra_layout) {
+    # let `Align` to determine how to draw
+    # 1. add default layer
+    # 2. add plot data
+    object <- .subset2(align, "Object") # `Align` object
+
     # we lock the Align object to prevent user from modifying this object
     # in `$draw` method, we shouldn't do any calculations in `$draw` method
-    x$lock()
-    on.exit(x$unlock())
+    object$lock()
+    on.exit(object$unlock())
 
-    # let `align` to determine how to draw
-    # 1. add default layer
-    # 2. add plot data
-    direction <- .subset2(x, "direction")
-    params <- .subset2(x, "params")
+    direction <- .subset2(object, "direction")
+    params <- .subset2(object, "params")
     draw_params <- params[
-        intersect(
+        vec_set_intersect(
             names(params),
             align_method_params(
-                x$draw,
-                c("panel", "index", "extra_panel", "extra_index")
+                object$draw,
+                c("plot", "panel", "index", "extra_panel", "extra_index")
             )
         )
     ]
@@ -27,12 +29,14 @@ align_build <- function(x, panel, index, controls, extra_layout) {
         extra_panel <- .subset2(extra_layout, "panel")
         extra_index <- .subset2(extra_layout, "index")
     }
-    plot <- inject(x$draw(
-        panel, index, extra_panel, extra_index, !!!draw_params
+    plot <- inject(object$draw(
+        .subset2(align, "plot"),
+        panel, index, extra_panel,
+        extra_index, !!!draw_params
     ))
 
     # only when user use the internal facet, we'll setup the limits
-    if (.subset2(x, "facet")) {
+    if (.subset2(align, "facet")) {
         # set up facets
         if (nlevels(panel) > 1L) {
             default_facet <- switch_direction(
@@ -53,7 +57,7 @@ align_build <- function(x, panel, index, controls, extra_layout) {
         }
         layout <- list(
             panel = panel, index = index,
-            labels = .subset2(x, "labels")
+            labels = .subset2(object, "labels")
         )
         plot <- plot + align_melt_facet(plot$facet, default_facet, direction) +
             switch_direction(
@@ -63,7 +67,7 @@ align_build <- function(x, panel, index, controls, extra_layout) {
             )
 
         # set up coord limits to align each observation
-        if (.subset2(x, "limits")) {
+        if (.subset2(align, "limits")) {
             plot <- plot +
                 switch_direction(
                     direction,
@@ -73,12 +77,12 @@ align_build <- function(x, panel, index, controls, extra_layout) {
         }
     }
     # remove axis titles, text, ticks used for alignment
-    if (isTRUE(.subset2(x, "no_axes"))) {
+    if (isTRUE(.subset2(align, "no_axes"))) {
         controls$plot_theme <- .subset2(controls, "plot_theme") +
             theme_no_axes(switch_direction(direction, "y", "x"))
     }
     plot <- plot_add_controls(plot, controls)
-    list(plot = plot, size = .subset2(x, "size"))
+    list(plot = plot, size = .subset2(align, "size"))
 }
 
 #' @importFrom ggplot2 ggproto

diff --git a/R/align-dendrogram.R b/R/align-dendrogram.R
@@ -393,7 +393,7 @@ AlignDendro <- ggproto("AlignDendro", Align,
                 ggplot2::labs(y = "height")
             )
     },
-    draw = function(self, panel, index, extra_panel, extra_index,
+    draw = function(self, plot, panel, index, extra_panel, extra_index,
                     # other argumentds
                     plot_cut_height, center, type, root) {
         direction <- .subset2(self, "direction")
@@ -464,7 +464,6 @@ AlignDendro <- ggproto("AlignDendro", Align,
             )
             node <- rename(node, c(x = "y", y = "x"))
         }
-        plot <- .subset2(self, "plot")
         # we do some tricks, since ggplot2 won't remove the attributes
         # we attach the `edge` data in `ggalign` attribute
         plot$data <- structure(node, ggalign = list(edge = edge))

diff --git a/R/align-gg.R b/R/align-gg.R
@@ -169,7 +169,7 @@ AlignGG <- ggproto("AlignGG", Align,
     },
 
     #' @importFrom stats reorder
-    draw = function(self, panel, index, extra_panel, extra_index) {
+    draw = function(self, plot, panel, index, extra_panel, extra_index) {
         data <- .subset2(self, "data")
         direction <- .subset2(self, "direction")
         axis <- to_coord_axis(direction)
@@ -204,7 +204,6 @@ AlignGG <- ggproto("AlignGG", Align,
                 order = FALSE
             )
         }
-        plot <- .subset2(self, "plot")
         plot$data <- restore_attr_ggalign(ans, data)
         plot
     }

diff --git a/R/align-initialize.R b/R/align-initialize.R
@@ -1,8 +1,9 @@
 #' `Align` is an environment, it won't be copied, and will modify in place
 #' @noRd
-align_initialize <- function(object, direction, position,
+align_initialize <- function(align, direction, position,
                              layout_data, layout_panel, layout_index,
                              layout_nobs, object_name) {
+    object <- .subset2(align, "Object")
     object$direction <- direction
     object$position <- position
     input_data <- .subset2(object, "input_data")
@@ -49,11 +50,21 @@ align_initialize <- function(object, direction, position,
     object$params <- params
 
     # prepare the data -------------------------------
-    ans <- align_initialize_layout(
+    layout <- align_initialize_layout(
         object, layout_nobs, direction,
         layout_panel, layout_index, object_name
     )
-    c(vec_set_names(ans, c("panel", "index")), list(nobs = layout_nobs))
+
+    # in the finally, Let us initialize the annotation plot -----
+    # must return a ggplot object
+    ggplot_params <- params[
+        vec_set_intersect(
+            names(params),
+            align_method_params(object$ggplot, character())
+        )
+    ]
+    plot <- inject(object$ggplot(!!!ggplot_params))
+    list(layout = vec_c(layout, list(nobs = layout_nobs)), plot = plot)
 }
 
 #' @importFrom rlang inject
@@ -66,15 +77,15 @@ align_initialize_layout <- function(object, layout_nobs, direction,
 
     # compute statistics ---------------------------------
     compute_params <- params[
-        intersect(names(params), align_method_params(object$compute))
+        vec_set_intersect(names(params), align_method_params(object$compute))
     ]
     object$statistics <- inject(
         object$compute(layout_panel, layout_index, !!!compute_params)
     )
 
     # make the new layout -------------------------------
     layout_params <- params[
-        intersect(names(params), align_method_params(object$layout))
+        vec_set_intersect(names(params), align_method_params(object$layout))
     ]
     layout <- inject(
         object$layout(layout_panel, layout_index, !!!layout_params)
@@ -146,19 +157,8 @@ align_initialize_layout <- function(object, layout_nobs, direction,
         ), call = call)
     }
 
-    # in the finally, Let us initialize the annotation plot -----
-    # must return a ggplot object
-    ggplot_params <- params[
-        intersect(
-            names(params),
-            align_method_params(object$ggplot, character())
-        )
-    ]
-    p <- inject(object$ggplot(!!!ggplot_params))
-    object$plot <- p
-
     # add annotation -------------------------------------
-    list(new_panel, new_index)
+    list(panel = new_panel, index = new_index)
 }
 
 ############################################################

diff --git a/R/layout-.R b/R/layout-.R
@@ -236,7 +236,9 @@ ggalign_stat.StackLayout <- function(x, what, ...) {
 }
 
 #' @export
-ggalign_stat.Align <- function(x, ...) .subset2(x, "statistics")
+ggalign_stat.align <- function(x, ...) {
+    .subset2(.subset2(x, "align"), "statistics")
+}
 
 #' @export
 ggalign_stat.default <- function(x, ...) {

diff --git a/R/layout-quad-add.R b/R/layout-quad-add.R
@@ -45,7 +45,7 @@ quad_layout_add.data.frame <- quad_layout_add.matrix
 
 ##################################################################
 #' @export
-quad_layout_add.Align <- function(object, quad, object_name) {
+quad_layout_add.align <- function(object, quad, object_name) {
     if (is.null(position <- quad@active)) {
         cli::cli_abort(c(
             "Cannot add {.var {object_name}} to {.fn {quad@name}}",