R/forward_warp.R

In jdknguyen/SMART: SMART (Semi-manual alignment to reference template): a pipeline for whole brain mapping projects

#' @title Forward warp back to atlas space
#' @description Loop through all segmented cells and perform forward warp loop
#'   back onto atlas space using registration and segmentation data. Options for
#'   plotting a schematic plot and saving forward warp images. For a whole
#'   brain, this function should be used after the segmentation list *segs* has
#'   been cleaned of duplicate cell counts. Note: if the imaging dataset is
#'   large this will be a time intensive process. Processing time will be
#'   printed once the function is finished.
#' @param setup (required) Setup list from [setup_pl()].
#' @param segs (required) Segmentation data.
#' @param regis (required) Registration data.
#' @param filetype (optional, default = "tif") Image type to save as. See the
#'   *type* argument in the function [savePlot()]. Options: "tif", "tiff",
#'   "wmf", "emf", "png", "jpg", "jpeg", "bmp","ps", "eps", "pdf"
#' @param savewarp (optional, default = TRUE) Save forward warp plot image in
#'   folder 'setup$savepaths$out_segmentation_warps'.
#' @param saveschematic (optional, default = TRUE) Save schematic plot images in
#'   folder 'setup$savepaths$out_segmentation_schem'.
#' @param plane (optional, default = "coronal") Atlas plane to register to.
#'   options: "coronal", "sagittal"
#' @param title (optional, default = FALSE) Title for the schematic plot.
#' @param mm.grid (optional, default = TRUE) Plot grid in schematic plot.
#' @param dev.size (optional, default = c(5.4, 4.465)). Dimensions of window.
#' @param pch (optional, default = 21) Graphical parameter for point shape.
#' @param cex (optional, default = 0.5) Graphical parameter for text size.
#' @param col (optional, default = "black") Graphical parameter for color.
#' @param scale.bar (optional, default = TRUE)  Display a measure bar in the
#'   schematic plot
#' @param region.colors (optional, default = TRUE)  Display Allen Brain Atlas
#'   colors in schematic plot.
#' @return Returns *dataset* a variable storing all mapped segmentation data to
#'   registration data.
#' @export
#' @md

forward_warp <- function (setup, segs, regis, filetype = c("tif", "tiff",
                                                           "wmf", "emf", "png", "jpg", "jpeg",
                                                           "bmp", "ps", "eps", "pdf"), savewarp = TRUE,
                          saveschematic = TRUE, plane = "coronal", title = FALSE,
                          mm.grid = TRUE, dev.size = c(5.4, 4.465), pch = 21, cex = 0.5,
                          col = "black", scale.bar = FALSE, region.colors = TRUE)
{
  tictoc::tic()
  filetype <- match.arg(filetype)
  if (length(setup) > 9) {
    regi_ind <- vector(mode = "numeric", length(segs$seg_z))
    for (n in 1:length(segs$seg_z)) {
      regi_ind[n] <- which.min(abs(setup$regi_z - segs$seg_z[n]))
    }
  }
  else {
    regi_ind <- 1:length(segs$seg_z)
  }
  if (length(setup$regi_z) > 1) {
    my_function <- approxfun(setup$regi_z, setup$regi_AP,
                             method = "linear")
    z_matched_AP <- my_function(segs$seg_z)
  }
  else {
    z_matched_AP <- setup$regi_AP
  }
  for (n in 1:length(segs$seg_z)) {
    toggle <- TRUE
    if(length(segs$segmentations[[n]]$soma$x) > 0){
      for(j in 1:length(segs$segmentations[[n]]$soma$x)){
        if(is.na(segs$segmentations[[n]]$soma$x[j]) == FALSE){
          toggle <- FALSE
        }
      }
    }
    if(toggle == FALSE){
      regi <- regis[[regi_ind[n]]]
      seg <- segs$segmentations[[n]]
      data <- wholebrain::inspect.registration(regi, seg, soma = TRUE,
                                               forward.warps = TRUE, batch.mode = TRUE)
      if (savewarp) {
        path <- paste0(setup$savepaths$out_segmentation_warps,
                       "/forwardwarp_z_", toString(segs$seg_z[n]),
                       "_AP_", toString(round(regi$coordinate,
                                              digits = 2)), "_plate_", platereturn(round(regi$coordinate,
                                                                                         digits = 2)), ".", filetype)
        savePlot(filename = path, type = filetype)
      }
      dev.off()
      data <- data[!data$id == 0, ]
      data$animal <- setup$anim_ID
      if (saveschematic) {
        wholebrain::schematic.plot(dataset = data, plane = plane,
                                   save.plots = FALSE, title = title, mm.grid = mm.grid,
                                   dev.size = dev.size, pch = pch, cex = cex, col = "black",
                                   scale.bar = scale.bar, region.colors = region.colors)
        path <- paste0(setup$savepaths$out_segmentation_schem,
                       "/schematic_z_", toString(segs$seg_z[n]),
                       "_AP_", toString(round(regi$coordinate,
                                              digits = 2)), "_plate_", platereturn(round(regi$coordinate,
                                                                                         digits = 2)), ".", filetype)
        savePlot(filename = path, type = filetype)
        dev.off()
      }
      data$AP <- z_matched_AP[n]
      if (n == 1) {
        dataset <- data
      }
      else {
        dataset <- rbind(dataset, data)
      }
    }
  }
  time <- tictoc::toc(quiet = TRUE)
  cat("\n", toString(round((time$toc - time$tic)/60,
                           digits = 2)), "min elapsed")
  return(dataset)
}