🚧 Update and export graph-related functions

- Cache edge_ages() for repeated sampling
- Why distance("unweighted") runs slower than "dijkstra"?

NAMESPACE

@@ -15,7 +15,10 @@ export(count_extant)
 export(dist_euclidean)
 export(dist_genealogy)
 export(dist_vaf)
+export(distances_from_origin)
 export(distances_mrs)
+export(distances_upstream)
+export(edge_lengths)
 export(evaluate_mrs)
 export(extract_history)
 export(filter_common_ancestors)
@@ -51,9 +54,11 @@ export(sample_bulk)
 export(sample_random_regions)
 export(sample_uniform_regions)
 export(save_serial_section)
+export(sink_cells)
 export(snapshot_surface)
 export(sort_vaf)
 export(stats_dispersion)
+export(subgraph_upstream)
 export(subtree)
 export(summarize_capture_rate)
 export(summarize_demography)

R/graph.R

@@ -19,14 +19,21 @@ as_symbolic_edgelist = function(population) {
 }
 
 #' @details
-#' `subtree` extracts subgraph among terminal nodes.
+#' `subtree()` and `subgraph_upstream()` extracts subgraph among terminal nodes.
 #' @param graph igraph
 #' @param nodes integer cell IDs
+#' @param vids vertex IDs
 #' @param trim whether to remove common ancestors older than MRCA.
 #' @rdname graph
 #' @export
 subtree = function(graph, nodes = integer(0L), trim = FALSE) {
   vids = igraphlite::as_vids(graph, nodes)
+  subgraph_upstream(graph, vids, trim = trim)
+}
+
+#' @rdname graph
+#' @export
+subgraph_upstream = function(graph, vids = integer(0L), trim = FALSE) {
   vids = upstream_vertices(graph, vids, trim = trim)
   igraphlite::induced_subgraph(graph, vids)
 }
@@ -59,6 +66,29 @@ internal_nodes = function(graph, nodes, sensitivity) {
   as.integer(names(counts)[(counts / n) > sensitivity])
 }
 
