R/ids.R

In natverse/fafbseg: Support Functions for Analysis of FAFB EM Segmentation

#' Convert between filenames and neuroglancer ids
#'
#' @description \code{swc2segmentid} converts an swc filename to a segment id
#'
#' @param x Input file or id
#' @param include.fragment Whether to include the sub identifier of the skeleton
#'   fragment (see details).
#' @return for \code{swc2segmentid} a numeric vector or matrix depending on the
#'   value of \code{include.fragment}
#' @export
#' @name fafbseg-ids
#' @importFrom stringr str_match
#' @examples
#' swc2segmentid("10001654273.1.swc")
#' swc2segmentid(sprintf("10001654273.%d.swc", 0:2), include.fragment=TRUE)
swc2segmentid <- function(x, include.fragment=FALSE) {
  res=str_match(basename(x), "^(\\d+)(\\.(\\d+)){0,1}\\.[Ss][Ww][Cc]$")
  if(isTRUE(include.fragment)) {
    res=res[,c(2,4), drop=FALSE]
    colnames(res)=c("segment", "fragment")
    res
  } else {
    res=res[,2]
  }
  mode(res)='numeric'
  res
}

#' @description \code{segmentid2zip} converts a segment id to the zip file that
#'   contains it
#' @export
#' @rdname fafbseg-ids
#' @examples
#' \dontrun{
#' # NB the default segmentation for fafbseg (flywire) no longer implies a local
#' # collection of skeletons, wrap calls in with_segmentation
#' with_segmentation("20190805", segmentid2zip(10001654273))
#' with_segmentation("20190805",
#'   segmentid2zip(swc2segmentid("10001654273.1.swc")))
#' }
segmentid2zip <- function(x) {
  divisor <- find_zip_divisor(getOption("fafbseg.skelziproot"))
  sprintf("%d.zip", as.numeric(x) %/% divisor)
}

#' @description \code{zip2segmentstem} converts a zip file to the initial part
#'   of the segment id i.e. the segment stem (see details).
#'
#' @details Segment ids are unique integers. There are about 8E8 in the current
#'   skeletonisation but it seems that the ids can still be > 2^31 (usually
#'   \code{.Machine$integer.max}). Therefore they will be stored in R as numeric
#'   values or the \code{bit64::integer64} values.
#'
#'   Each segmentation has keen skeletonised however this usually results in
#'   multiple skeleton fragments which have been written out as separate SWC
#'   files: \code{"named <segment id>.<fragment>.swc"}
#'
#'   Each segment id is mapped onto a zip file by dividing by a divisor and
#'   discarding the remainder. Peter Li's data release of 2018-10-02 switched
#'   from 1E5 to 1E6.
#' @export
#' @rdname fafbseg-ids
#' @importFrom tools file_path_sans_ext
zip2segmentstem <- function(x) {
  as.integer(file_path_sans_ext(basename(x)))
}


# ids should be integers >= 0
# this does not check that they are also valid 64 bit ints which might be good
valid_id <- function(x, na.ok=FALSE) {
  if(is.integer64(x) || is.integer(x)) {
    if(na.ok)
      return(x>=0 | is.na(x))
    else
      return(!is.na(x) & x>=0)
  }
  if(is.numeric(x)) {
    inrange=(x>=0 & x < 2^53)
    inrange[is.na(inrange)]=FALSE
    if(na.ok)
      return(inrange | is.na(x))
    else
      return(inrange)

    return(checkmate::test_double(x, lower=0, upper=(2^53-1), any.missing = na.ok))
  }
  cx=as.character(x)
  res=grepl("^\\d{1,19}$", cx, perl = TRUE, useBytes = TRUE)
  if(na.ok) {
    nares=grepl("^(|NA|NAN|NULL)$", cx, perl = TRUE, ignore.case = TRUE)
    res[is.na(cx) | nares]=TRUE
  }
  res
}


