From 706aafa35447f1551cc9dee9c89f34a19e414b78 Mon Sep 17 00:00:00 2001
From: jiajic <72078254+jiajic@users.noreply.github.com>
Date: Fri, 20 Oct 2023 11:19:50 -0400
Subject: [PATCH] update: `sankeyPlot`

- make `sankeyPlot` a generic function.
- Allow either simple sankey creation from a single set of metadata by providing two column names to `x` param or complex creation if a `giottoSankeyPlan` object is supplied instead
---
 NAMESPACE         |   2 +-
 R/plot_sankey.R   | 184 ++++++++++++++++++++++++++++++++++------------
 man/sankeyPlot.Rd |  35 ++++++++-
 3 files changed, 170 insertions(+), 51 deletions(-)

diff --git a/NAMESPACE b/NAMESPACE
index ccdf4cb..e4b51ba 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -34,7 +34,6 @@ export(plotUMAP_2D)
 export(plotUMAP_3D)
 export(plot_output_handler)
 export(sankeyLabel)
-export(sankeyPlot)
 export(sankeySet)
 export(sankeySetAddresses)
 export(set_default_color_continuous)
@@ -81,6 +80,7 @@ export(violinPlot)
 exportClasses(giottoSankeyPlan)
 exportMethods("+")
 exportMethods("sankeyRelate<-")
+exportMethods(sankeyPlot)
 exportMethods(sankeyRelate)
 import(GiottoClass)
 import(GiottoUtils)
diff --git a/R/plot_sankey.R b/R/plot_sankey.R
index 6fcf5bf..73abcfc 100644
--- a/R/plot_sankey.R
+++ b/R/plot_sankey.R
@@ -47,13 +47,13 @@ setMethod('show', signature = 'giottoSankeyPlan', function(object) {
   cat('\n')
 })
 
-# generics
+# generics ####
 setGeneric('sankeyRelate', function(x, ...) standardGeneric('sankeyRelate'))
 setGeneric('sankeyRelate<-', function(x, add, value) standardGeneric('sankeyRelate<-'))
+setGeneric('sankeyPlot', function(gobject, x, ...) standardGeneric('sankeyPlot'))
 
 
-
-# methods
+# methods ####
 
 
 #' @name sankeyRelate
