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.

EDIT: there's now an index gallery at https://xarray-indexes.readthedocs.io/ - https://github.com/xarray-contrib/xarray-indexes

• 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 PintIndex 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 vylor/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
• Forecast Model Run Collections (FMRC) - https://rolodex.readthedocs.io/forecast.html

[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())

[DamienIrving]

I'm very excited about the rolodex.ForecastIndex described at: https://xarray-indexes.readthedocs.io/earth/forecast.html

Apologies if I'm missing something very obvious, but how does one install rolodex in order to run import rolodex.forecast?

From the rolodex documentation page I can't find installation instructions or even a GitHub repo corresponding to that documentation page or a rolodex package?

[kmuehlbauer]

I think this is the related GitHub repo: https://github.com/dcherian/rolodex

[dcherian]

my goal is to move it to https://github.com/abkfenris/xarray_fmrc but haven't had time to do that yet.

[matthiasschabel]

Following up on a suggestion in a napari issue, I kind of took the idea and ran with it in an AI-fuelled binge, implementing xarrayrf : (https://github.com/matthiasschabel/xarrayrf). It's an xarray extension that attaches a reference frame to a DataArray. It is obviously alpha/prerelease, and the README and design docs are essentially unadulterated Claude-speak right now and need significant reworking. This example should give a flavor of what I'm trying to accomplish, which is to encapsulate reference frames and transformations in a clean and extensible way : https://github.com/matthiasschabel/xarrayrf/blob/main/examples/xarrayrf_tour.ipynb. It has currently got adapters for DICOM, NifTi, NGFF, and GeoTIFF.

Some features so far:

• Named spaces. Every transform declares its source and target. Array coordinates play the role of napari's data space, and reference frames carry an identity, so physical vs world (the NGFF distinction linked above) is an explicit frame-to-frame transform rather than a position in a chain.
• More than scale/translate. Transforms follow a small protocol (map points, with an optional inverse and Jacobian); affines and compositions are built in. Coordinates may be nonuniform, e.g. irregular slice positions.
• Dimension roles. The frame names its spatial dims. Time, channel and anything else are nonspatial by rule, and transforms are applied by axis name, not position.
• Slicing. The frame lives in an xarray index, so ordinary isel, sel and arithmetic either carry the geometry or refuse. A selected plane keeps its position as a scalar coordinate, which gives the rectangular data-to-world matrix discussed above; a plane from a sheared volume gets the correct per-slice offset without special handling.
• Canvas to data. Exact inversion where an inverse exists, refusal where it doesn't (a plane embedded in a volume). Invertibility is a capability, not a requirement.
• Resampling off the GPU. Explicit, lazy resampling onto another grid or frame.

Things not yet implemented but on the radar:

• per-timepoint or per-section transforms (the stabilization and EM alignment cases above)
• nonlinear transforms (deformable registration, GIS, astronomy, etc...)
• out-of-core oblique reslicing
• point-set layers