#' Helper function to turn diverse inputs into neuroglancer segment ids
#'
#' @param x Neuroglancer ids either as a vector of ids (character vector or
#'   \code{integer64} recommended) OR a scene specification as raw JSON. format
#'   (character vector), the path to a file on disk, a neuroglancer scene URL
#'   (which embeds a JSON scene specification in a single URL), or an R list
#'   generated by parsing one of the above.
#' @param as_character Whether to return segments as character rather than
#'   numeric vector (the default is character for safety).
#' @param include_hidden Whether to include \code{hiddenSegments} (typically for
#'   flywire scenes).
#' @param must_work if \code{TRUE}, the default, then an error will be generated
#'   if the scene has no segments or if invalid ids are present. Explicit NAs
#'   will trigger this error when \code{must_work=T}, but will be returned as id
#'   0 without error when \code{must_work=F}.
#' @param unique When \code{TRUE} drops any duplicated ids with a warning
#' @param ... Additional arguments passed to \code{\link{ngl_decode_scene}}
#'
#' @return Numeric (or character) vector of segment ids, taken from the first
#'   segmentation layer (with a warning) if the scene contains more than one.
#' @details For simple vector inputs, \code{ngl_segments} will treat the values
#'   \code{"NA", "NAN", "NULL"} (all case insensitive) and \code{""} as NA.
#'
#' @export
#' @family neuroglancer-urls
#' @examples
#' # -> character
#' ngl_segments(c(10950626347, 10952282491, 13307888342))
#' # turns these into numeric
#' ngl_segments(c("10950626347", "10952282491", "13307888342"), as_character=FALSE)
#'
#' \donttest{
#' u="https://ngl.flywire.ai/?json_url=https://globalv1.flywire-daf.com/nglstate/5409525645443072"
#' ngl_segments(u, as_character = TRUE)
#' sc=ngl_decode_scene(u)
#' # set segments
#' ngl_segments(sc) <- c("720575940621039145")
#' # or a shortcut to add ids
#' sc=sc+c("720575940621039145", "720575940626877799")
#' sc
#' \dontrun{
#' # paste resultant URL to clipboard to use in neuroglancer
#' clipr::write_clip(as.character(sc))
#' }
#'
#' # you can also modify the URL directly
#' ngl_segments(u)=c("720575940621039145", "720575940626877799")
#' }
#'
#' \dontrun{
#' browseURL(u)
#'
#' ## Summary of different classes of input
#' # from clipboard
#' ngl_segments(clipr::read_clip())
#' # URL
#' ngl_segments("<ngl-scene-url>")
#' # path to file on disk
#' ngl_segments("/path/to/scene.json")
#' # R list
#' ngl_segments(scenelist)
#' }
ngl_segments <- function(x, as_character=TRUE, include_hidden=FALSE,
                         must_work=TRUE, unique=FALSE, ...) {
  checkvalid <- function(x) {
    if(must_work && length(x)==0)
      stop("Sorry. There are no valid segments.")
    vidx=valid_id(x, na.ok = F)
    if(must_work && !all(vidx))
      stop("Sorry. There are some invalid segments.")
    x[!vidx]=0L
    if(unique && anyDuplicated(x)) {
      nx=length(x)
      x=unique(x)
      warning("ngl_segments: Dropping ", nx - length(x)," duplicate ids", call. = F)
    }
    x
  }
  if(is.factor(x) || is.logical(x))
    x=as.character(x)
  if(is.numeric(x)) {
    x <- checkvalid(x)
    return(if(as_character) id2char(x) else as.numeric(x))
  }

  if(is.character(x)) {
    if(isTRUE(substr(x[1],1,1)=="{") && jsonlite::validate(x))
      x=ngl_decode_scene(x, ...)
    else {
      x[!nzchar(x)]=NA
      # character vector of segment ids
      if(all(valid_id(x, na.ok = TRUE)) || length(x)!=1) {
        x <- checkvalid(x)
        return(if(as_character) x else as.numeric(x))
      } else {
        x=ngl_decode_scene(x, ...)
      }
    }
  }

  layers <- ngl_layers(x)
  nls <- ngl_layer_summary(layers)
  flywireseglayers=nls$type=="segmentation_with_graph"
  if(any(flywireseglayers)){
    layers <- layers[flywireseglayers]
    nls <- ngl_layer_summary(layers)
  }

  nallsegs=if(include_hidden) nls$nsegs+nls$nhidden else nls$nsegs

  if(sum(nallsegs)==0) {
    if(must_work) stop("Sorry. No segments entry in this list!")
    # make an empty list
    segments=numeric()
  } else {
    # extract the chosen layer
    sl <- layers[[min(which(nallsegs>0))]]

    if(sum(nls$nsegs>0) > 1)
      warning("Sorry. More than one segments entry in this list:\n",
        paste(nls$name[nallsegs>0], collapse = '\n'))

    segments=unlist(sl[['segments']])
    if(include_hidden)
      segments <- union(segments, unlist(sl[['hiddenSegments']]))
  }
  segments <- checkvalid(segments)
  if(as_character) as.character(segments) else as.numeric(segments)
}

