Slicing of structured curvilinear grids

[dcherian]

Is your feature request related to a problem?

Currently NDPointIndex (e.g. with KDTree) only allows advanced indexing.

Can we make it support slicing too (with appropriate semantics)?

Describe the solution you'd like

I just wrote this up

import xarray as xr

ds = xr.tutorial.open_dataset("ROMS_example").set_xindex(("lat_rho", "lon_rho"), xr.indexes.NDPointIndex,)

# intended slicer
slicers = dict(lat_rho=slice(28, 29), lon_rho=slice(-91, -89))

To get this to work, I did this:

import itertools

index = ds.xindexes["lat_rho"]
edges = tuple((slicer.start, slicer.stop) for slicer in slicers.values())
vectorized_sel = {
    name: xr.DataArray(dims=("pts",), data=data)
    for name, data in zip(
        slicers.keys(), map(np.asarray, zip(*itertools.product(*edges)))
    )
}
idxrs = index.sel(vectorized_sel, method="nearest").dim_indexers
new_slicers = {
    name: slice(array.min().item(), array.max().item()) for name, array in idxrs.items()
}
subset =ds.sel(new_slicers)

Output

As you can see: this effectively defines slicing as an operation that "returns a continuous slice of data such that every point within the bounding box defined by the slice objects is returned." Does this make sense?

ds.salt.isel(s_rho=-1, ocean_time=0).plot(x="lon_rho", y="lat_rho")
ds.salt.isel(s_rho=-1, ocean_time=0, **new_slicers).plot(
    x="lon_rho", y="lat_rho", cmap="Blues", add_colorbar=False
)

x0, x1 = slicers["lon_rho"].start, slicers["lon_rho"].stop
y0, y1 = slicers["lat_rho"].start, slicers["lat_rho"].stop
import matplotlib.pyplot as plt

plt.plot([x0, x1, x1, x0, x0], [y0, y0, y1, y1, y0], color="r")

Describe alternatives you've considered

We could mask out data outside the box, but that then becomes a copy instead of a view.

Additional context

No response

[keewis]

conceptually, kdtrees should support range queries, but whether or not the specific implementation (e.g. scipy.spatial.KDTree) supports that is a different question.

[dcherian]

Doesn't seem like it's supported; but the definition of range query also returns all points within the box, which doesn't necessary yield a dense array of the same dimensionality (it's ragged in the general case). Also, it seems a lot cheaper to simply query the corner points

[keewis]

that makes sense; point indexes like kdtree or balltree work on an unstructured list of points, so they will only ever return a list of points, as well.

Which also means that your trick only works for structured grid coordinates.

[dcherian]

You are right. My approach assumes sortedness so perhaps it can hidden behind a specific method or we set a flag at construction time

[dcherian]

Thinking about this more, I think we need a different index structure that recognizes the structured nature of this grid. I am not sure what that structure is.

cc @rsignell @klindsay28 @hetland @ChrisBarker-NOAA, all of whom have lots of experience with this kind of thing ;). For reference, we've recently added KDTree indexing to Xarray (example) and are trying to figure out how to make something useful for structured curvilinear grids.

[ChrisBarker-NOAA]

Actually, I'm not sure you need anything fancy for this, though it depends a bit on what you want to do.

The trick is that a structured curvilinear grid like this is representable as simple rectangular arrays, but a "rectangle" in array index space is not (necessarily) rectangular in "world" coordinates.

So if you want to subset (slice) an unstructured grid to a rectangle in world coordinates, you have to get either:

A) A non-logically rectangular subset, which isn't really a slicing operation

or

B) A larger subset that is logically rectangular, but is not rectangular in world coordinates.

I've wanted to implement A) for ages, but the result would no longer be a structured grid, and wouldn't fit well into the xarray data model.

Most folks (actually, all that I know of) do (B) -- and it works fine, if a bit odd in some pathological cases.

it seems a lot cheaper to simply query the corner points

Exactly -- that's (B) and not only cheaper and easier, but necessary if you want to return a logically rectangular result.

Actually, just the corner points is not good enough. Imagine a concave curvature along one side -- if you get the corners, you may miss a few of the cells that are inside the box. but it's still a simple <,> comparison -- no need for a fancy spatial index [*].

This has been implemented in numerous place, but one example in the xarray_subset_grid project:

https://github.com/ioos/xarray-subset-grid

See:

https://github.com/ioos/xarray-subset-grid/blob/main/xarray_subset_grid/grids/sgrid.py

NOTE: I'm a bit wary of using this proposed syntax exactly -- in this case, it maps fairly well to slicing. But in other cases, not so much. For example, I've wanted for years for xarray to have an abstraction for "subsetting" that is not directly mapped to the assumptions of rectangular arrays -- for example, for subsetting unstructured grids.

To do that, it would make sense (to me) to have a clearly different operation, even if it was, in the case of rectangular grids, exactly the same, and in the case of curvilinear grids, almost the same.

That would open the door not just to unstructured grids, but also to subsetting to, say, a polygon rather than a simple bounding box.

That being said this particular idea seems harmless and an useful :-)

[*] I don't think you need it in this case, but if you do need a fancy spatial index at some point, you might want to consider celltree:

https://github.com/NOAA-ORR-ERD/cell_tree2d

[it hasn't been touched in years -- not because it's not maintained, but because it is stable -- it's been working fine for years -- we've been keeping conda-forge build up to date: https://anaconda.org/conda-forge/cell_tree2d/files)

[dcherian]

I've wanted for years for xarray to have an abstraction for "subsetting" that is not directly mapped to the assumptions of rectangular arrays -- for example, for subsetting unstructured grids.

Chris, we have this abstraction now! Here is one complex example.

[ChrisBarker-NOAA]

@dcherian : very cool! I'll need to dig into that and see if it can be used for the, e.g. subsetting a unstructured grid case.

Then maybe the problem-at-hand should be done this way, rather than as slices?

[keewis]

I believe the idea is to treat slices along indexed dimensions as a bounding box, then pass that to the wrapped index. It may also be possible to create virtual coordinates (indexed by the same index) to allow indexing by polygons / shapely.box objects / bbox arrays.

[hetland]

So, I think there are a few issues. One is making sure things are selected efficiently, and I believe this is the basis of the comment by @keewis. But after there is some reasonable bounding box, what next.

In my trials, I have found it hard to use kdtrees to interpolate to arbitrary points -- the general case of slicing. The problem is that no matter how many closest point you select with kdtree, one or many, the interpolation will 'jump' as you march along the points in an interpolated transect. There is no way to smooth this out, as two adjacent points eventually will have different input points and this creates a discontinuity in input information.

The best and most performant solution that I have found is based on cell_tree2d. The gridded package has an example of how to use this: https://github.com/pyoceans/gridded