R/domain_geozarr.R

In geozarr: GeoZarr Conventions for Geospatial Data in Zarr Stores

#' GeoZarr domain object
#'
#' @description This class implements a GeoZarr domain object. A GeoZarr domain
#'   object is a `zarr_domain` descendant object identifying groups and arrays
#'   in a `zarr` object that are formatted using GeoZarr conventions.
#'
#'   This domain supports the standard conventions for GeoZarr data sets,
#'   specifically the `cs` and `spatial` conventions, as well as two other
#'   formats for geospatial Zarr data that predate GeoZarr: the format used by
#'   the Python XArray library (Zarr v.2 and v.3) and the NCZarr format (Zarr
#'   v.2).
#' @docType class
#' @export
zarr_domain_geozarr <- R6::R6Class('zarr_domain_geozarr',
  inherit = zarr::zarr_domain,
  cloneable = FALSE,
  private = list(
    is_xarray_group = function(name, metadata, parent, store) {
      if (metadata$node_type == 'array')
        return(FALSE)

      # List immediate children of this group
      prefix <- if (is.null(parent)) ''
                else paste0(parent$prefix, name, '/')
      children <- store$list_dir(prefix)
      if (!length(children)) return(FALSE)

      # Fetch metadata for each child, keep only arrays (skip sub-groups)
      child_meta <- lapply(children, \(child) {
        childprefix <- paste0(prefix, child, '/')
        store$get_metadata(childprefix)
      })
      child_meta <- Filter(\(m) !is.null(m) && m$node_type == "array", child_meta)
      if (length(child_meta) == 0L) return(FALSE)

      # Every array must have dimension names, in either v2 or v3 form
      all(vapply(child_meta, \(m) {
        !is.null(m$dimension_names) ||             # v3 native field
        !is.null(m$attributes$`_ARRAY_DIMENSIONS`) # v2 attribute
      }, logical(1L)))
    },

    # Detect if the referenced array is XArray formatted. Dimension variables are
    # excluded.
    is_xarray_array = function(name, metadata, parent, store) {
      if (is.null(parent) || metadata$node_type == 'group')
        return(FALSE)

      attrs <- metadata$attributes %||% list()

      # v.2: unambiguous
      if (!is.null(attrs$`_ARRAY_DIMENSIONS`))
        return(TRUE)

      # v.3: dimension_names is necessary but not sufficient
      dims <- metadata$dimension_names
      if (is.null(dims) || length(dims) == 0L)
        return(FALSE)

      # Exclude dimension coordinate arrays (1D, same name as its dimension)
      if (length(dims) == 1L && dims[1L] == name)
        return(FALSE)

      # XArray data variables always carry the "coordinates" attribute
      # Secondary arrays (coordinate values, boundaries) never do
      coords_attr <- attrs$coordinates
      !is.null(coords_attr) && is.character(coords_attr) && nchar(coords_attr) > 0L
    }
  ),
  public = list(
    #' @description Create a new GeoZarr domain instance. The GeoZarr domain
    #'   instance manages the groups and arrays in the Zarr store that it refers
    #'   to. This instance provides access to all objects in the Zarr store.
    #' @returns A `zar_domain_geozarr` object.
    initialize = function() {
      super$initialize('GeoZarr')
    },

    #' @description This method will create a `geozarr_array` for an array node
    #'   and a `geozarr_group` for a group node with GeoZarr conventions
    #'   declared in its attributes. Either the "spatial" or "cs" convention
    #'   has to be declared or the Zarr store has to be formatted using XArray
    #'   or NCZarr or this domain will decline to manage the node.
    #' @param name The name of the node.
    #' @param metadata List with the metadata of the node.
    #' @param parent The parent node of this new node. May be `NULL` for a root
    #'   node.
    #' @param store The store to persist data in.
    #' @return A `geozarr_array` or `geozarr_group` instance if supported,
    #'   `FALSE` otherwise.
    build = function(name, metadata, parent, store) {
      conv <- metadata$attributes$zarr_conventions
      if (is.null(conv)) {
        # Check for XArray
        if (private$is_xarray_array(name, metadata, parent, store)) {
          return(geozarr_array$new(name, metadata, parent, store))
        }

        # Check for NCZarr

        # No fun
        return(FALSE)
      }

      gz_conv <- GeoZarr.options$conventions
      for (cv in seq_along(conv)) {
        if (conv[[cv]]$name %in% gz_conv$name) {
          if (metadata$node_type == 'array')
            return(geozarr_array$new(name, metadata, parent, store))
          else
            return(geozarr_group$new(name, metadata, parent, store))
        }
      }

      # `spatial` convention: array may use parent attributes
      if (metadata$node_type == 'array' && inherits(parent, 'geozarr_group') &&
          'spatial' %in% parent$attributes$zarr_conventions)
        return(geozarr_array$new(name, metadata, parent, store))

      FALSE
    }
  )
)

