How can I implement a simple "Range Index" using custom Indexes?

[rabernat]

Motivated by discussions in the GeoZarr spec and Pangeo Discourse, I've convinced myself that it's important for Xarray to be able to represent "analytic" coordinates, i.e. coordinates defined via a mathematical formula and never explicitly materialized into a discrete array.

This is impossible today, even with the new flexible indexes API, and I think this is something we need to fix.

To boil it down to the simplest possible example, imagine we have a 1D array representing evenly spaced values in the interval (0, 1).

I can easily create an accessor that allows me to query this data without ever explicitly materializing the full array of coordinates into memory. Here's an example

import xarray as xr

npoints = 20

# define a 1D dataset with dimension `x` and no dimension coordinate 
ds = xr.DataArray(
    data=list(range(npoints)),
    dims=['x'],
    name="foo"
).to_dataset()

# define a simple convention for how to encode the range information in metadata
# Note: there are many different ways to do this, what I picked here is arbitrary
# The point is that there is no variable called x
ds.attrs["bounds"] = {
    "x": {"left": 0.0, "right": 1.0}
}

# define an accessor

@dataclass(frozen=True)
class BoundsIndexer:
    left: float
    right: float
    npoints: int

    def get_nearest_point(self, value):
        if value < self.left or value > self.right:
            raise ValueError("Value outside of bounds")
        return int((value - self.left) / (self.right - self.left) * ( self.npoints - 1))

@xr.register_dataset_accessor("bounds")
class BoundsAccessor:
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

        self._indexes = {
            k: BoundsIndexer(v['left'], v['right'], self._obj.dims[k])
            for k, v in self._obj.attrs["bounds"].items()
        }

    def sel(self, **kwargs):
        assert len(kwargs) == 1, "Only one dimension can be selected"
        dim, value = kwargs.popitem()
        idx = self._indexes[dim].get_nearest_point(value)
        return self._obj.isel(**{dim: idx})

# use it
assert ds.bounds.sel(x=0.5).foo.values == 9

But what i want is to override Xarray's native .sel method with this index!

That is not possible with the custom indexes API, because xr.Index requires a class method called from_variables which generated an index from existing variables. In this case, there ARE no variables.

cc @benbovy @jhamman @dcherian

[dcherian]

#8473 (comment)
#7041

https://notebooksharing.space/view/37a45e01bbfe69d4242808ef9e88bf3dbcbd297d0221493fc4a4ae77e980e809#displayOptions=

I think the fundamental requirement is a variable of some sort. It seems like a scalar value can be used.

[rabernat]

Doh, yes my question is basically a dupe of #8473

[rabernat]

But what if the bounds are 2D? And defined by a 6-component affine transform matrix. (This is the actual gdal-style use case.)

The requirement that a variable exists feels arbitrary and limiting.

[dcherian]

In the Xarray data model, you'd want a Variable to propagate these attributes around in any case. Dataset attrs are lost when you extract a DataArray for example.

So it needs to be on the DataArray, and then to limit duplication, you'd instead land up with cf-style scalar coordinate variable (e.g. x_bounds) with all the info, and an attribute on each dataarray linking that array to the scalar coordinate variable (eg. geospatial_x_coordinate or something). You can set an index on x_bounds then, and use .sel(x_bounds=...)

This is the idea behing spatial_ref for example.

[benbovy]

Yes spatial_ref is a good example on how to handle this case I think.

While it would be technically feasible to allow Xarray indexes without any corresponding coordinate (but at least related to one or more dimensions), it makes the Xarray model slightly more complex probably with more complications in the API and implementation details.

[benbovy]

Why not having a variable wrapping a lazy array (e.g., a dask array or any other duck array) for the range index? This may never get materialized into a discrete array fully in memory, but this could also be done if needed.

