helper-funky.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_multidim <- 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)

  dimension_names <- stringr::str_extract(colnames(cell_positions), "comp_[0-9]*") %>% discard(is.na)

  positions <- cell_positions %>%
    select(cell_id, !!dimension_names) %>%
    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"))
  waypoint_edges$length <- sqrt(rowSums((waypoint_edges[,paste0(dimension_names, "_from")] - waypoint_edges[,paste0(dimension_names, "_to")])**2))

  lst(
    positions = waypoint_positions,
    edges = waypoint_edges
  )
}