Skip to content

Generate spatial coords with rasterix.RasterIndex #855

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
benbovy wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Generate spatial coords with rasterix.RasterIndex #855

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d68e5b4
add use_raster_index global option (opt-in)
benbovy Feb 20, 2025
a3eecd9
refactor generate spatial coords
benbovy Feb 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion rioxarray/_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from packaging import version
from rasterio.errors import NotGeoreferencedWarning
from rasterio.vrt import WarpedVRT
from xarray import Dataset, IndexVariable
from xarray import Coordinates, Dataset, IndexVariable
from xarray.backends.common import BackendArray
from xarray.backends.file_manager import CachingFileManager, FileManager
from xarray.backends.locks import SerializableLock
Expand Down Expand Up @@ -1212,6 +1212,7 @@ def open_rasterio(

# Get geospatial coordinates
if parse_coordinates:
coords = Coordinates(coords)
coords.update(
_generate_spatial_coords(
affine=riods.transform, width=riods.width, height=riods.height
Expand Down
10 changes: 10 additions & 0 deletions rioxarray/_options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,18 @@

EXPORT_GRID_MAPPING = "export_grid_mapping"
SKIP_MISSING_SPATIAL_DIMS = "skip_missing_spatial_dims"
USE_RASTER_INDEX = "use_raster_index"

OPTIONS = {
EXPORT_GRID_MAPPING: True,
SKIP_MISSING_SPATIAL_DIMS: False,
USE_RASTER_INDEX: False,
}
OPTION_NAMES = set(OPTIONS)

VALIDATORS = {
EXPORT_GRID_MAPPING: lambda choice: isinstance(choice, bool),
USE_RASTER_INDEX: lambda choice: isinstance(choice, bool),
}


Expand Down Expand Up @@ -60,6 +63,13 @@ class set_options: # pylint: disable=invalid-name
If True, it will not perform spatial operations on variables
within a :class:`xarray.Dataset` if the spatial dimensions
are not found.
use_raster_index: bool, default=False
If True, this option will generate (lazy) spatial coordinates wrapping the
affine :meth:`~rioxarray.rioxarray.XRasterBase.transform` and will associate
them with a custom ``RasterIndex`` Xarray index.
Otherwise this will generate spatial coordinates using explict values
computed by forward transformation and will associate each of the 1-dimensional
spatial coordinates with a default (pandas) index.


Usage as a context manager::
Expand Down
44 changes: 30 additions & 14 deletions rioxarray/rioxarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,23 +121,39 @@ def affine_to_coords(

def _generate_spatial_coords(
*, affine: Affine, width: int, height: int
) -> dict[Hashable, Any]:
) -> xarray.Coordinates:
"""get spatial coords in new transform"""
new_spatial_coords = affine_to_coords(affine, width, height)
if new_spatial_coords["x"].ndim == 1:
return {
"x": xarray.IndexVariable("x", new_spatial_coords["x"]),
"y": xarray.IndexVariable("y", new_spatial_coords["y"]),
}
return {
"xc": (("y", "x"), new_spatial_coords["x"]),
"yc": (("y", "x"), new_spatial_coords["y"]),
}

if get_option("use_raster_index"):
try:
from rasterix import RasterIndex
except (ImportError, ModuleNotFoundError) as err:
raise ImportError(
"The rasterix package is needed for generating "
"the spatial coordinates with a RasterIndex"
) from err

raster_index = RasterIndex.from_transform(affine, width, height)
Copy link

@dcherian dcherian Apr 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should also delete the transform from the grid_mapping variable like so:

grid_mapping = ds.rio.grid_mapping
obj[grid_mapping].attrs.pop("GeoTransform")

following https://github.com/dcherian/rasterix/blob/main/design_notes/raster_index.md#read-time and our discussions. This will trigger some more changes (from a quick search, _cached_transform will need to pull the Affine from RasterIndex)

return xarray.Coordinates.from_xindex(raster_index)

else:
new_spatial_coords = affine_to_coords(affine, width, height)
if new_spatial_coords["x"].ndim == 1:
coords = {
"x": xarray.IndexVariable("x", new_spatial_coords["x"]),
"y": xarray.IndexVariable("y", new_spatial_coords["y"]),
}
else:
coords = {
"xc": (("y", "x"), new_spatial_coords["x"]),
"yc": (("y", "x"), new_spatial_coords["y"]),
}
return xarray.Coordinates(coords)


def _get_nonspatial_coords(
src_data_array: Union[xarray.DataArray, xarray.Dataset]
) -> dict[Hashable, Union[xarray.Variable, xarray.IndexVariable]]:
) -> xarray.Coordinates:
coords: dict[Hashable, Union[xarray.Variable, xarray.IndexVariable]] = {}
for coord in set(src_data_array.coords) - {
src_data_array.rio.x_dim,
Expand All @@ -162,7 +178,7 @@ def _get_nonspatial_coords(
src_data_array[coord].values,
src_data_array[coord].attrs,
)
return coords
return xarray.Coordinates(coords)


def _make_coords(
Copy link

@dcherian dcherian Apr 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It'd be nice to have an accessor function to set this index even if I opened a Zarr file for example. Is there something like ds.rio.assign_coords? I didn't see anything similar in the docs. Perhaps we should have ds.rio.assign_indexes(raster:bool=True, crs: bool=True) that does both RasterIndex and CRSIndex

Copy link
Member

@snowman2 snowman2 Apr 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you think that would be a good fit to add to a rasterix accessor? If you do that, you just set decode_coords=False and assign the index/coords post load.

Copy link

@dcherian dcherian Apr 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It needs a bunch of extra info : x dim name, y dim name, transform, and so it's pretty convenient to just do all that with the .rio accessor. It seems like a prudent and convenient addition

Copy link
Member

@snowman2 snowman2 Apr 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inside rasterix, it could look for the .rio , .odc, or .geo accessors inside the method. It could be done similar to open_dataset where it automatically figures out how to open files based on what's available. It also allows the user select their preferred engine. If they don't exist, then raise an error.

Copy link

@dcherian dcherian Apr 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is an option but to me it feels a lot cleaner to add a method under the rio namespace. (and similarly under the other namespaces if they thought it was a good idea). At least for rioxarray it's just a small wrapper around an existing function, so it doesn't seem like meaningful extra maintenance burden to me.

Copy link
Author

@benbovy benbovy Apr 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inside rasterix, it could look for the .rio , .odc, or .geo accessors inside the method

Hmm this wouldn't be ideal IMHO, assuming that the goal of the rasterix (and xproj) project is to provide reusable utilities to downstream packages like rioxarray, odc-geo, etc. even though technically looking for those accessors is possible without having cyclic dependencies.

Under the same assumption, rasterix and xproj should also be flexible / not strongly opinionated about the names of the spatial coordinates and dimensions, etc.

Expand All @@ -172,7 +188,7 @@ def _make_coords(
dst_width: int,
dst_height: int,
force_generate: bool = False,
) -> dict[Hashable, Any]:
) -> xarray.Coordinates:
"""Generate the coordinates of the new projected `xarray.DataArray`"""
coords = _get_nonspatial_coords(src_data_array)
if (
Expand Down
Loading