+#' @details
+#' `distances_from_origin()` and `distances_upstream()` are a shorthand of
+#' `distances(..., mode = 2L)`.
+#' @param weights A numeric or logical. Edge attribute "weight" is used if TRUE.
+#' @rdname graph
+#' @export
+distances_from_origin = function(graph, nodes = sink_cells(graph), weights = numeric(0L), trim = FALSE) {
+  vids = igraphlite::as_vids(graph, nodes)
+  distances_upstream(graph, vids, weights = weights, trim = trim)
+}
+
+#' @rdname graph
+#' @export
+distances_upstream = function(graph, vids = numeric(0), weights = numeric(0L), trim = FALSE) {
+  src = if (trim) {
+    common_ancestors(graph, vids)[1L]
+  } else {
+    igraphlite::Vsource(graph)
+  }
+  res = igraphlite::distances(graph, vids, src, weights = weights, mode = 2L)
+  as.vector(res)
+}
+
 upstream_vertices = function(graph, vids, trim = FALSE) {
   vlist = neighborhood_in(graph, vids)
   vids = unique(unlist(vlist, use.names = FALSE))
@@ -69,6 +99,17 @@ upstream_vertices = function(graph, vids, trim = FALSE) {
   vids
 }
 
+get_mrca = function(graph, nodes) {
+  ca = common_ancestors(graph, igraphlite::as_vids(graph, nodes))
+  igraphlite::as_vnames(graph, ca[1L])
+}
+
+common_ancestors = function(graph, vids) {
+  vlist = neighborhood_in(graph, vids)
+  vids = unique(unlist(vlist, use.names = FALSE))
+  Reduce(intersect, vlist)
+}
+
 neighborhood_out = function(graph, vids) {
   igraphlite::neighborhood(graph, vids, order = 1073741824L, mode = 1L)
 }
@@ -91,22 +132,16 @@ paths_to_source = function(graph, nodes = integer(0L)) {
   lapply(res, function(x) vnames[x])
 }
 
-shortest_dist_from_source = function(graph, vids = numeric(0)) {
-  res = igraphlite::distances(graph, vids, igraphlite::Vsource(graph), mode = 2L, algorithm = "unweighted")
-  as.integer(res)
-}
-
-distances_from_origin = function(graph, nodes = integer(0L)) {
-  vids = igraphlite::as_vids(graph, nodes)
-  shortest_dist_from_source(graph, vids)
-}
-
-sink_nodes = function(graph) {
+#' @details
+#' `sink_cells()` returns the cell IDs of the sink nodes.
+#' @rdname graph
+#' @export
+sink_cells = function(graph) {
   igraphlite::Vnames(graph)[igraphlite::is_sink(graph)]
 }
 
 sinks = function(graph, nodes) {
-  .sink = sink_nodes(graph)
+  .sink = sink_cells(graph)
   paths = paths_to_sink(graph, nodes)
   lapply(paths, function(x) x[x %in% .sink])
 }
@@ -133,3 +168,15 @@ count_sink = function(graph, nodes = integer(0L)) {
     sum(igraphlite::is_sink(graph))
   }
 }
+
+copy_edge_attr = function(subgraph, graph) {
+  if (ncol(graph$Eattr) > 0L) {
+    eattr_names = names(igraphlite::edge_attr(graph))
+    lhs = as.data.frame(subgraph) |>
+      dplyr::select(!dplyr::any_of(eattr_names))
+    rhs = as.data.frame(graph)
+    subgraph$Eattr = dplyr::left_join(lhs, rhs, by = c("from", "to")) |>
+      dplyr::select(!c("from", "to"))
+  }
+  invisible(subgraph)
+}

R/ms.R

@@ -38,18 +38,33 @@ mutate_clades = function(graph, mu = NULL, segsites = NULL) {
     dplyr::mutate(carriers = paths_to_sink(graph, .data$origin))
 }
 
-genetic_distance = function(graph, vids = igraphlite::Vsink(graph), mu = 0, accel = 0) {
-  segments = if (accel > 0) {
-    age = distances_from_origin(graph)
-    (1 + accel)**age
-  } else {
-    rep_len(1.0, graph$vcount)
+#' @details
+#' `edge_lengths()` calculates lengths of edges on a given genealogy tree.
+#' @param accel assumption undocumented yet
+#' @rdname ms
+#' @export
+edge_lengths = function(graph, mu = 0, accel = 0) {
+  # distances(): "unweighted" is slower than "dijkstra"
+  lens = rep_len(1, graph$ecount)
+  if (accel > 0) {
+    lens = (1 + accel)**edge_ages(graph)
   }
   if (mu > 0) {
-    segments = stats::rpois(length(segments), mu * segments)
+    lens = stats::rpois(graph$ecount, mu * lens)
   }
-  segments[igraphlite::is_source(graph)] = 0
-  ancestors = neighborhood_in(graph, vids)
-  names(ancestors) = igraphlite::Vnames(graph)[vids]
-  purrr::map_dbl(ancestors, \(v) sum(segments[v]))
+  lens
+}
+
+edge_ages = function(graph) {
+  age = igraphlite::edge_attr(graph, "age")
+  if (is.null(age)) {
+    age = distances_upstream(graph, graph$to, trim = FALSE)
+    igraphlite::edge_attr(graph, "age") = age
+  }
+  age
+}
+
+mutate_edges = function(graph, mu = 0, accel = 0) {
+  igraphlite::edge_attr(graph, "weight") = edge_lengths(graph, mu = mu, accel = accel)
+  graph
 }

R/plot-genealogy.R

@@ -22,7 +22,9 @@ augment_genealogy = function(graph, mu = 0, accel = 0) {
       TRUE ~ "internal"
     ))
   if (mu > 0 || accel > 0) {
-    named_dist = genetic_distance(graph, igraphlite::V(graph), mu = mu, accel = accel)
+    elengths = edge_lengths(graph, mu = mu, accel = accel)
+    udist = distances_upstream(graph, weights = elengths)
+    named_dist = stats::setNames(udist, vnames)
     res$d = named_dist[as.character(res$node)]
     res$d_parent = named_dist[as.character(res$parent)]
   }