R/orient.R

In dynverse/scvelo: An R wrapper around the great "scvelo" python package

#' Reorients the edges of the milestone network to the cell's RNA velocity vectors
#'
#' @inheritParams dynwrap::common_param
#' @inheritParams dynwrap::add_expression
#' @param velocity The velocity object as generated by [get_velocity()].
#'
#' @return The trajectory with oriented *milestone_network* and *progressions*
#'
#' @importFrom dynwrap get_expression flip_edges remove_root
#' @importFrom Matrix tril triu t
#' @importFrom methods is
#'
#' @export
# removed example for now
# @examples
# # we create a simple linear trajectory
# library(tibble)
# library(dynwrap)
#
# cell_ids <- c("a", "b", "c", "d", "e")
# pseudotime <- setNames(seq_along(cell_ids), cell_ids)
# expression <- as.matrix(data.frame(
#   a = pseudotime,
#   b = pseudotime ** 2,
#   c = log(pseudotime)
# ))
# expression_future <- as.matrix(data.frame(
#   a = (pseudotime + 1),
#   b = (pseudotime + 1) ** 2,
#   c = log(pseudotime + 1)
# ))
#
# # the milestone network is "wrong" in the sense that B and A are oriented in the opposite direction
# milestone_network <- tibble::tribble(
#   ~from, ~to, ~length, ~directed,
#   "B", "A", 1, TRUE,
#   "B", "C", 1, TRUE
# )
# progressions <- tibble::tribble(
#   ~cell_id, ~from, ~to, ~percentage,
#   "a", "B", "A", 1,
#   "b", "B", "A", 0.5,
#   "c", "B", "A", 0,
#   "d", "B", "C", 0.5,
#   "e", "B", "C", 1
# )
#
# trajectory <- dynwrap::wrap_expression(
#   counts = expression,
#   expression = expression,
#   expression_future = expression_future
# )
#
# trajectory <- dynwrap::add_trajectory(
#   trajectory,
#   milestone_network = milestone_network,
#   progressions = progressions
# )
#
# trajectory_oriented <- orient_topology_to_velocity(trajectory)
#
# # the edge is now correctly oriented
# trajectory_oriented$milestone_network
orient_topology_to_velocity <- function(
  trajectory,
  velocity = trajectory$velocity,
  expression_source = trajectory,
  expression = dynwrap::get_expression(expression_source, "expression"),
  expression_future = dynwrap::get_expression(expression_source, "expression_future")
) {
  # dummy proofing
  assert_that(
    is(trajectory, "dynwrap::with_trajectory"),
    !is.null(expression),
    !is.null(expression_future)
  )

  if (nrow(trajectory$divergence_regions)) {
    stop("Orienting topologies with divergence regions doesn't work yet")
  }

  assert_that(
    !is.null(velocity$transition_matrix),
    nrow(velocity$transition_matrix) == nrow(expression),
    nrow(velocity$transition_matrix) == nrow(expression)
  )

  transition_matrix <- velocity$transition_matrix

  flip_fractions <- pmap_dbl(trajectory$milestone_network, function(from, to, ...) {
    progressions_edge <-
      trajectory$progressions %>%
      filter(from == !!from, to == !!to) %>%
      arrange(percentage)

    transition_matrix_ordered <- Matrix::t(transition_matrix[progressions_edge$cell_id, progressions_edge$cell_id])

    sum(Matrix::tril(transition_matrix_ordered, 1)) / sum(Matrix::triu(transition_matrix_ordered, 1))
  })

  trajectory <- dynwrap::flip_edges(
    trajectory,
    trajectory$milestone_network %>%
      mutate(fraction = flip_fractions) %>%
      filter(fraction < 1)
  )

  # remove the root
  trajectory <- trajectory %>% remove_root()

  # tada!
  trajectory
}