ENH: Spatial Indexes

[snowman2]

Discussion: How to apply spatial indexes in rioxarray

Related:

• https://github.com/xarray-contrib/xproj
• https://github.com/xarray-contrib/rasterix
• Experimentally support CRSIndex #588
• Generate spatial coords with rasterix.RasterIndex #855
• Add raster index #846
• https://xdggs.readthedocs.io/en/latest/user_guide/conventions.html

Current method for applying indexes:

#888 (comment)

ds_raster_zarr = xr.open_dataset("store.zarr")
# build xarray indexes from Zarr/CF conventions
ds_raster_idx = (
    ds_raster_zarr
    .set_xindex("crs", xproj.CRSIndex, crs=ds_raster_zarr.rio.crs)
    .assign_coords(xarray.Coordinates.from_xindex(rasterix.RasterIndex.from_transform(ds_raster_zarr.rio.transform()))
)

# use xarray geospatial indexes
ds_result = (
    ds_raster_idx
    .isel(x=..., y=...)
    ...
)

Challenges:

• Metadata from other conventions may become invalid/conflict when performing xarray operations (DataArray Inheritance from Parent Dataset for Zarr Conventions #888 (comment)).

Potential strategy:

• Stage 1: Read only support & provide examples for writing spatial indexes
• Stage 2: Write support

[keewis]

I've used the following pattern for a very similar problem in xdggs (and also pint-xarray): reading from disk is done with the usual xarray backends, followed by a decode step which extracts and strips the relevant metadata and creates the index:

import xarray as xr
import xdggs

ds = xr.open_dataset(..., engine="zarr").dggs.decode(convention="zarr")

After that, xarray operations modify the index as needed, which makes sure the spatial metadata stays up-to-date.

Before writing back to disk we'd encode the index:

ds.dggs.encode(convention="zarr").to_zarr(...)

rioxarray could also follow that pattern, just that its decode method (or however you call it) would create xproj and rasterix indexes.

(There is a proposal to add custom dataset coders to xarray, but this will need quite a bit of thought to figure out how to avoid having multiple coders conflict, so I wouldn't count on this being available anytime soon)

cc @benbovy, @dcherian, @scottyhq

[snowman2]

decode step which extracts and strips the relevant metadata and creates the index

That is a nice strategy. It removes any possible invalidation of input metadata when performing xarray operations.

[dcherian]

It removes any possible invalidation of input metadata when performing xarray operations.

Yes, this is the way IMO. rasterix.assign_index will delete any attributes it interprets to create the index. It keeps track of the transform internally as the Dataset/DataArray are manipulated. What we haven't built is the encode step that destroys the index and adds attributes for a specific convention.

I think .rioxarray should check for RasterIndex in .xindexes and grab the transform from that.
So,

ds = (
    xr.open_dataset("geozarr.zarr")
    # assigns RasterIndex, deletes GeoZarr attributes
    .pipe(rasterix.assign_index)
)

ds.rio.transform  # grabs transform from RasterIndex

The Xarray backend piece can optionally create a new RasterIndex from scratch on data read.

[benbovy]

Thanks @snowman2 for opening this issue and @keewis @dcherian for commenting.

I also think that it is a good strategy to remove any ambiguity and try working with a single source of georeferencing info in Xarray operations.

The idea behind the XarrayIndexesConvention proposed in #888 was similar, assuming we can easily clear and switch between conventions (#882 (comment)), but I agree that it might be better leaving conventions for IO metadata and have a decode() / encode() explicit API here. An Xarray backend option is good as well.

[benbovy]

I'd be happy to take another stab at implementing support for Xarray spatial indexes in Rioxarray (superseding #855).

There seems to be a consensus on the following workflow:

1. decode step
   • extract/strip the revelant geospatial attributes and create the spatial indexes + coordinates
2. manipulate Dataset/DataArray objects using the Xarray (+ Rioxarray, etc.) API
   • the geospatial information is accessed / propagated / updated via the spatial indexes
3. encode step
   • destroy the spatial indexes, maybe drop their associated coordinates, and write the geospatial information as attributes

What would be the best plan forward?

For the decode step, rasterix.assign_index() can already read and strip geospatial attributes related to a few conventions (e.g., CF, Zarr proj/spatial), but the logic it uses is duplicate with Rioxarray's recent convention classes. The latter are implemented in a cleaner & composable way and they support writing attributes. Wouldn't be better to rely on those convention classes here? We'd need to extend the convention protocol in order to clear/strip attributes.

Here is a possible plan (expanded from @snowman2's top comment):

• Stage 1: implement read only support & provide examples for writing spatial indexes
  • ds.rio.transform and ds.rio.crs first check for the presence of spatial indexes. If found, the transform/crs are read from the indexes. Otherwise, it falls back to reading transform/crs using the conventions
  • examples for writing spatial indexes are based on 3rd-party API, e.g., rasterix.assign_index()
• Stage 2: extend Rioxarray's Convention protocol to clear geospatial metadata in a Dataset/DataArray, then add implementations for the existing conventions (CF, Zarr)
• Stage 3: implement the decode step, using the convention classes for reading/stripping attributes
  • add a ds.rio.decode() method, or a decode_* option to xarray.open_dataset(..., engine="rasterio"), or both
• Stage 4: implement the encode step, using the convention classes for writing attributes
  • add a ds.rio.encode() method (there's no Xarray plugin support for writing datasets yet)

[snowman2]

I think implement read only support for spatial indexes would be the best place to start for rioxarray. I would like the logic to check whether the indexes have the crs or transform attributes instead of a specific class. That implementation is agnostic to the index class and allows for alternate libraries with similar functionality to be used.

I think the encode/decode logic should live in rasterix and xproj. This makes it usable outside of rioxarray and ensures consistency across all libraries that use it. It also opens up the possibility of connecting that logic with custom dataset coders to xarray that may be coming in the future.