Custom index examples

[benbovy]

Let's gather here examples (prototypes) of custom indexes that have been proposed so far. This is all very much WIP, some examples may rely on features that have not yet landed in the main branch while other examples may be outdated. Anyway, it should provide a good basis for more experiments with Xarray indexes. To @pydata/xarray or anyone: feel free to edit this comment and add to the list below.

• default indexes (for reference): PandasIndex and PandasMultiIndex
• Helper class to implement a custom meta-index based on pandas index(es): notebook
• basic array index: notebook
• kd-tree index: notebook
• unit-aware (pint) index: implement the PintMetaIndex xarray-contrib/pint-xarray#163
• periodic boundary index: Periodic Boundary Index #7031
• interval index: SciPy Sprint: Creating Flexible Indexes #6783 (comment)
• functional index: notebook
• CRS index (raster 2d): notebook and Experimentally support CRSIndex corteva/rioxarray#588
• out-of-core (dask) index: Low memory/out-of-core index? #1650 / Supporting out-of-core computation/indexing for very large indexes #1094 / notebook
• Staggered grid (xGCM) index: https://github.com/dcherian/xindexes/blob/main/xgcmindex.ipynb; accessors through encapsulation xgcm/xgcm#225 (comment)
• Shapely (2.0) geometry index: https://github.com/xarray-contrib/xvec - notebook
• World Coordinate System (WCS) index (astropy): notebook, Create a to/from xarray method/function sunpy/ndcube#222
• DGGS (Discrete Global Grid System) indexes: https://github.com/xarray-contrib/xdggs
• Regular or range index: notebook, Regular (linspace) Coordinates/Index #8473

[SimonHeybrock]

1. Not sure if this would work implicitly via the PintMetaIndex you linked to, but maybe worth mentioning: We would like something like pandas.IntervalIndex but with units.

2. In the field I work in the automatic alignment that indexes provide (such as automatically returning the intersection from da1 + da2 with an index-coord for 'x') is not really required. Instead it may create (some) overhead, or risk running long computations where the user would actually just have preferred getting an early exception. Therefore, I am playing with the idea of having a StrictIndex (probably bad name, just to have something to refer to) which:

   • Is a very light wrapper around a NumPy array or duck array, as we want to be able to do math on this to compute derived coords.
   • Allows for label-based indexing (sel), but performance is not optimized for that as few queries are expected.
   • Raises an exception when trying to join with something that is not identical (length, labels, order of labels). Basically, da1 + da2 should just raise unless da1.coords['x'] is identical to da2.coords['x'] in the NumPy-array-equal sense.

[benbovy]

Early prototype of such an ArrayIndex in this notebook (I added it to the list above).

[SimonHeybrock]

Very nice! I gave this a try with a pint.Quantity (instead of a np.ndarray) and I got it to work with very minor changes, not specific to Pint. I can now do (no need for @keewis's PintMetaIndex in this case!):

da.sel(x=np.array(2.2) * pint.Unit('m'))

da.sel(x=2.2 * pint.Unit('m')) or da.sel(x=pint.Quantity(2.2, 'm')) do not work without further changes, since xr.core.utils.is_scalar returns False, but I presume this could be solved somehow.

[SimonHeybrock]

There is one subtlety I ran into: Using np.argmax(self.array == label) in ArrayIndex.sel, we will get index 0 on a units mismatch, but we'd actually like an exception. Adding a special check for Pint would not be so nice, and I was wondering if this is something applicable to other duck arrays as well? An ugly hack that works is to use

class ArrayIndex(Index):
  ...
  def sel(self, ...):
    _ = self.array[0] == label  # trigger exception if incompatible
    indexer = np.argmax(self.array == label)
    ...

Is there a better and generic way to ensure that a label is "compatible" with a certain index?

[benbovy]

Yes, when experimenting with this ArrayIndex I've also found the returned index 0 annoying in a few other cases too.

Is there a better and generic way to ensure that a label is "compatible" with a certain index?

There is a vague suggestion of a protocol here.

[keewis]

in this specific case I'd actually set force_ndarray_like = True on the application registry, which would convert the scalar to a 0d array (if you use pint-xarray that will happen automatically to fix some other issues)

[TomNicholas]

Spotted another custom index experiment in the wild by @Huite - unstructured grid indexes. (More context here)

[TomNicholas]

Histogram integration in xarray via flox would benefit from using pandas.IntervalIndex to represent both ends of each histogram bin. I think there are currently some bugs if you try to create an xarray object with an IntervalIndex.

[TomNicholas]

@benbovy is an IntervalIndex something which needs a wrapper similar to your MultiPandasIndex example? cc @dcherian

Helper class to implement a custom meta-index based on pandas index(es): notebook

[dcherian]

IntervalIndex seems a little hard to support, because it wraps a Pandas ExtensionArray. I tried some things here (sorry not well documented). One other thing to try could be a numpy structured array.

[benbovy]

I think one of the best approaches would be to have a custom Xarray IntervalIndex that only supports a unique 1-d coordinate.

That said, an extension array adaptor like the one at the end of @dcherian's notebook would be nice to have. As another example, geopandas uses extension arrays for the geometry columns so this could be useful for other cases like https://github.com/martinfleis/xvec.

Perhaps something like this could work for IntervalIndex:

• Expose a custom xarray.indexes.IntervalIndex that wraps a pd.IntervalIndex
  • Like for PandasIndex, its corresponding 1-d coordinate wraps the pandas index in a xarray.core.indexing.PandasIndexingAdapter
  • It could actually be implemented as a subclass of PandasIndex?
• How to use set_xindex?
  • obj.set_xindex("x", IntervalIndex) where "x" is a 1-d coordinate (if it wraps an pd.arrays.IntervalArray or a pd.IntervalIndex just reuse it, otherwise create a new pandas interval index with default bounds).
  • obj.set_xindex("x", IntervalIndex, use_bounds_coord="x_bounds") where use_bounds_coord is an option specific to the Xarray IntervalIndex. In this case the "x_bounds" coordinate remains unindexed, it could be manually removed after creating the index.
• Getting back the bounds coordinate, e.g., with something like obj.assign(x_bounds=obj.xindexes["x"].get_bounds_coord())