R/position-brain.R

In LCBC-UiO/ggbrain: Plotting Tool for Brain Atlases

# position ----

#' Reposition brain slices
#'
#' Function for repositioning
#' pre-joined atlas data (i.e. data and atlas
#' already joined to a single data frame).
#' This makes it possible for users to
#' reposition the geometry data for the atlas
#' for control over final plot layout. For even
#' more detailed control over the positioning,
#' the "hemi" and "side" columns should be
#' converted into factors and ordered by wanted
#' order of appearance.
#'
#' @param data sf-data.frame of joined brain atlas and data
#' @param position position formula for slices
#'
#' @return sf-data.frame with repositioned slices
#' @export
#'
#' @examples
#' reposition_brain(dk, hemi ~ side)
#' reposition_brain(dk, side ~ hemi)
#' reposition_brain(dk, hemi + side ~ .)
#' reposition_brain(dk, . ~ hemi + side)
reposition_brain <- function(data, position = "horizontal"){
  data <- as.data.frame(data, stringsAsFactors = FALSE)
  frame_2_position(data, position)
}


#' Alter brain atlas position
#'
#' Function to be used in the position argument in geom_brain
#' to alter the position of the brain slice/views.
#'
#' @param position formula describing the rows ~ columns organisation.
#'
#' @export
#' @return a ggproto object
#' @importFrom ggplot2 ggproto
#' @examples
#' library(ggplot2)
#' ggplot() +
#'   geom_brain(atlas = dk, aes(fill = region),
#'              position = position_brain(. ~ side + hemi ),
#'              show.legend = FALSE)
#'
#' ggplot() +
#'   geom_brain(atlas = dk, aes(fill = region),
#'              position = position_brain(side ~ hemi ),
#'              show.legend = FALSE)
position_brain <- function(position = "horizontal") {
  ggproto(NULL, PositionBrain, position = position)
}

PositionBrain <- ggplot2::ggproto(
  "PositionBrain", ggplot2:::Position,
  position = hemi + side ~ .,

  setup_params = function(self, data) {
    list(position = self$position)
  },

  compute_layer = function(self, data, params, layout) {

    df3 <- frame_2_position(data, params$position)
    bbx <- sf::st_bbox(df3$geometry)

    # rescale layout to reflect new coordinates
    if(is.null(layout$coord$limits$y))
      layout$coord$limits$y <- bbx[c(2,4)]

    # rescale layout to reflect new coordinates
    if(is.null(layout$coord$limits$x))
      layout$coord$limits$x <- bbx[c(1,3)]

    data <- df3

    df3
  }
)

# geometry movers ----

position_formula <- function(pos, data){
  chosen <- all.vars(pos, unique = FALSE)
  chosen <- chosen[!grepl("\\.", chosen)]

  if(any(duplicated(chosen)))
    stop("Cannot position brain with the same data as columns and rows",
         call. = FALSE)

  if(unique(data$type) == "cortical"){
    if(length(chosen) < 2)
      stop("position formula not correct. ",
           "Missing ", paste0(c("side","hemi")[!c("side","hemi") %in% chosen], collapse=" & "), "",
           call. = FALSE
      )
    position <- if(length(grep("\\+", pos)) > 0){
      ifelse(grep("^\\.", pos) == 2,
             "columns", "rows")
    }else{
      chosen
    }
  }

  if(all(sum(grepl("\\.|~", pos)) != 2 & position %in% c("rows", "columns")))
    stop("Formula for a single row or column must contain both a . and ~",
         call. = FALSE)

  list(
    position = position,
    chosen = chosen
  )

}


frame_2_position <- function(data, pos){
  dfpos <- split_data(data, pos)
  df2 <- lapply(dfpos$data, gather_geometry)
  posi <- ifelse(length(dfpos$position) > 1, "grid", dfpos$position)
  df3 <- switch(posi,
                rows    = stack_vertical(df2),
                columns = stack_horizontal(df2),
                grid    = stack_grid(df2, dfpos$position[1], dfpos$position[2])
  )

  df4 <- st_as_sf(df3$df)
  attr(sf::st_geometry(df4), "bbox") = df3$box

  df4
}

