R/vendor/dynwrap.R

In aertslab/SCopeLoomR: Build .loom files (compatible with SCope) and extract data from .loom files.

# Source: https://github.com/dynverse/dynplot/blob/master/R/plot_dimred.R
#' Project the waypoints
#' @inheritParams add_cell_coloring
#' @param waypoints The waypoints to use for projecting, as generated by [dynwrap::select_waypoints()]
#' @param trajectory_projection_sd The standard deviation of the gaussian kernel
#' @param color_trajectory How to color the trajectory, can be "nearest" for coloring to nearest cell, or "none"
#' 
#' @importFrom stats dnorm
project_waypoints <- function(
  traj,
  cell_positions,
  waypoints = dynwrap::select_waypoints(traj),
  trajectory_projection_sd = sum(traj$milestone_network$length) * 0.05,
  color_trajectory = "none"
) {
  testthat::expect_setequal(cell_positions$cell_id, colnames(waypoints$geodesic_distances))
  
  # project waypoints to dimensionality reduction using kernel and geodesic distances
  # rate <- 5
  # trajectory_projection_sd <- sum(traj$milestone_network$length) * 0.05
  # dist_cutoff <- sum(milestone_network$length) * 0.05
  # k <- 3
  # weight_cutoff <- 0.01
  
  # weights <- waypoints$geodesic_distances %>% stats::dexp(rate = 5)
  weights <- waypoints$geodesic_distances %>% stats::dnorm(sd = trajectory_projection_sd)
  testthat::expect_true(all(!is.na(weights)))
  # weights <- waypoints$geodesic_distances < dist_cutoff
  # weights[weights < weight_cutoff] <- 0
  
  weights <- weights / rowSums(weights)
  positions <- cell_positions %>%
    select(cell_id, comp_1, comp_2) %>%
    slice(match(colnames(weights), cell_id)) %>%
    column_to_rownames("cell_id") %>%
    as.matrix()
  
  # make sure weights and positions have the same cell_ids in the same order
  testthat::expect_equal(colnames(weights), rownames(positions))
  
  # calculate positions
  waypoint_positions <- (weights %*% positions) %>%
    as.data.frame() %>%
    rownames_to_column("waypoint_id") %>%
    left_join(waypoints$waypoints, "waypoint_id")
  
  # add color of closest cell
  if (color_trajectory == "nearest") {
    testthat::expect_true("color" %in% colnames(cell_positions))
    
    waypoint_positions <- waypoint_positions %>%
      mutate(closest_cell_ix = (weights %>% apply(1, which.max))[waypoint_id]) %>%
      mutate(closest_cell_id = colnames(weights)[closest_cell_ix]) %>%
      mutate(color = (cell_positions %>% select(cell_id, color) %>% deframe())[closest_cell_id])
  }
  
  # positions of different edges
  waypoint_edges <- waypoints$waypoint_network %>%
    left_join(waypoint_positions %>% rename_all(~paste0(., "_from")), c("from" = "waypoint_id_from")) %>%
    left_join(waypoint_positions %>% rename_all(~paste0(., "_to")), c("to" = "waypoint_id_to")) %>%
    mutate(length = sqrt((comp_1_to - comp_1_from)**2 + (comp_2_to - comp_2_from)**2))
  
  # add arrows to every milestone to milestone edge
  # an arrow is placed at the waypoint which is in the middle from the start and end
  waypoint_edges <- waypoint_edges %>%
    group_by(from_milestone_id, to_milestone_id) %>%
    mutate(
      distance_to_center = (comp_1_from - mean(c(max(comp_1_from), min(comp_1_from))))^2 + (comp_2_from - mean(c(max(comp_2_from), min(comp_2_from))))^2,
      arrow = row_number() == which.min(distance_to_center)
    )
  
  lst(
    positions = waypoint_positions,
    edges = waypoint_edges
  )
}