id2char <- function(x) {
  as.character(bit64::as.integer64(x))
}

#' @export
#' @rdname ngl_segments
#' @param value Segment ids in any form understandable by \code{ngl_segments}.
#'   This can include character/numeric/int64 ids, a URL, parsed neuroglancer
#'   scene etc.
#' @description \code{ngl_segments<-} replaces neuroglancer segments in a
#'   neuroglancer scene parsed by \code{\link{ngl_decode_scene}}.
#' @details \code{ngl_segments<-} chooses the FlyWire style
#'   \code{segmentation_with_graph} layer if it exists otherwise the first
#'   visible segmentation layer. Note that \code{hiddenSegment will be removed
#'   in this process}.
`ngl_segments<-` <- function(x, value) {
  was_char <- is.character(x)
  baseurl <- if(was_char) x else NULL
  # choose first non hidden layer to add segments
  x=ngl_decode_scene(x)
  layers=ngl_layers(x)
  nls=ngl_layer_summary(layers)
  # this is flywire specific, but always what you want
  sel=which(nls$type=="segmentation_with_graph")
  # if we can't find that, then go with standard approach
  if(length(sel)==0)
    sel=which(nls$visible & grepl("^segmentation", nls$type))
  if(length(sel)==0)
    stop("Could not find a visible segmentation layer!")
  if(length(sel)>1) {
    warning('Multiple segmentation layers. Choosing first!')
    sel=sel[1]
  }
  if(is.null(value))
    value <- character()
  newsegs=ngl_segments(value, as_character = TRUE, must_work = FALSE)
  if(!all(valid_id(newsegs)))
    warning("There are ", sum(!valid_id(newsegs)), " invalid segments")
  x[['layers']][[sel]][['segments']]=newsegs
  if(nls$nhidden[sel]>0)
    x[['layers']][[sel]][['hiddenSegments']]=NULL
  if(was_char) as.character(x, baseurl=baseurl) else x
}

#' Extract and manipulate layers in a neuroglancer scene
#'
#' @description \code{ngl_layers} extract the neuroglancer layers with
#'   convenience options for selecting layers by characteristics such as
#'   visibility, type etc.
#' @param x a neuroglancer scene object (see \code{\link{ngscene}}) or an
#'   existing \code{nglayers} object (which you probably want to subset).
#' @param subset an expression (evaluated in the style of subset.dataframe)
#'   which defined
#' @return A list of layers with additional class \code{nglayers}
#'
#' @export
#' @aliases nglayers
#'
#' @seealso \code{\link{ngl_decode_scene}}, \code{\link{ngl_layers}},
#'   \code{\link{ngl_segments}}, \code{\link{ngl_encode_url}}
#' @examples
#' \donttest{
#' u="https://ngl.flywire.ai/?json_url=https://globalv1.flywire-daf.com/nglstate/5409525645443072"
#' sc=ngl_decode_scene(u)
#' sc
#' names(ngl_layers(sc))
#' str(ngl_layers(sc))
#'
#' str(ngl_layers(sc, nsegs>0))
#' str(ngl_layers(sc, visible==TRUE))
#' str(ngl_layers(sc, !visible))
#' # flywire segmentation
#' str(ngl_layers(sc, type=="segmentation_with_graph"))
#' # image or segmentation
#' str(ngl_layers(sc, type %in% c("image", "segmentation_with_graph")))
#' }
ngl_layers <- function(x, subset=NULL) {
  layers <- if(inherits(x, 'nglayers')) x
  else {
    if(!is.ngscene(x))
      stop("Unable to extract layer information from ", deparse(substitute(x)),
           " as it is not an ngscene object!")
    layers=x[['layers']]
    class(layers)=c("nglayers", "list")
    layers
  }

  # record the layers as names for ease of manipulation in R
  # these attributes should be stripped off by ngl_encode_url
  df <- ngl_layer_summary(layers)
  names(layers) <-  if(any(is.na(df$name))) seq_along(layers) else df$name

  e <- substitute(subset)
  if(!is.null(e)) {
    r <- eval(e, df, parent.frame())
    if(is.character(r)) r=match(r, df$name)
    layers=layers[r]
    class(layers)=c("nglayers", "list")
  }

  layers
}