This is almost possible today (needs work in #8124).

It requires that the variable exists before the index is created, but this could be done for example by a custom method (in an accessor) that creates both the coordinate and the index.

[rabernat]

Because I don't want Dask involved in this (there are no chunks) and I don't want to even pretend that there is an array there.

[benbovy]

Makes sense for Dask. But having an array (lazy / virtual or not) still makes sense, unless you have no (array-based) data variable you want to select along the indexed dimension(s)?

[rabernat]

Thanks for the help and feedback. I realize this is a bit redudant with examples that @benbovy and others have already shared. However, if we want this capability to be widely adopted, we need the API to be easy to use. So hopefuly me sharing my struggles is productive! 🙃

Here's something I got working

import xarray as xr
import numpy as np
from xarray.indexes import Index
from xarray.core.indexes import IndexSelResult


@dataclass(frozen=True)
class BoundsIndexer:
    left: float
    right: float
    npoints: int

    def get_nearest_point(self, value):
        if value < self.left or value > self.right:
            raise ValueError("Value outside of bounds")
        return int((value - self.left) / (self.right - self.left) * ( self.npoints - 1))
    

@dataclass(frozen=True)
class BoundsIndex(Index):
    
    indexes: dict[str, BoundsIndexer]
    
    @classmethod
    def from_variables(cls, variables, *, options):
        # assume there is only one variable
        assert len(variables) == 1 
        # and that variable name is "bounds"
        bounds = variables["bounds"]
        indexes = {
            k: BoundsIndexer(v['left'], v['right'], v["npoints"])
            for k, v in bounds.attrs.items()
        }
        return cls(indexes)
    
    def sel(self, labels, method=None, tolerance=None):
        if method != "nearest":
            raise ValueError("Only nearest method is supported")
        
        # feels super redundant to have to do this
        bounds = labels["bounds"]

        results = {}
        for dim, q in bounds.items():
            results[dim] = self.indexes[dim].get_nearest_point(q)

        return IndexSelResult(results)

Now I can create a dataset like this

ds = xr.Dataset(
    data_vars={"foo": (("y", "x"), np.arange(npoints * npoints).reshape(npoints, npoints))},
    coords={"bounds":
        ((), (), {
            "x": {"left": 0.0, "right": 1.0, "npoints": npoints},
            "y": {"left": 0.0, "right": 2.0, "npoints": npoints}
        })
    }
).set_xindex("bounds", BoundsIndex)

which looks like this

<xarray.Dataset>
Dimensions:  (y: 20, x: 20)
Coordinates:
  * bounds   object ()
Dimensions without coordinates: y, x
Data variables:
    foo      (y, x) int64 0 1 2 3 4 5 6 7 8 ... 392 393 394 395 396 397 398 399
Indexes:
    bounds   BoundsIndex

And I can query it like this

ds.sel(bounds={'x': 0.5, 'y': 0.2}, method="nearest")

What I like about this

• All the coordinate information is stored in metadata, no explicit array
• There is a single variable where this information lives. This is important in the case of more complex affine transforms where x and y coords and not orthogonal to the dataset dimensions

What feels weird

• I don't like having to write ds.sel(bounds={'x': 0.5, 'y': 0.2}. I'd much rather write ds.sel(x=0.5, y=0.2). Is that possible somehow?
• I don't like having to put "npoints" in the attrs. It's redundant (and therefore possibly inconsistent.) We already have a dimension called x on the dataset. I just want to use the size of that dimension to create my index. But that information is not available from within from_variables, because there is no variable called x!

In both cases, it seems like things could be improved by somehow making indexes aware of / associated with dimensions, rather than only variables.

---

I'm also coming around to the idea that there could be lazily materialized x and y arrays for this dataset.

The other thing we will need is some sort of plugin / entry-point mechanism for encoding / decoding so that we can make these indexes created automatically upon loading conforming datasets.

[Illviljan]

• I don't like having to write ds.sel(bounds={'x': 0.5, 'y': 0.2}. I'd much rather write ds.sel(x=0.5, y=0.2). Is that possible somehow?

xarray tries to support Hashables for better or worse. Replace "x" with 3 and "y" with 200 then ds.sel(3=0.5, 200=0.2) isn't valid python code anymore.

[rabernat]

I think that's a red herring. That syntax already works today, so nothing new. The point is I don't want to wrap it all inside bounds; I just want to refer to x directly.

[keewis]

isn't valid python code anymore.

note that technically you can work around that limitation:

ds.sel(**{3: 0.5, 200: 0.2})

works. I.e. it is only a limitation of the syntax, not in the way parameters are received by a function.

I agree, though, that API-wise you'd want to be able to select along the "virtual" x and y coordinates directly.

[benbovy]

I think a good starting point for supporting indexes associated with dimensions might be to update the Xarray data model like so:

1. An Xarray index shall be associated with at least one coordinate or one dimension. It may be associated with either several coordinates or several dimensions but it cannot be associated with a mix of both coordinates and dimensions
2. A coordinate shall be associated with either zero or one index. A dimension shall be associated directly with either zero or one index ; a dimension may be indirectly associated with several indexes via distinct (sets of) coordinates ; a dimension cannot be associated with two indexes both directly and indirectly via a dimension coordinate (i.e., an Xarray Dataset or DataArray cannot contain both an index associated with a dimension "x" and another index associated with the dimension coordinate "x")

The restriction in (1) prevents unnecessarily complex API: an Xarray Index subclass shall implement either Index.from_variables(variables, * options) or Index.from_dimensions(dim_sizes, *, options). Or maybe it may implement both, but only one of these two methods will be called when setting a new index object.

The restrictions in (2) prevent ambiguities that would otherwise arise in ds.sel(x=...), ds.drop_indexes("x"), ds.set_xindex("x", ...), etc.

[benbovy]

We could probably make things easier with harder restrictions:

• If a Dataset / DataArray has an index associated with one or more dimensions, then it cannot have dimension coordinates for any of those dimensions.

This removes ambiguities with ds.xindexes["x"], where it would be otherwise hard to figure out whether the index is associated with the "x" dimension or the "x" coordinate.

[rabernat]

Another point of reference for storing coordinate information without a corresponding variable: OME-Zarr Coordinate Transformations.

It would be amazing if we could open OME-Zarr in Xarray and expose useful indexes!

[rabernat]

Here's a very concrete, practical task that could be used to advance this issue:

Try to modify RioXarray's generate_spatial_coords to create a custom Range index from the affine parameters instead of a explicit x, y coordinate variables.

If we can make this work well today with no changes to Xarray, that tells us that the current API is up to the job. If not, it will suggest where we need to make changes. I believe that the crux will come down to the issue around whether we can associate custom indexes with dimensions, rather than coordinates. (See #8955 (reply in thread)).

[benbovy]

You can modify generate_spatial_coords() so that it returns a xr.Coordinates object with a Range index and (possibly lazy) x, y coordinates. This would work pretty well today with no changes to Xarray.

For generate_spatial_coords() to return a Range index with no x, y coordinates, this would require some changes as detailed in
#8955 (reply in thread). I have the feeling that supporting Xarray indexes associated with dimension could be pretty easy to implement, but I might be wrong.

[rabernat]

Let's try the easy thing first then.

[arthur-st]

Would like to chime in with a use case from ionospheric physics, in case that's useful. I've worked with some radio datasets that are sampled at, e.g., 20 MHz, meaning that what one is working with is a billions-length 1D dataset that's float64 or complex128. Given start time t_0, it's possible to compute time t_i = t_0 + sampling_period * idx_i for any position, and so a space-efficient way to record this for .sel(time=slice(t_x, t_y) would just solve the formula for idx_i.