Meta-indexes with >2 dimensions

[wpbonelli]

Hi, I've recently been poking around in the custom index stuff, and found some things to ask about.

Say you have a dataset representing a structured grid with some data living on grid cells. Dimensions x/y are natural, but say you don't have any explicitly spatial labels at the moment and want to be able to index rows/cols with i/j. You could do

ds = xr.Dataset(
    {
        "a": (("rows", "cols"), np.zeros((3, 3))),
        "b": (("cells",), np.ones((9,)))
    },
    coords={
        "i": ("rows", [0, 1, 2]),
        "j": ("cols", [0, 1, 2]),
        "n": ("cells", np.arange(9)),
    },
)

ds = ds\
    .set_xindex("i", PandasIndex)\
    .set_xindex("j", PandasIndex)\
    .set_xindex("n", PandasIndex)

ds
>>> <xarray.Dataset> Size: 264B
Dimensions:  (x: 3, y: 3, cells: 9)
Coordinates:
  * i        (x) int64 24B 0 1 2
  * j        (y) int64 24B 0 1 2
  * n        (cells) int64 72B 0 1 2 3 4 5 6 7 8
Dimensions without coordinates: x, y, cells
Data variables:
    a        (x, y) float64 72B 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    b        (cells) float64 72B 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0

print(ds.a.sel(i=0))
>>> <xarray.DataArray 'a' (y: 3)> Size: 24B
array([0., 0., 0.])
Coordinates:
    i        int64 8B 0
  * j        (y) int64 24B 0 1 2
Dimensions without coordinates: y

print(ds.b.sel(n=0))
>>> <xarray.DataArray 'b' ()> Size: 8B
array(1.)
Coordinates:
    n        int64 8B 0

You can write a meta-index to do this too (lifted from docs with some asserts removed)

class MetaIndex(Index):
    def __init__(self, indices):
        self._indices = indices

    @classmethod
    def from_variables(cls, variables):
        return {k: PandasIndex.from_variables({k: v}) for k, v in variables.items()}

    def create_variables(self, variables=None):
        idx_vars = {}
        for index in self._indices.values():
            idx_vars.update(index.create_variables(variables))
        return idx_vars

    def sel(self, labels):
        results = []
        for k, index in self._indices.items():
            if k in labels:
                results.append(index.sel({k: labels[k]}))
        return merge_sel_results(results)

As suggested by @benbovy here you can do

def alias(dataset: xr.Dataset, old_name: str, new_name: str) -> PandasIndex:
    size = dataset.sizes[old_name]
    return PandasIndex(pd.RangeIndex(size, name=new_name), dim=old_name)

ds = ds.drop_indexes(["i", "j", "n"])
ds = ds.assign_coords(xr.Coordinates.from_xindex(MetaIndex(
    {
        "i": alias(ds, "x", "i"),
        "j": alias(ds, "y", "j"),
    }
)))

And it works.

ds.a.sel(i=0)
<xarray.DataArray 'a' (y: 3)> Size: 24B
array([0., 0., 0.])
Coordinates:
    i        int64 8B 0
    j        (y) int64 24B 0 1 2
Dimensions without coordinates: y

Note the absence of the star before the j coordinate in the result, indicating that i and j are not associated with an index? But it looks like the index is there

ds.a._indexes
>>> {'i': <__main__.MetaIndex object at 0x148892850>, 'j': <__main__.MetaIndex object at 0x148892850>}

Now, if I pass not 2 but 3 indexes to MetaIndex, it breaks down

ds = ds.drop_indexes(["i", "j"])
ds = ds.assign_coords(xr.Coordinates.from_xindex(MetaIndex(
    {
        "i": alias(ds, "x", "i"),
        "j": alias(ds, "y", "j"),
        "n": alias(ds, "cells", "n"),
    }
)))

print(ds.a.sel(i=0))
>>> Exception has occurred: KeyError
"no index found for coordinate 'i'"

print(ds.b.sel(n=0))
>>> Exception has occurred: KeyError
"no index found for coordinate 'n'"

But if I remove the first two entries in the dict argument to MetaIndex, and alias only n for cells it works again. It seems like there is a constraint that a meta-index can have at most 2 indexes internally? Or maybe I'm misdiagnosing.

I also saw this and wondered in general how to implement to_pandas_index() for a meta-index having multiple dimensions. Without to_pandas_index I found that the .indexes property crashes so I have been inspecting ._indexes instead.

I know this is all under development. I am excited to take advantage of it. Just trying to wrap my head around it. Thanks!

[benbovy]

Note the absence of the star before the j coordinate in the result, indicating that i and j are not associated with an index? But it looks like the index is there

The index is still in the input object but normally it is not in the result ds.a.sel(i=0)._indexes.

By default the index is dropped after the selection unless MetaIndex.isel(...) is implemented such that it returns a new index instance.

if I pass not 2 but 3 indexes to MetaIndex, it breaks down

It breaks down in your example because the selection is done on the DataArrays ds.a and ds.b and not on the Dataset ds. Each of those DataArrays have only a subset of the dimensions linked to the meta-index, which is by default dropped (not all index coordinates are propagated to those DataArrays). #10137 will allow overriding this behavior.

Without to_pandas_index I found that the .indexes property crashes so I have been inspecting ._indexes instead.

.indexes is the legacy property that returns a mapping of pandas.Index objects, whereas ._indexes returns a mapping of xarray.Index objects, although it is preferable to access the latter via the .xindexes public property.

A possible implementation of MetaIndex.to_pandas_index() would be to check if there is only one wrapped index and convert it. But note that this method is only needed for using the custom index with Xarray features that haven't been refactored yet for using xarray.Index objects instead of pandas.Index objects (Xarray's most used features have already been refactored).

[wpbonelli]

Thank you @benbovy just to make sure I understand, if I am reading filter_indexes_from_coords correctly, a dataset does not propagate index-linked coordinates to a dataarray unless for all coordinates in the index the coordinate dimensions are a subset of the array dimensions. So in this example, the index with i/j/n matches neither array a nor b, so none of its coordinates get through.

And after #10137 I could override should_add_coord_in_dataarray to get this to work. But maybe the lesson here is to use separate indexes for dimensions/coords that never appear together in arrays. Generally, it seems better to keep indexes separate where possible, than to use a meta-index, unless combining them is "natural" in some way like for rasters or CRS. Would you agree?

Thank you for the heads up about .xindexes etc. And about to_pandas_index. So the latter is temporary, it seems.

[benbovy]

Yes, all true!

to_pandas_index will likely stay for a while, though.