@@ -442,14 +442,24 @@ sankey_relation_pair = function(g, gsp, rel_idx, node_idx_start = 0) {
 #' @name sankeyPlot
 #' @description
 #' Create a sankey plot from a giotto object. Pulls from information in the
-#' metadata.
+#' metadata. Simple 1 to 1 sankeys can be generated from a single spatial unit
+#' and feature type using the `spat_unit`, `feat_type`, `meta_type`, `cols`,
+#' and (optionally) `idx` params. More complex and cross spatial unit/feature
+#' type sankeys can be set up using the `sankey_plan` param which accepts a
+#' `giottoSankeyPlan` object.
 #' @inheritParams data_access_params
 #' @inheritDotParams networkD3::sankeyNetwork -Links -Nodes -Source -Target -Value -NodeID
-#' @param sankey_plan giottoSankeyPlan object
+#' @param x giottoSankeyPlan object or character vector referring to source and
+#' target columns in metadata
 #' @param meta_type build sankey on cell or feature metadata
+#' @param spat_unit spatial unit of metadata
+#' @param feat_type feature type of metadata
+#' @param meta_type whether to use cell or feature metadata
+#' @param idx table subset index for 1 to 1 comparisons
 #' @examples
 #' \dontrun{
 #' g = GiottoData::loadGiottoMini("vizgen")
+#' # with giottoSankeyPlan
 #' leiden = sankeySet(spat_unit = 'aggregate',
 #'                    feat_type = 'rna',
 #'                    col = 'leiden_clus')
@@ -459,51 +469,133 @@ sankey_relation_pair = function(g, gsp, rel_idx, node_idx_start = 0) {
 #' plan = leiden + louvain
 #' sankeyRelate(plan) = c(0,1)
 #' sankeyPlot(g, plan)
+#'
+#' # with single set of metadata
+#' activeSpatUnit(g) = 'aggregate'
+#' sankeyPlot(g, c('louvain_clus', 'leiden_clus'))
 #' }
-#' @export
 #' @keywords plotting sankey
-sankeyPlot = function(gobject,
-                      sankey_plan,
-                      meta_type = c('cell', 'feat'),
-                      ...) {
-  checkmate::assert_class(gobject, 'giotto')
-  GiottoUtils::package_check("networkD3")
-  meta_type = match.arg(meta_type, choices = c('cell', 'feat'))
-  sankey_plan@data_type = meta_type
-
-  # iterate through sankey relations in the giottoSankeyPlan
-  node_idx_start = 0
-  links_dt = data.table::data.table()
-  nodes = c()
-
-  for (rel_i in seq(nrow(sankeyRelate(sankey_plan)))) {
-
-    rel_data = sankey_relation_pair(
-      g = gobject,
-      gsp = sankey_plan,
-      rel_idx = rel_i,
-      node_idx_start = node_idx_start
-    )
 
-    # append data
-    links_dt = rbind(links_dt, rel_data$links)
-    nodes = c(nodes, rel_data$nodes)
 
-    # update start index
-    node_idx_start = links_dt[, max(target)]
+#' @rdname sankeyPlot
+#' @export
+setMethod(
+  'sankeyPlot',
+  signature(gobject = 'giotto',
+            x = 'giottoSankeyPlan'),
+  function(gobject,
+           x,
+           meta_type = c('cell', 'feat'),
+           ...) {
+    checkmate::assert_class(gobject, 'giotto')
+    GiottoUtils::package_check("networkD3")
+    meta_type = match.arg(meta_type, choices = c('cell', 'feat'))
+    x@data_type = meta_type
+
+    # iterate through sankey relations in the giottoSankeyPlan
+    node_idx_start = 0
+    links_dt = data.table::data.table()
+    nodes = c()
+
+    for (rel_i in seq(nrow(sankeyRelate(x)))) {
+
+      rel_data = sankey_relation_pair(
+        g = gobject,
+        gsp = x,
+        rel_idx = rel_i,
+        node_idx_start = node_idx_start
+      )
+
+      # append data
+      links_dt = rbind(links_dt, rel_data$links)
+      nodes = c(nodes, rel_data$nodes)
+
+      # update start index
+      node_idx_start = links_dt[, max(target)]
+    }
+
+    # create nodes table
+    nodes = data.table::data.table(name = nodes)
+
+    networkD3::sankeyNetwork(
+      Links = links_dt,
+      Nodes = nodes,
+      Source = 'source',
+      Target = 'target',
+      Value = 'value',
+      NodeID = 'name',
+      ...
+    )
   }
+)
 
-  # create nodes table
-  nodes = data.table::data.table(name = nodes)
-
-  networkD3::sankeyNetwork(
-    Links = links_dt,
-    Nodes = nodes,
-    Source = 'source',
-    Target = 'target',
-    Value = 'value',
-    NodeID = 'name',
-    ...
-  )
-}
+
+
+#' @rdname sankeyPlot
+#' @export
+setMethod(
+  'sankeyPlot',
+  signature(gobject = 'giotto',
+            x = 'character'),
+  function(gobject,
+           x,
+           spat_unit = NULL,
+           feat_type = NULL,
+           meta_type = c('cell', 'feat'),
+           idx = NULL,
+           ...) {
+
+    checkmate::assert_character(x, len = 2L)
+
+    # Data type being compared. Either cell or feat
+    meta_type = match.arg(meta_type, choices = c('cell', 'feat'))
+    # Determines which type of metadata to get
+    meta_get_fun = switch(
+      meta_type,
+      'cell' = getCellMetadata,
+      'feat' = getFeatureMetadata
+    )
+    id_col = switch(
+      meta_type,
+      'cell' = 'cell_ID',
+      'feat' = 'feat_ID'
+    )
+
+    # get metadata
+    # Defaults for spat_unit and feat_type are set inside of the getter if they
+    # are provided as NULL
+    meta_cm = meta_get_fun(
+      gobject = g,
+      spat_unit = spat_unit,
+      feat_type = feat_type,
+      output = 'cellMetaObj',
+      copy_obj = TRUE,
+      set_defaults = TRUE
+    )
+
+    # perform subset
+    # 1. subobject subsetting does not drop ID col
+    # 2. drop to data.table
+    if (!is.null(idx)) {
+      meta_cm = meta_cm[idx]
+    }
+    test_dt = meta_cm[][, x, with = FALSE]
+
+    res = sankey_compare(data_dt = test_dt)
+    links_dt = res$links
+
+    # create nodes table
+    nodes = data.table::data.table(name = res$nodes)
+
+    networkD3::sankeyNetwork(
+      Links = links_dt,
+      Nodes = nodes,
+      Source = 'source',
+      Target = 'target',
+      Value = 'value',
+      NodeID = 'name',
+      ...
+    )
+  }
+)
 
diff --git a/man/sankeyPlot.Rd b/man/sankeyPlot.Rd
index 8de9a79..92e1ef9 100644
--- a/man/sankeyPlot.Rd
+++ b/man/sankeyPlot.Rd
@@ -2,16 +2,28 @@
 % Please edit documentation in R/plot_sankey.R
 \name{sankeyPlot}
 \alias{sankeyPlot}
+\alias{sankeyPlot,giotto,character-method}
 \title{Create a sankey plot}
 \usage{
-sankeyPlot(gobject, sankey_plan, meta_type = c("cell", "feat"), ...)
+\S4method{sankeyPlot}{giotto,giottoSankeyPlan}(gobject, x, meta_type = c("cell", "feat"), ...)
+
+\S4method{sankeyPlot}{giotto,character}(
+  gobject,
+  x,
+  spat_unit = NULL,
+  feat_type = NULL,
+  meta_type = c("cell", "feat"),
+  idx = NULL,
+  ...
+)
 }
 \arguments{
 \item{gobject}{giotto object}
 
-\item{sankey_plan}{giottoSankeyPlan object}
+\item{x}{giottoSankeyPlan object or character vector referring to source and
+target columns in metadata}
 
-\item{meta_type}{build sankey on cell or feature metadata}
+\item{meta_type}{whether to use cell or feature metadata}
 
 \item{...}{
   Arguments passed on to \code{\link[networkD3:sankeyNetwork]{networkD3::sankeyNetwork}}
@@ -42,14 +54,25 @@ browser on the client so don't push it too high.}
     \item{\code{sinksRight}}{boolean. If \code{TRUE}, the last nodes are moved to the
 right border of the plot.}
   }}
+
+\item{spat_unit}{spatial unit of metadata}
+
+\item{feat_type}{feature type of metadata}
+
+\item{idx}{table subset index for 1 to 1 comparisons}
 }
 \description{
 Create a sankey plot from a giotto object. Pulls from information in the
-metadata.
+metadata. Simple 1 to 1 sankeys can be generated from a single spatial unit
+and feature type using the \code{spat_unit}, \code{feat_type}, \code{meta_type}, \code{cols},
+and (optionally) \code{idx} params. More complex and cross spatial unit/feature
+type sankeys can be set up using the \code{sankey_plan} param which accepts a
+\code{giottoSankeyPlan} object.
 }
 \examples{
 \dontrun{
 g = GiottoData::loadGiottoMini("vizgen")
+# with giottoSankeyPlan
 leiden = sankeySet(spat_unit = 'aggregate',
                    feat_type = 'rna',
                    col = 'leiden_clus')
@@ -59,6 +82,10 @@ louvain = sankeySet(spat_unit = 'aggregate',
 plan = leiden + louvain
 sankeyRelate(plan) = c(0,1)
 sankeyPlot(g, plan)
+
+# with single set of metadata
+activeSpatUnit(g) = 'aggregate'
+sankeyPlot(g, c('louvain_clus', 'leiden_clus'))
 }
 }
 \keyword{plotting}