#' Set the GeoZarr convention with array details in the metadata of the array
#'
#' This function will write the GeoZarr convention with all details into the
#' metadata of a Zarr array.
#'
#' @param metadata A `list` with the basic metadata of the array.
#' @param coord_sys The [CoordinateSystem] instance of the GeoZarr array.
#' @param crs Optional, a `list` with 1 or more CRS definitions. Named list
#'   elements can be "compound",  "spatial" (X-Y), "vertical" or "temporal".
#'   Those elements are a `list` themselves with one or more of the named
#'   elements "code", "wkt2" or "projjson" (for the attributes in the `proj`
#'   convention) with their value being the attribute to register. In the `cs`
#'   convention, a "compound" CRS will be associated with the coordinate system
#'   of the GeoZarr array, the other options will be associated with the
#'   specific CRSs of the coordinate system. In the `spatial` convention there
#'   can only be 1 CRS which must be of type "compound" or "spatial" and which
#'   is registered at the root of the "attributes" of the GeoZarr array.
#' @param external_group Optional, the path to the group, relative to the
#'   location of the array, that stores any external arrays with coordinate
#'   values.
#' @param registration Optional, the registration point of the array for use
#'   with the "spatial" convention. Defaults to "pixel".
#' @return A `list` with the metadata updated with convention attributes.
#' @export
#' @examples
#' ab <- zarr::array_builder$new()
#' ab$data_type <- "int32"
#' ab$shape <- c(1000L, 20L)
#'
#' crd1 <- CoordinatesPacked$new("X_coordinates", "EAST", "m", c(0, 1000), 1000L)
#' ax1 <- CoordinateSystemAxis$new("Axis1", "X", crd1)
#' crd2 <- CoordinatesPacked$new("Y_coordinates", "NORTH", "m", c(0, 1000), 20L)
#' ax2 <- CoordinateSystemAxis$new("Axis2", "Y", crd2)
#'
#' cs <- CoordinateSystem$new("CS", list(Axis1 = ax1, Axis2 = ax2))
#'
#' crs <- list(spatial = list(code = "EPSG:4326"))
#'
#' set_convention(ab$metadata(), cs, crs)
set_convention <- function(metadata, coord_sys, crs = NULL, external_group, registration = 'pixel') {
  meta <- metadata
  atts <- meta$attributes %||% list()
  axes <- coord_sys$axes

  # Drop any existing information
  meta$dimension_names <- NULL
  atts$zarr_conventions <- NULL
  if (length(atts)) {
    atts <- atts[!startsWith(names(atts), c('spatial:', 'proj:'))] # Drop any old spatial and proj elements
    atts$cs <- NULL # Remove any previous cs information
  }

  # dimension_names
  meta <- append(meta, list(dimension_names = vapply(axes, function(ax) ax$name, character(1L), USE.NAMES = FALSE)))

  # Axis abbreviation
  ax_abbr <- vapply(axes, function(ax) ax$abbreviation, FUN.VALUE = character(1), USE.NAMES = FALSE)
  X_axis <- which(ax_abbr == 'X')
  Y_axis <- which(ax_abbr == 'Y')
  if (!length(X_axis) && !length(Y_axis))
    stop('Cannot convert to GeoZarr: No X and/or Y axes found', call. = FALSE)

  # Set GeoZarr convention attributes
  if (length(X_axis) && length(Y_axis) && length(ax_abbr) <= 3L &&
      !('Z' %in% ax_abbr) && !('T' %in% ax_abbr) &&
      inherits(axes[[X_axis]]$coordinates$values, 'CoordinateValuesNumericPacked') && # == numeric & regular
      inherits(axes[[Y_axis]]$coordinates$values, 'CoordinateValuesNumericPacked') &&
      axes[[Y_axis]]$coordinates$values$raw[2L] < 0) {                                # == Y values descending
    # spatial convention
    # X + Y, optionally a band, no others, and X + Y coordinates are numeric and regular
    spatial <- zarr_conv_spatial$new()
    atts <- spatial$register(atts)

    dimensions <- c(axes[[Y_axis]]$name, axes[[X_axis]]$name)
    spatial$dimensions <- dimensions
    spatial$set_coordinates(shape = c(axes[[X_axis]]$length, axes[[Y_axis]]$length),
                            x = axes[[X_axis]]$coordinates$values$raw,
                            y = axes[[Y_axis]]$coordinates$values$raw,
                            registration = registration)
    atts <- c(atts, spatial$as_list())

    if (!is.null(crs) && is.list(crs) && length(crs) == 1L &&
        all(names(crs) %in% c('compound', 'spatial')) && length(crs[[1L]]) &&
        all(names(crs[[1L]]) %in% c('code', 'wkt2', 'projjson'))) {
      proj <- zarr_conv_proj$new(crs[[1L]])
      atts <- proj$register(atts)
      atts <- proj$write(atts)
    } else
      stop('Argument `crs` when provided must be a `list` with formatted CRS elements', call. = FALSE)
  } else {
    # cs convention
    cs_conv <- zarr_convention_cs$new()
    atts    <- cs_conv$register(atts)

    # Direction lookup by axis abbreviation
    cs_direction <- c(X = 'EAST', Y = 'NORTH', Z = 'UP', T = 'FUTURE', OTHER = 'OTHER')

    has_external <- FALSE
    axis_defs <- lapply(axes, function(ax) {
      # Values
      values <- ax$coordinates$raw

      packed <- if (inherits(ax$coordinates, 'CoordinatesTime')) {
        # Time coordinates are always held in unpacked form so must check if they are regular
        if (.is_regular(values)) {
          values <- c(values[1L], values[2L] - values[1L])
          TRUE
        } else FALSE
      } else inherits(ax$coordinates, 'CoordinatesPacked')

      values_def <- if (packed)
        cs_conv$values_regular(values[1L], values[2L])
      else if (ax$length <= GeoZarr.options$max_explicit)
        cs_conv$values_explicit(values)
      else {
        # External coordinate values: Write coordinate values to an external array.
        # The name of the external array is the same as the name of the axis. The
        # actual writing to the external array should be done in the calling code.
        has_external <<- TRUE
        cs_conv$values_external(paste0(external_group, '/', ax$name))
      }

      # Boundary values
      bnds <- ax$coordinates$bounds_raw
      bnds_def <- if (is.vector(bnds))
        cs_conv$boundaries_regular(bnds[1L], bnds[2L])
      else if (is.matrix(bnds))
        # External boundary values: Write boundary values to an external array.
        # The name of the external array is `<axis_name>_bounds`. The
        # actual writing to the external array should be done in the calling code.
        cs_conv$values_external(paste0(external_group, '/', paste0(ax$name, '_bounds')))
      else NULL

      # Time
      time_def <- if (inherits(ax$coordinates, 'CoordinatesTime')) {
        def <- strsplit(ax$coordinates$time$calendar$definition, ' ', fixed = TRUE)[[1L]]
        cs_conv$time(unit = def[1L], epoch = def[3L], calendar = ax$coordinates$time$calendar$name)
      } else NULL

      # Coordinates and axis
      coords_def <- cs_conv$coordinates(values_def, unit = ax$coordinates$unit,
                                        boundaries = bnds_def, time = time_def)
      abbr <- ax$abbreviation
      if (abbr == ' ') abbr <- 'OTHER'
      direction  <- cs_direction[[abbr]]
      if (abbr == 'OTHER') abbr <- ''
      cs_conv$axis(list(coords_def), abbreviation = abbr, direction = direction)
    })

    # Do we need the ref convention?
    if (has_external) {
      ref_conv <- zarr::zarr_convention_ref$new()
      atts <- ref_conv$register(atts)
    }

    # Group axes into separate CRS objects by axis category
    if (is.null(crs)) {
      cs_conv$add_crs(axes = axis_defs[c(X_axis, Y_axis)], type = 'planar')
      cs_conv$add_crs(axes = axis_defs[which(ax_abbr == 'Z')], type = 'vertical')
      cs_conv$add_crs(axes = axis_defs[which(ax_abbr == 'T')], type = 'temporal')
      cs_conv$add_crs(axes = axis_defs[which(ax_abbr == 'OTHER')])
    } else {
      proj <- if (is.null(crs$spatial)) NULL else zarr_conv_proj$new(crs$spatial)
      cs_conv$add_crs(axes = axis_defs[c(X_axis, Y_axis)], type = 'planar', id = proj)
      proj <- if (is.null(crs$vertical)) NULL else zarr_conv_proj$new(crs$vertical)
      cs_conv$add_crs(axes = axis_defs[which(ax_abbr == 'Z')], type = 'vertical', id = proj)
      proj <- if (is.null(crs$temporal)) NULL else zarr_conv_proj$new(crs$temporal)
      cs_conv$add_crs(axes = axis_defs[which(ax_abbr == 'T')], type = 'temporal', id = proj)
      cs_conv$add_crs(axes = axis_defs[which(ax_abbr == 'OTHER')])

      if (!is.null(crs$compound)) {
        proj <- zarr_conv_proj$new(crs$compound)
        cs_conv$cs_crs <- proj
      }
      if (!is.null(proj))
        atts <- proj$register(atts)
    }

    atts <- c(list(cs = cs_conv$as_list()), atts)
  }

  meta$attributes <- NULL # Remove any previous attributes
  meta$attributes <- atts
  meta
}