#' @export
#' @description \code{ngl_layers<-} sets the layers element of a
#'   \code{\link{ngscene}} object, taking care of name/class details.
#' @rdname ngl_layers
#' @param value a list specifying one or more neuroglancer layers. This will
#'   usually come from a json fragment or another parsed neuroglancer scene. See
#'   examples.
#' @examples
#' # get a sample flywire neuroglancer scene
#' sc=ngl_decode_scene(system.file("flywire-annotations.json" ,
#'   package = 'fafbseg'))
#' sc
#' # save a copy
#' sc.orig <- sc
#' # remove a layer
#' ngl_layers(sc)=ngl_layers(sc)[-3]
#' # or using convenient - notation
#' sc.noann <- sc.orig - "annotation"
#'
#' # reverse layer order
#' ngl_layers(sc)=ngl_layers(sc)[2:1]
#'
#' # keep visible only
#' ngl_layers(sc) <- ngl_layers(sc, visible)
#' # visible + multiple segments
#' ngl_layers(sc) <- ngl_layers(sc, visible & nsegs>0)
#' # flywire segmentation
#' ngl_layers(sc) <- ngl_layers(sc, type=="segmentation_with_graph")
#'
#' # combine layers using + convenience method
#' sc.noann + ngl_layers(sc.orig)['annotation']

#' \dontrun{
#' # combine layers from two scenes
#' ngl_layers(sc) <- c(ngl_layers(sc), ngl_layers(sc2))
#' ngl_layers(sc) <- c(ngl_layers(sc)[-(3:4)], ngl_layers(sc2)[3:4])
#' ngl_layers(sc) <- c(ngl_layers(sc), ngl_layers(sc2)[-1])
#' ngl_layers(sc) <- c(ngl_layers(sc), ngl_layers(sc2)['annotation'])
#' sc
#'
#' # another way to add a single scene
#' ngl_layers(sc)[[4]] <- ngl_layers(sc2)[[4]]
#'
#' # add a new layer to a scene by parsing some JSON from the clipboard
#' # note the double brackets are essential here
#' ngl_layers(sc)[['jfrc_mesh']] <- jsonlite::fromJSON(clipr::read_clip())
#' }
`ngl_layers<-` <- function(x, value) {
  # note that in cases of expressions like
  x[['layers']] <- value
  # this looks a bit odd, but looks after class, names etc
  x[['layers']]=ngl_layers(x)
  x
}

#' @export
`[.nglayers` <- function(x, i) {
  structure(NextMethod("["), class = class(x))
}

null2na <- function(x) sapply(x, function(y) if(is.null(y)) NA else y,USE.NAMES = F)

