SciPy Sprint: Creating Flexible Indexes

[TomNicholas]

Using Flexible Indexes Sprint

I'm at SciPy right now, and running an xarray code Sprint right now!
It's in-person here but I hope we could involve people remotely too!

Times: 9am-6pm CDT July 16th and 17th

Sprint topic

I thought we could propose one focusing on the recent Flexible Indexes refactor - that's in a state where we're ready to try out making custom indexes.

Specific use case ideas

I have already found a few scipy attendees here who are interested in making various types of custom indexes, in particular:

• periodic index (think wrapping longitude around the earth) @dopplershift possibly?
• out-of-core index (chunked index for very long and thin data, a genomics guy here is interested @nvictus)
• pint index (@keewis you made a pint index somewhere didn't you?)
• KDTtree-like indexer for radar data that @mgrover1 had an idea for
• Staggered grid index as an alternative implementation of xGCM - I'm tempted to try this one (cc @rabernat and @jbusecke)

Links

Indexes refactor project which links to more suggestions https://github.com/pydata/xarray/projects/1

Any other ideas for indexes?

@pydata/xarray what do you think?

cc @scottyhq @JessicaS11

EDIT: @benbovy has come through and given us a temporary API to set custom indexes! 🚀

EDIT2: and @benbovy wrote some docs to help us! What a hero

EDIT3: Invite to shiny new discord! https://discord.gg/yduZSMYe

[dcherian]

Some index ideas:

• Add indexes to the repr
• IntervalIndex!
  • (it already works mostly -- assign pd.IntervalIndex to .indexes, and then .sel works ).
  • It'd be nice to add little helper things like .mid, .left, .right that return xarray objects.
  • Add support to .plot
• experiment with xoak some more

Another topic could be material for tutorial.xarray.dev (https://github.com/xarray-contrib/xarray-tutorial/issues)

• domain-specific examples/demos (even just linking out to existing material would be great).
• cloud data access tutorial
• other "data cleaning" examples.

Update docstrings with "Examples" and "See Also" sections: #5816 (comment)

[TomNicholas]

Great ideas! The tutorial / examples are good too for newer contributors I think.

[weiji14]

Thanks for the suggestions, have started a draft "cloud data access" tutorial at xarray-contrib/xarray-tutorial#132 😃

[keewis]

you made a pint index somewhere didn't you?

yes, here it is: xarray-contrib/pint-xarray#163 (see here for an example of how to try it out using the internal API, we don't have public API to set custom indexes yet).

[TomNicholas]

Messaged at the same time as you @benbovy !

There's still quite a lot to do before flexible indexes may be used "in-production"

So does this mean this sprint idea isn't viable? Or just that it will be tricker because there is less guidance available for someone trying to make a custom index right now?

[benbovy]

So does this mean this sprint idea isn't viable?

I think it is viable and it is a great idea!! (I was writing another comment :-))

Possible ways to quickly get some "public, temporary API" to play with custom indexes:

• A DataArray / Dataset alternative constructor that internaly calls the the constructor with fastpath=True (or maybe just use the constructor with fastpath=True?)

• Some DataArray / Dataset alternative set_index method with a simple signature (i.e., set one index for one or more coordinates)

[TomNicholas]

Thanks @benbovy !

(or maybe just use the constructor with fastpath=True?)

This seems pretty good, except that fastpath=True is going to skip checks that we ideally would like to run right? Otherwise our sprinters are more likely to get weird errors from invalid Dataset objects later on?

A DataArray / Dataset alternative constructor that internaly calls the the constructor with fastpath=True (or maybe just use the constructor with fastpath=True?)

Some DataArray / Dataset alternative set_index method with a simple signature (i.e., set one index for one or more coordinates)

I'm happy to try a PR to add either a new temporary constructor or a new set_index method later today (though I might need some help)

[benbovy]

A quick prototype (not fully tested):

import warnings

from xarray.core.indexes import Index 


def set_xindex(self, coord_names, index_cls, **kwargs):
    """Temporary API for creating and setting a new, custom index from
    existing coordinate(s).
    
    Parameters
    ----------
    coord_names : str or list
        Name(s) of the coordinate(s) used to build the index.
        If several names are given, their order matters.
    index_cls : class
        Xarray index subclass.
    **kwargs
        Options passed to the index constructor. Not working for now
        (not sure yet how to do it).
    
    """
    warnings.warn("This is temporary API to experiment with custom indexes")
    
    if not issubclass(index_cls, Index):
        raise TypeError(f"{index_cls} is not a subclass of xarray.core.indexes.Index")
    
    if isinstance(coord_names, str):
        coord_names = [coord_names]
    
    invalid_coords = set(coord_names) - self._coord_names
    
    if invalid_coords:
        raise ValueError(f"those coordinates don't exist in Dataset: {invalid_coords}")
    
    # we could be more clever here (e.g., drop-in index replacement if index
    # coordinates do not conflict), but let's not allow this for now
    indexed_coords = set(coord_names) & set(self._indexes)
    
    if indexed_coords:
        raise ValueError(f"those coordinates already have an index: {indexed_coords}")
    
    coord_vars = {k: self._variables[k] for k in coord_names}
    
    # note: extra checks (e.g., all coordinates must have the same dimension(s))
    # should be done in the implementation of Index.from_variables
    index = index_cls.from_variables(coord_vars)
    
    # in case there are index coordinate variable wrappers
    # (e.g., for PandasIndex we create coordinate variables that wrap pd.Index).
    # `coord_vars` passed as argument is for propagating coordinate metadata
    new_coord_vars = index.create_variables(coord_vars)
    
    # reorder variables and indexes so that coordinates having the same index
    # are next to each other
    variables = {}
    for k, v in self._variables.items():
        if k not in coord_names:
            variables[k] = v
    
    for k in coord_names:
        variables[k] = new_coord_vars.get(k, self._variables[k])
        
    indexes = {}
    for k, v in self._indexes.items():
        if k not in coord_names:
            indexes[k] = v
    
    for k in coord_names:
        indexes[k] = index
    
    return self._construct_direct(
        variables=variables,
        coord_names=self._coord_names,
        dims=self._dims,
        indexes=indexes
    )

Quick test:

# monkey patch
xr.Dataset.set_xindex = set_xindex

# dummy custom Index (which is just a pandas index)
from xarray.core.indexes import PandasIndex


class DummyIndex(PandasIndex):
    ...

# actually set a pandas index for a non-dimension coordinate
ds = xr.Dataset(coords={"x": ("y", [0, 1, 2])}).set_xindex("x", DummyIndex)

ds2.xindexes
# Indexes:
# x: <__main__.DummyIndex object at 0x16ecfbf40>

# sel works!
ds2.sel(x=[0, 1])
# <xarray.Dataset>
# Dimensions:  (y: 2)
# Coordinates:
#   * x        (y) int64 0 1
# Dimensions without coordinates: y
# Data variables:
#     *empty*

[TomNicholas]

Incredible! Thanks @benbovy ! I could presumably easily make a set_xindexes for DataArray from this too

[benbovy]

I'd love to see what would come out of a sprint on creating custom indexes! Too sad that I won't be much available this weekend (even remotely), but I'll try to see tomorrow if I can prepare some quick documentation or code that could be useful for the sprint.

periodic index (think wrapping longitude around the earth) @dopplershift possibly?

For the case of latitude / longitude, you could check https://github.com/benbovy/pys2index which allows wrapping queries around the Earth, although it currently only supports point-wise indexing. It is used in xoak but the latter is not yet reusing Xarray's flexible indexes.

FWIW, in addition to continuing some work on flexible indexes, I'll start (probably sometime during this summer) working on Python bindings for https://github.com/paleolimbot/s2geography (thanks to a NumFOCUS small development grant), which could then be reused as a backend engine for Xarray geographical indexes.

[fhk]

Would love to see - https://h3geo.org/

[fhk]

this is the closest I have found - https://github.com/joaofig/geo-spoke/blob/master/geo/geospoke.py

[TomNicholas]

@fhk that looks cool! I don't quite understand though - could you perhaps give a pseudocode idea of what you want to achieve in user code?

[Cadair]

Hello!

I am also at SciPy and am very interested in attempting you use the flexible indicies stuff to work with Astro functional coordinate transforms.

[martindurant]

Thanks for the mention. I'm not there, unfortunately, but happy to take part in the conversation.

[martindurant]

Also worth mentioning, that several other xarray-like file formats contain coordinate information that would be accessible with functional indexes. Mostly these are affine (nifti, geoTIFF...), but more general forms should be doable too.

Is a coordinate transform with cross-terms on the table? With static array coordinates, we have examples where variable V(point_index), and point_index->(x, y) . But astro coordinates commonly have quadratic and higher terms and cross-terms, because mirrors are bendy and detectors flat.

[benbovy]

xref this discussion (with some suggestions regarding the API): #3620

[Cadair]

My main goal is that you can provide an object with callables to convert from pixel to "world" coordinates etc, so the actual transform can be as complex as it needs be.

[benbovy]

Is a coordinate transform with cross-terms on the table? With static array coordinates, we have examples where variable V(point_index), and point_index->(x, y) . But astro coordinates commonly have quadratic and higher terms and cross-terms, because mirrors are bendy and detectors flat.

I'm not familiar with astro coordinates, but since an Xarray custom index could be defined for an arbitrary number of coordinates with each an arbitrary number of dimensions I guess this could be supported.

Custom indexes could also be coupled with Dataset / DataArray accessors, e.g., to generate new coordinates on the fly or extend the API that consumes those indexes (i.e., bypass coordinates and deal directly with the index), like suggested in #3620 (comment).

several other xarray-like file formats contain coordinate information that would be accessible with functional indexes.

This makes me think that we could connect flexible storage backends with flexible indexes. For example, a storage backend that directly creates the right index(es) from the information contained in a given file format.

Since indexes are now 1st class citizens of the Xarray data-model we can directly attach to them some information that is relevant to the end-user (e.g., CRS, WRS, etc.). Sometimes it might actually make more sense than trying to attach the same information as coordinate metadata.

[benbovy]

I've put together some preliminary documentation on custom indexes here: https://hackmd.io/Zxw_zCa7Rbynx_iJu6Y3LA?view

It's a very rough draft but I hope it will be helpful for the sprint!

EDIT: Do not hesitate to edit / comment the draft!

[TomNicholas]

This is amazing @benbovy. One big question I have is how should we test our indexes? Is there any potential for a common set of tests to run or will all the tests need to be specific to that particular Index type?

[benbovy]

I think that for "meta-indexes" there would likely be some common testing patterns, i.e., to ensure that everything is passed correctly to/from the encapsulated PandasIndex objects. This is another reason why a meta-index helper (base) class (with tests) would be a useful addition in Xarray.

For the other cases, I think that the tests will be more specific to each custom index. There's a lot of flexibility, e.g., for the results of label-based selection it is possible to return new indexes, new variables, drop variables, etc., so it might be hard to write common tests that would work for all the indexes.

[keewis]

Is there any interest in making this a hybrid event? I'd be able to help out tomorrow in that case.

[TomNicholas]

I thought I had replied to this but perhaps my email didn't send... I was going to suggest emailing the user mailing list but I think it might be too later for that now.

Do we have a suitable discord room or whatever to use for remote participants?

[keewis]

I've created the scipy22 branch with the first commit being the temporary API from #6783 (reply in thread).

[keewis]

It's probably still worth announcing, but my impression is that the mailing list does not reach enough people. Maybe something else as well. Twitter, maybe?

[TomNicholas]

Ohh yeah twitter! Let's do that now (I don't have access to the twitter account though afaik)

