Add raster index #846

benbovy · 2025-02-20T17:14:21Z

Tests added
Fully documented, including docs/history.rst for all changes and docs/rioxarray.rst for new API

This work in progress PR adds a custom RasterIndex Xarray index that is built on top of pydata/xarray#9543 (available in Xarray's main branch).

Opt-in feature via rioxarray.set_option(use_raster_index=True)
- Alternatively or in addition we could expose this option as an IO backend option. The global option is handy for the other places where new spatial coordinates are generated (e.g., reproject).
RasterIndex should work with both x/y independent axis 1D coordinates (rectilinear affine transform with no rotation) vs. x/y dependent axis 2D coordinates. In the latter case data selection is more limited.
The spatial coordinates freshly generated with a RasterIndex are lazy (in both 1D and 2D cases): they wrap the affine transform.
There are still quite a few things to implement and test.
I skipped implementing CRS-aware functionality here (Experimentally support CRSIndex #588) as I'm focusing on the coordinate transform (also this PR is already getting big), but I guess this could be added later on top of this.

TODO on the Xarray Index API implementation (could also be done after this PR, if needed):

RasterIndex.from_variables: compute the affine transformation from explicit coordinate values, although maybe not very useful?
RasterIndex.concat: possible to compute a new affine transformations by combining the transformations of multiple raster indexes?
RasterIndex.join: like for concat we could perhaps combine the transformations in some simple cases (check that the two affine transforms are compatible)?
RasterIndex.reindex_like: maybe this could be implemented using a forward transformation of the current index chained with a reverse transformation of the other index?
RasterIndex.rename: must propagate the new coordinate and/or dimension names to the underlying AffineTransform objects!
RasterIndex.copy (not sure we need to override Index.copy)
RasterIndex.__getitem__ (not sure we really need to implement it)

Here is an example notebook that I'll update as I continue working on this PR.

There are a few issues that we need to fix in Xarray so this PR can work properly (I've added a few temporary workarounds when possible):

cc @scottyhq @dcherian @maxrjones

Based on coordinate transform examples copied and adapted from pydata/xarray#9543.

Use a mapping of coordinate name(s) to underlying index(es) behind the hood.

I encountered pre-commit/pre-commit#1375 as I didn't have python 3.10 installed in my environment so pre-commit couldn't find it. I applied this suggestion: pre-commit/pre-commit#1375 (comment).

Create a RasterIndex if option use_raster_index is True.

see pydata/xarray#10063

benbovy · 2025-02-20T17:17:16Z

rioxarray/raster_index.py

+            return False
+
+        return all(
+            index.equals(other._wrapped_indexes[k])  # type: ignore[arg-type]


Need to relax the typing constraints of xarray.core.indexes.CoordinateTransformIndex.equals() (the other argument should be any xarray.Index) so we can remove the # type: ignore here.

dcherian · 2025-02-20T18:05:49Z

rioxarray/raster_index.py

+from xarray.core.indexing import IndexSelResult, merge_sel_results
+
+
+class AffineTransform(CoordinateTransform):


I think the Affine* classes should live in Xarray. They are useful for both geo and bio-imaging domains at least (and possibly astro too).

Yes I agree these are useful in many domains! The wrapped affine.Affine object is here restricted to 2D affine transformation of the plane so it might be limited for general usage, though. These classes are also thin wrappers.

They are "thin" in code but 'deep' in understanding of the Xarray model. That's why I think we should vendor it. Even so, vendoring will save a bunch of your time as you iterate on this since you wouldn't need to keep code in two places in-sync.

Maybe the whole content of raster_index.py including the RasterIndex class could (eventually) be moved into Xarray or some dedicated package? There's nothing IO specific and it could be useful in other downstream packages like odc-geo, geoxarray, etc.

benbovy · 2025-02-21T09:19:46Z

rioxarray/raster_index.py

+        self.axis_transform = transform
+        self.dim = transform.dim
+
+    def isel(  # type: ignore[override]


Again here we should make the return type of xarray.core.Index.isel() more permissive (i.e., any Index object or None) so we can remove the # type: ignore comment.

benbovy · 2025-02-21T15:05:33Z

rioxarray/raster_index.py

+        if new_indexes:
+            # TODO: if there's only a single PandasIndex can we just return it?
+            # (maybe better to keep it wrapped if we plan to later make RasterIndex CRS-aware)
+            return RasterIndex(new_indexes)
+        else:
+            return None


When it is possible, it would be convenient to return a PandasIndex instead of a RasterIndex for reusing the selection results in further Xarray operations that require alignment with other Xarray objects. The reason is that Xarray alignment is based on strict Index types, e.g., it cannot compare or align a RasterIndex (encapsulating a single PandasIndex) with another PandasIndex.

However, we can only return a PandasIndex in a limited number of cases, e.g.,

it is possible for da.isel(x=[...], y=0) (or da.isel(x=0, y=[...])) to return a "x" (or "y") coordinate with a PandasIndex and to drop the index for the "y" (or "x") scalar coordinate.

it is not possible for da.isel(x=[...], y=[...]) to return both "x" and "y" coordinates each with a PandasIndex. This is because xarray.Index.isel() doesn't currently allow returning multiple index objects.

Returning a PandasIndex in some cases and a RasterIndex in other cases would make the behavior complex and pretty confusing.

Alternatively, we could always return a RasterIndex for both "x" and "y" coordinates even when selecting it with scalar values. This would provide simple behavior and might facilitate CRS-aware operations even with scalar spatial coordinates, but this would be even less Xarray-alignment friendly.

So in short: it is not straightforward to choose the right behavior here. It is a trade-off.

but this would be even less Xarray-alignment friendly.

I think this is fine. Users can reset_index once they are sure what they're about to do makes sense

snowman2 · 2025-04-10T16:41:37Z

This has been an in-progress MR for a while. Any plans to keep working on it?

dcherian · 2025-04-10T19:19:56Z

Yes, we landed on keeping the index separate since it's generally useful outside just rioxarray. Are you ok with adding another dependency for this feature?

A few of us are meeting Monday morning 4/21 8.30 MT. Would you like to join?

snowman2 · 2025-04-10T20:12:02Z

A few of us are meeting Monday morning 4/21 8.30 MT. Would you like to join?

Sure, I would be happy to join.

benbovy · 2025-04-16T13:07:44Z

I'm closing this PR.

RasterIndex has been implemented in https://github.com/dcherian/rasterix. See #855 for integration within Rioxarray.

benbovy added 11 commits February 19, 2025 16:36

wip add RasterIndex

550c35f

Based on coordinate transform examples copied and adapted from pydata/xarray#9543.

refactor RasterIndex (cleaner logic)

1031770

Use a mapping of coordinate name(s) to underlying index(es) behind the hood.

fix running pre-commit for all python3 versions

7c96f49

I encountered pre-commit/pre-commit#1375 as I didn't have python 3.10 installed in my environment so pre-commit couldn't find it. I applied this suggestion: pre-commit/pre-commit#1375 (comment).

implement to_pandas_index + lint

587199a

add use_raster_index global option (opt-in)

2977f6c

refactor generate spatial coords

06c453e

Create a RasterIndex if option use_raster_index is True.

misc. fixes

d390aed

add RasterIndex repr

f540461

temp workaround for xarray isel issue

6851b91

see pydata/xarray#10063

fix affine slicing

19fe15a

fix isel with arbitray array values

09c2d7a

benbovy commented Feb 20, 2025

View reviewed changes

dcherian reviewed Feb 20, 2025

View reviewed changes

typo

158830b

benbovy commented Feb 21, 2025

View reviewed changes

benbovy mentioned this pull request Feb 21, 2025

More permissive Index typing pydata/xarray#10067

Merged

benbovy commented Feb 21, 2025

View reviewed changes

fix vectorized (fancy) isel

a5f9772

benbovy mentioned this pull request Mar 26, 2025

Understanding concerns with CF encoding of CRS zarr-developers/geozarr-spec#20

Open

benbovy mentioned this pull request Apr 10, 2025

Add some notes on RasterIndex design xarray-contrib/rasterix#4

Merged

This was referenced Apr 14, 2025

Add RasterIndex (affine) xarray-contrib/rasterix#5

Merged

Generate spatial coords with rasterix.RasterIndex #855

Draft

benbovy closed this Apr 16, 2025

		from xarray.core.indexing import IndexSelResult, merge_sel_results


		class AffineTransform(CoordinateTransform):

Add raster index #846

Add raster index #846

Uh oh!

Conversation

benbovy commented Feb 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

benbovy Feb 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

benbovy Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

dcherian Feb 21, 2025

Choose a reason for hiding this comment

Uh oh!

benbovy Feb 24, 2025

Choose a reason for hiding this comment

Uh oh!

benbovy Feb 21, 2025

Choose a reason for hiding this comment

Uh oh!

benbovy Feb 21, 2025

Choose a reason for hiding this comment

Uh oh!

dcherian Apr 3, 2025

Choose a reason for hiding this comment

Uh oh!

snowman2 commented Apr 10, 2025

Uh oh!

dcherian commented Apr 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

snowman2 commented Apr 10, 2025

Uh oh!

benbovy commented Apr 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

benbovy commented Feb 20, 2025 •

edited

Loading

benbovy Feb 20, 2025 •

edited

Loading

dcherian commented Apr 10, 2025 •

edited

Loading