ngl_layer_summary <- function(layers) {
  if(is.ngscene(layers))
    layers=layers[['layers']]

  sources=sapply(layers, function(x) unlist(x[['source']],use.names = F)[1])
  types=sapply(layers, "[[", "type")
  # nb layers are visible by default
  visible=sapply(layers, function(y) {vis=y$visible;if(is.null(vis)) TRUE else vis})
  nsegs=sapply(layers, function (y) length(y[['segments']]))
  nhidden=sapply(layers, function (y) length(y[['hiddenSegments']]))
  names=sapply(layers, "[[", "name")
  if(length(names)!=length(layers)) names <- as.character(seq_along(layers))

  st = data.frame(
    name = null2na(names),
    type = null2na(types),
    visible = visible,
    nsegs = null2na(nsegs),
    nhidden = null2na(nhidden),
    source = null2na(sources),
    row.names = NULL,
    stringsAsFactors = F
  )
  st
}

ngl_segmentation <- function(x=getOption('fafbseg.sampleurl'), rval=c('url', 'full'), ...) {
  rval=match.arg(rval)
  layers <- ngl_layers(ngl_decode_scene(x), ...)
  st <- ngl_layer_summary(layers)
  # remove any layers without defined sources
  st=st[!is.na(st$source),,drop=FALSE]
  seglayer=grep('seg', st$type)
  if(length(seglayer)==0) {
    NULL
  } else if(rval=='url') {
    st$source[[seglayer[1]]]
  } else {
    layers[[seglayer[1]]]
  }
}

#' @export
#' @description \code{+.ngscene} adds segments or layers to a neuroglancer scene
#' @rdname ngl_layers
#' @section Using + and -: There are shortcut methods that allow you to add or
#'   subtract segments or layers from neuroglancer scenes. These are designed
#'   for convenience in interactive use, but may be a bit fragile for unusual
#'   inputs.
#'
#' @param y Segments or layers to add or remove from a neuroglancer scene.
#'   Segments are provided as character vectors or by applying
#'   \code{\link{ngl_segments}} to a more complex object. Layers to remove
#'   should be the layer name. Layers to add should be in the form of an R list
#'   returned by ng_layers or a JSON fragment copied from neuroglancer.
`+.ngscene` <- function(x, y) {
  if(!is.list(y)) {
    if(all(valid_id(y))) {
      y=ngl_segments(y, as_character = TRUE)
      ngl_segments(x) <- union(ngl_segments(x, must_work = F), y)
      return(x)
    } else {
      parsed <- try(jsonlite::fromJSON(y, simplifyVector = TRUE, simplifyDataFrame = FALSE), silent = TRUE)
      if(inherits(parsed, "try-error"))
        stop("Please supply valid 64 bit integer ids or valid JSON")
      y=parsed
    }
  }
  # if we've got this far we have a list
  if(!inherits(y, 'nglayers')) {
    y=list(y)
  }
  ngl_layers(x) <- c(ngl_layers(x), y)
  x
}

#' @export
#' @description \code{-.ngscene} removes segments or whole layers from a
#'   neuroglancer scene. It does not complain if the segment is not present.
#' @rdname ngl_layers
`-.ngscene` <- function(x, y) {
  if(!is.character(y) && !is.numeric(y))
    stop("I do not yet handle complex input")

  layers <- ngl_layers(x)
  layer_like <- all(y %in% names(layers))
  if(layer_like) {
    if(any(valid_id(y)))
      warning("Assuming that ", deparse(substitute(y)), " identifies neuroglancer layer(s)!")
    tokeep=setdiff(names(layers), y)
    ngl_layers(x)=ngl_layers(x)[tokeep]
  } else {
    if(!all(valid_id(y)))
      stop("Please supply valid 64 bit integer ids!")
    y=ngl_segments(y, as_character = TRUE)
    ngl_segments(x) <- setdiff(ngl_segments(x), y)
  }
  x
}


#' @export
print.ngscene <- function(x, ...) {
  layerdf=ngl_layer_summary(x)
  segs=suppressWarnings(ngl_segments(x, must_work = FALSE))
  segs.all=suppressWarnings(ngl_segments(x, include_hidden = T, must_work = FALSE))

  cat(
    "neuroglancer scene with ",
    nrow(layerdf),
    " layers and ",
    length(segs.all),
    " segments (of which ",
    length(segs),
    " shown)\n",
    sep=""
  )
  print(layerdf)
  invisible(x)
}