Also I just made a discord https://discord.gg/AwgGwrkn

[keewis]

I don't have access to the twitter account though afaik

@dcherian might be able to help with that?

[TomNicholas]

Great thank you Justus! Any remote help is greatly appreciated as everything is very busy at the moment here. I intended this discussion to advertise the sprint, so we could at least pin this? I could also send an email to the users mailing list? Very late notice though - what does @xarray_dev think?

…

On Fri, 15 Jul 2022, 17:20 Justus Magin, ***@***.***> wrote: That sounds like good advice! I'll try to sort the docstrings according to difficulty tomorrow, most of them *should* be very easy tasks. We should probably also edit the first post in this discussion to be that list of issues. Should we advertise the sprint to people outside of scipy? If so, where would that be? I suppose we will only get spontaneous sprinters, but it might be worth it anyways? — Reply to this email directly, view it on GitHub <https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_pydata_xarray_discussions_6783-23discussioncomment-2D3159868&d=DwMCaQ&c=009klHSCxuh5AI1vNQzSO0KGjl4nbi2Q0M1QLJX9BeE&r=qdISi9HqjazmE0DcySuXts3OlnplnLfKjH4hpzAV0xo&m=hy4v72GPuC8PqxeHL_80IY4_rVb8Q2c_Eyr_YFobpu3PXnSw9BufgE9E6Nzeygqq&s=uDQ3E2jMNtCjkJvz0kKkCajptcEAC18w8ZCCLQjVI90&e=>, or unsubscribe <https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_notifications_unsubscribe-2Dauth_AISNPI7HIYVOMCCRR6MIOBDVUHP3LANCNFSM53S52YNA&d=DwMCaQ&c=009klHSCxuh5AI1vNQzSO0KGjl4nbi2Q0M1QLJX9BeE&r=qdISi9HqjazmE0DcySuXts3OlnplnLfKjH4hpzAV0xo&m=hy4v72GPuC8PqxeHL_80IY4_rVb8Q2c_Eyr_YFobpu3PXnSw9BufgE9E6Nzeygqq&s=PjPSJswvHQ0gddeOqIlBQ1tBJDftG_NmRdTynNkmawY&e=> . You are receiving this because you were mentioned.Message ID: ***@***.***>