Flexible indexes refactoring notes

[benbovy]

As a preliminary step before I take on the refactoring and implementation of flexible indexes in Xarray for the next few months, I reviewed the status of https://github.com/pydata/xarray/projects/1 and started compiling partially implemented or planned changes, thoughts, etc. into a single document that may serve as a basis for further discussion and implementation work.

It's still very much work in progress (I will update it regularly in the forthcoming days) and it is very open to discussion (we can use this PR for that)!

I'm not sure if Xarray's root folder is a good place for this document, though. We could move this into a new repository in xarray-contrib (that could also host other enhancement proposals) if that's necessary.

I'm looking forward to getting started on this and to getting your thoughts/feedback!

[keewis]

this is a good summary, @benbovy.

One thing I'm missing is duckarray support, though. Not sure if this is realistic, but I'm hoping to reduce the maintenance burden on duckarray support libraries (such as pint-xarray) as much as possible: subclassing every new index class (or having the index provider explicitly add duckarray support) seems a bit too much work.

place: I think moving it into the documentation should be fine but a separate repository would probably be a bit cleaner (especially if we expect to write more design documents).

[shoyer]

This looks great, thanks for writing this up!

I would suggest keeping this in the xarray repository rather than letting it get lost somewhere else (it is very much about xarray's code, after all). Perhaps we could make a new sub-directory design_notes?

[benbovy]

Thanks for your comments @keewis and @shoyer!

I think it's better for now to keep having this discussion and these notes in the Xarray repository, for more visibility. We could still move this elsewhere later if this PR becomes too cluttered, as there are potentially many aspects we can discuss about.

[benbovy]

One thing I'm missing is duckarray support, though. Not sure if this is realistic, but I'm hoping to reduce the maintenance burden on duckarray support libraries (such as pint-xarray) as much as possible: subclassing every new index class (or having the index provider explicitly add duckarray support) seems a bit too much work.

I haven't looked much at pint-xarray yet, so I'm not sure to understand. Why would you need to subclass every new index class?

If you are referring to the issue that you describe in your comment #525 (comment), the refactoring should decouple indexes from the coordinates, leaving the latter "just" as if they were regular variables (thus with duckarray support). What is currently possible with non-index coordinates should be possible with all coordinates. Actually, I'm not sure that we'll need to keep IndexVariable after the refactoring.

Or maybe you're referring to unit-aware indexing (what @shoyer mentioned in #525 (comment))? In this case I'm not sure how we could do that without having specific index classes for that purpose. Maybe some pre/post indexing hooks in Xarray that could be used, e.g., to convert indexer units into the coordinate units?

[keewis]

I was referring to both, but I agree that the first point would be fixed by the decoupling. The second point is just as important, though. We're currently adding wrapper methods to pint-xarray which handle the unit-awareness, but I was hoping we could change the API so those would not be necessary. For the implementation: hooks would probably work. Other options might be decorator functions or context managers?

[jthielen]

It's great to be able to follow along with the discussion here! I'm definitely interested in seeing where the duck array index support ends up.

One use-case motivated question: the flexible indexes refactoring has also been pointed to as the resolution to #2233, where multidimensional coordinates have the same name as one of their dimensions. I wasn't quite able to tell through the narrative here if that has been addressed along the way yet or not ("A. only 1D coordinates with a name matching their dimension name" for implicit index creation does seem to get close though). So, would it be worth directly addressing #2233 here, or should that wait?

[benbovy]

For the implementation: hooks would probably work. Other options might be decorator functions or context managers?

Or similarly to _repr_inline_:

class MyDuckArray:
    ...

    def _sel_(self, indexer):
        """Prepare the label-based indexer to conform to this coordinate array."""
        ...
        return new_indexer

    ...

[benbovy]

One use-case motivated question: the flexible indexes refactoring has also been pointed to as the resolution to #2233, where multidimensional coordinates have the same name as one of their dimensions. I wasn't quite able to tell through the narrative here if that has been addressed along the way yet or not ("A. only 1D coordinates with a name matching their dimension name" for implicit index creation does seem to get close though). So, would it be worth directly addressing #2233 here, or should that wait?

I think #2233 will be addressed by the index refactoring here. I don't see any issue with multidimensional coordinates having the same name as one of their dimensions once indexes are decoupled from dimensions/coordinates. I might still be missing something, though.

[mathause]

One thing I did not see discussed is alignment (or did I miss this?). Currently alignment is based on the "indexes" or well I guess pd.Index(dimension-coords). I assume not all potential indexes are useful for this? Alignment will still be based on 1D dimension-coords?

One (potential) edge case are MultiIndex which allow to align several non-dimension coordinates.

import xarray as xr
da1 = xr.DataArray([1, 2, 3], dims="x", coords=dict(time=("x", [1, 2, 3]), exp=("x", ["a", "a", "b"])))
da2 = xr.DataArray([2, 3, 4], dims="x", coords=dict(time=("x", [1, 2, 3]), exp=("x", ["a", "a", "a"])))
da1 = da1.set_index(x=("time", "exp"))
da2 = da2.set_index(x=("time", "exp"))

da1 + da2

<xarray.DataArray (x: 2)>
array([3, 5])
Coordinates:
  * x        (x) MultiIndex
  - time     (x) int64 1 2
  - exp      (x) object 'a' 'a'

[benbovy]

Alignment hasn't been discussed yet here, but it should! Some quick thoughts:

• support for alignment should probably be optional for an Xarray index wrapper.
  • like pandas.Index, the index wrapper classes that support it should implement .equals(), .union() and/or .intersection()
  • support might be partial if that makes sense (outer, inner, left, right, exact...).
  • index equality might involve more than just the labels, like the CRSIndex proposed in Checking non-dimensional coordinates for equality #2996
  • some indexes might implement inexact alignment, like in Alignment with tolerance2 #4489 or a KDTree index that selects nearest-neighbors within a given tolerance
  • alignment may be "multi-dimensional", i.e., the KDTree example above vs. dimensions aligned independently of each other
• we need to decide what to do when one dimension has more than one index that supports alignment
  • we should probably raise unless the user explicitly specify which index to use for the alignment
• we need to decide what to do when one dimension has one or more index(es) but none support alignment
  • either we raise or we fail back (silently) to alignment based on dimension size
• for inexact alignment, the tolerance threshold might be given when building the index and/or when performing the alignment
• are there cases where we want a specific index to perform alignment and another index to perform selection? It would be tricky to support that unless we allow multiple indexes per coordinate. "Meta" indexes (Flexible indexes refactoring notes #4979 (comment)) would help but then I'm worried about the possible explosion of index wrapper classes.

[dcherian]

There are also high-level methods that could use indexes in non-trivial ways. These methods become complicated when considering nD versions of IntervalIndex (or CellIndex that represents the appropriate bounds variable, #1475). I'm raising these so we can define the scope of the API.

1. resample (CFTimeIndex and a DatetimeIntervalIndex)
2. DatetimeAccessor & TimedeltaAccessor properties (CFTimeIndex and a DatetimeIntervalIndex)
3. interp & interpolate_na,
   • with IntervalIndex, these become regridding operations. Should we support hooks for these operations?
4. differentiate, integrate, polyfit
   • raise an error if not a "simple" 1D index?
5. pad
6. coarsen has to make choices about output index labels.
7. sortby
8. plotting
   1. plot.pcolormesh "infers" interval breaks along axes, which are really inferred bounds for the appropriate indexes.
   2. plot.step again uses bounds. In fact, we may even want step to be the default 1D plotting function if the axis has bounds attached.

Another perhaps-unintended use-case is that various accessors will try to use .indexes to cache state (#3268 (comment)).

An example is CRSIndex, where the crs applies to all variables in a dataset but isn't associated with a dimension on any of the other variables (Maybe this is a good definition for a PropertyIndex). This would be used by rioxarray, salem and friends. Also consider a possible XgcmIndex that basically contains an Xgcm.Grid object (also satisfies the PropertyIndex definition). We would want to propagate such PropertyIndexes in all operations, and also optionally check them during alignment (?).

[benbovy]

There are also high-level methods that could use indexes in non-trivial ways.

Thanks @dcherian for listing those methods here, that's something worth to keep in mind! I think that for now it would be reasonable to restrict those methods to the indexes that are currently available in Xarray instead of trying to extend the API of Xarray index wrappers in order to support those special cases. I guess it's ok for "default" or "common" xarray indexes to provide extra functionality that could not be implemented in 3rd party indexes, as well as it would be ok for 3rd-party indexes to provide non-standard, extra functionality that would be reused for methods implemented in DataArray/Dataset accessors.

Maybe this is a good definition for a PropertyIndex

Are the CRSIndex and XgcmIndex examples really independent of any coordinate in the DataArray/Dataset? Looks like in #2996 a CRSIndex could be bound to x and y coordinates and a XgcmIndex could be bound to x, y, x_c, y_c, face, etc. coordinates?

[benbovy]

Are the CRSIndex and XgcmIndex examples really independent of any coordinate in the DataArray/Dataset? Looks like in #2996 a CRSIndex could be bound to x and y coordinates and a XgcmIndex could be bound to x, y, x_c, y_c, face, etc. coordinates?

@rabernat @jbusecke (xgcm) @snowman2 @fmaussion @djhoese (crs) it would be interesting to have your thoughts here.

What would be the pros and cons of:

• Refactoring xgcm.Grid into a XgcmGridIndex? The index would typically be assigned to the Dataset coordinates that are also specified in the coords argument of xgcm.Grid and all other arguments (except the Dataset itself) would become index options. xgcm.Grid methods would then be accessible via Dataset accessor(s) (or eventually just replaced by xarray's corresponding methods).

• Refactoring a crs attribute (either a "public" Dataset/DataArray attribute or hidden behind an accessor) into some CRSIndex that would typically be assigned to x/y or lat/lon coordinates?

A major advantage is that using a custom index, there's no need to encapsulate a Dataset/DataArray into a higher level structure (e.g., xgcm.Grid) and there would be more control on how it is propagated from one xarray object to another compared to an attribute or via a "stateful" accessor (e.g., crs). Another advantage is that Xarray selection and/or alignment can be customized. But that can be also a downside: unless we allow multiple indexes per coordinate, such XgcmGridIndex and CRSIndex would then have the responsibility of handling selection and alignment for all their corresponding coordinates. That may not be a big deal, though: XgcmGridIndex and CRSIndex could simply encapsulate pandas.Index instances for all (or a subset) of their coordinates.

Are there any other challenges and/or opportunities? (sorry, it has probably been already discussed elsewhere. There's too many places to look for :-) ).

[snowman2]

For reference for how rioxarray does things: https://corteva.github.io/rioxarray/stable/getting_started/crs_management.html

>>> import xarray, rioxarray
>>> xda = xarray.DataArray(1)
>>> xda.rio.write_crs(4326, inplace=True)
<xarray.DataArray ()>
array(1)
Coordinates:
    spatial_ref  int64 0
Attributes:
    grid_mapping:  spatial_ref
>>> xda.spatial_ref
<xarray.DataArray 'spatial_ref' ()>
array(0)
Coordinates:
    spatial_ref  int64 0
Attributes:
    crs_wkt:                      GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["...
    semi_major_axis:              6378137.0
    semi_minor_axis:              6356752.314245179
    inverse_flattening:           298.257223563
    reference_ellipsoid_name:     WGS 84
    longitude_of_prime_meridian:  0.0
    prime_meridian_name:          Greenwich
    geographic_crs_name:          WGS 84
    grid_mapping_name:            latitude_longitude
    spatial_ref:                  GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["...

[shoyer]

A major advantage is that using a custom index, there's no need to encapsulate a Dataset/DataArray into a higher level structure (e.g., xgcm.Grid) and there would be more control on how it is propagated from one xarray object to another compared to an attribute or via a "stateful" accessor (e.g., crs)

I agree, this would be really nice.

One challenge is that often it is not advisable to explicitly build the coordinate arrays that correspond to such grids For example, consider a satellite image: 2D lat/lon arrays could be as expensive to store as the image itself, even though the values can be computed on the fly with very cheap arithmetic.

To fill this gap, we need first-class support for custom lazy arrays in xarray. If you read the documentation for the backend refactor (#4810), you'll see that we do have a minimal version of this internally in the form of core.indexing.LazilyIndexedArray. Ideally we would not only expose this functionality publicly, but would even factor it out into a separate lazy "duck array" library that could be used independently of xarray.

[benbovy]

Ah yes, making it more easily reusable would be welcome indeed. I guess that such lazy arrays will be already needed for the creation of coordinates from the levels of an existing pandas.MultiIndex.

[djhoese]

2D lat/lon arrays could be as expensive to store as the image itself, even though the values can be computed on the fly with very cheap arithmetic.

Just wanted to mention in case it comes up later, this is true for some datasets and for others the lon/lats are not uniformly spaced so they can't be calculated (just based on the way the satellite instrument works). They have to be loaded from the original dataset (on-disk file). For a while in the Satpy library we were storing 2D dask arrays for the lon/lat coordinates until we realized xarray was sometimes computing them and we didn't want that.

[dcherian]

I think that for now it would be reasonable to restrict those methods to the indexes that are currently available in Xarray instead of trying to extend the API of Xarray index wrappers in order to support those special cases.

Fully agree. We should raise nice error messages when possible. I just wanted to raise awareness about this issue (i.e. methods that use indexes in non-trivial ways).

Are the CRSIndex and XgcmIndex examples really independent of any coordinate in the DataArray/Dataset?

Good point! I hadn't thought of it that way.

A major advantage is that using a custom index, there's no need to encapsulate a Dataset/DataArray into a higher level structure (e.g., xgcm.Grid) and there would be more control on how it is propagated from one xarray object to another compared to an attribute or via a "stateful" accessor (e.g., crs)

Yes! I'm v. happy to see this discussion is happening :)

But that can be also a downside: unless we allow multiple indexes per coordinate, such XgcmGridIndex and CRSIndex would then have the responsibility of handling selection and alignment for all their corresponding coordinates. That may not be a big deal, though: XgcmGridIndex and CRSIndex could simply encapsulate pandas.Index instances for all (or a subset) of their coordinates.

For XgcmIndex this is even desirable: xgcm/xgcm#200, Re CRSIndex I think you're correct, it could just forward to whatever Index class it's wrapping

[benbovy]

For reference for how rioxarray does things: https://corteva.github.io/rioxarray/stable/getting_started/crs_management.html

That's good to know, thanks! Like it may create specific indexes for time coordinates, I could imagine Xarray's decode_cf(decode_coords=True) to eventually return some kind of CRSIndex from any variable referred to in grid_mapping attribute.

[benbovy]

Thanks everyone for your comments so far!! They have been really helpful in improving the notes!

This is now ready for another round of review. I've tried to include all the points raised in the discussion above. I also marked all the conversations as resolved even though it's still open for discussion! (it's just a way to "reset" them for more clarity). I'll move the notes into a design_notes folder just before merging this PR.

With the last commits, I think that the notes now cover most of the aspects regarding the use of indexes in Xarray. The goal with these notes is not to settle every detail of the refactoring (decisions can be made while iterating on the implementation), but rather describe the big picture and outline the main opportunities and challenges. Referring to the notes will help throughout the implementation. Hopefully it will allow more Xarray users and devs sharing their point of views to make sure we're not missing anything important here.

One thing that is not in the notes: to which acceptable extent this refactoring may introduce breaking changes? I think that it will be hard to avoid any breaking change. That said, as the index refactoring would rather bring internal data structures to the light I don't expect many things to break (at least, not the things that 90% of Xarray users often rely on). Hardest part will probably be to ensure a smooth transition while updating the API that is too specific to pandas.MultiIndex into something that is more index-agnostic...

[shoyer]

One thing that is not in the notes: to which acceptable extent this refactoring may introduce breaking changes? I think that it will be hard to avoid any breaking change. That said, as the index refactoring would rather bring internal data structures to the light I don't expect many things to break (at least, not the things that 90% of Xarray users often rely on). Hardest part will probably be to ensure a smooth transition while updating the API that is too specific to pandas.MultiIndex into something that is more index-agnostic...

My opinion is that breaking changes are OK if done with care. The main thing to keep in mind is that users dealing with change typically will not directly benefit from it.

So any breaking changes need to be intentional, and happen all at once in a major release. We also need to issue warnings when possible.

Depending on the magnitude of these changes, this could be a good reason to declare "xarray 1.0"

[shoyer]

+1

[rabernat]

This is really great @benbovy! My comments are mostly just clarification questions and small rephrasing suggestions. The discussion above caught most of the big stuff.

[rabernat]

Here and elsewhere it sounds like it could be useful to define an index precedence order.

[benbovy]

Yes this could be a solution to all situations where the index to pick is ambiguous (although it is not yet clear to me how often such situation could happen). This would introduce more complexity, though:

• we need to extend the API so that users can define this order
• we need to ensure strict precedence order
• does it make sense to use the same precedence order for all operations? (e.g., alignment vs. selection)

[rabernat]

I think indexes should optionally support interpolation.

[benbovy]

That would certainly be useful.

[benbovy]

Thanks for your reviews @shoyer and @rabernat! I've updated the notes according to your comments.

From my point of view this is ready, but we can leave this PR open for another few days in case anyone else wants to add some comments (@pydata/xarray).

Next week I'll start with the implementation (I'll take on the open PR in https://github.com/pydata/xarray/projects/1 and will do only internal refactoring making sure that all tests are passing).

[dcherian]

shall we merge?

[benbovy]

shall we merge?

Yes! I wanted to wait for the bi-weekly dev meeting but I've just missed it (PST/PDT -> 🤦 -> sorry).

Let's merge this and continue the discussion in follow-up issues/PRs.