split_data <- function(data, position){
  if(inherits(position, "formula")){
    pos <- position_formula(position, data)
    df2 <- dplyr::group_by_at(data, pos$chosen)
    df2 <- dplyr::group_split(df2)
    pos <- pos$position
  }else{
    if(length(position) == 1){
      if(position %in% c("horizontal", "vertical"))
        position <- default_order(data)
    }
    pos <- as.data.frame(strsplit(position, " "),stringsAsFactors = FALSE)
    if(unique(data$type) == "cortical"){
      k <- cbind(pos[2,] %in% data$side,
                 pos[1,] %in% data$hemi)
      k <- sapply(1:nrow(k), function(x) sum(k[x,]))
      pos <- pos[ifelse(k == 2, TRUE, FALSE)]

      df2 <- lapply(pos, function(x){
        data[data$hemi == x[1] & data$side == x[2],]
      })
    }else{
      df2 <- lapply(pos, function(x){
        data[data$side == x,]
      })
    }
    pos <- unique(ifelse(position == "vertical", "rows", "columns"))
  }

  return(list(data = df2, position = pos))
}

gather_geometry <- function(df){
  bbx <- sf::st_bbox(df$geometry)
  df$geometry <- df$geometry - bbx[c("xmin", "ymin")]
  df
}

stack_horizontal <- function(df){
  sep <- get_sep(df)

  bx <- list()
  for(k in 1:length(df)){
    df[[k]]$geometry <- df[[k]]$geometry + c((k-1)*sep[1], 0)
    bx[[k]] <- sf::st_bbox(df[[k]]$geometry )
  }

  list(
    df = do.call(rbind, df),
    box = get_box(bx)
  )
}

stack_vertical <- function(df){
  sep <- get_sep(df)

  bx <- list()
  for(k in 1:length(df)){
    df[[k]]$geometry <- df[[k]]$geometry + c(0, (k-1)*sep[2])
    bx[[k]] <- sf::st_bbox(df[[k]]$geometry )
  }

  list(
    df = do.call(rbind, df),
    box = get_box(bx)
  )
}

stack_grid <- function(df, rows, columns){
  bx <- list()
  sep <- get_sep(df)

  if(length(df) == 4){
    # switch columns if they are not the same
    # so grid aligns properly
    if(unique(df[[3]][columns]) != unique(df[[2]][columns])){
      df <- list(df[[1]], df[[2]],
                 df[[4]], df[[3]])
    }
  }

  # move second and fourth on x
  for(k in c(2,3)){
    df[[k]]$geometry <- df[[k]]$geometry + c(sep[1],0)
  }

  # move third and fourth on y
  for(k in c(3,4)){
    df[[k]]$geometry <- df[[k]]$geometry + c(0,sep[2])
  }

  bx <- lapply(df, function(x) sf::st_bbox(x$geometry ))
  df <- do.call(rbind, df)

  df[, c("xmin", "xmax", "ymin", "ymax")] <- NULL

  list(
    df = df,
    box = get_box(bx)
  )
}

get_box <- function(bx){
  bx <- do.call(rbind, bx)
  pad <- max(bx)*.01
  bx <- c(-pad, -pad,
          max(bx[,"xmax"]) + pad,
          max(bx[,"ymax"]) + pad)
  x <- stats::setNames(bx, c("xmin", "ymin", "xmax", "ymax"))
  class(x) <- "bbox"
  x
}

get_sep <- function(data){
  sep <- sapply(data, function(x) sf::st_bbox(x$geometry))
  sep <- c(max(sep[3,]), max(sep[4,]))
  c("x" = sep[1] + sep[1]*.2, "y" = sep[2] + sep[2]*.2)
}

default_order <- function(data){
  if(unique(data$type) == "cortical")
    return(c("left lateral", "left medial", "right medial", "right lateral"))
  unique(data$side)
}