Repository navigation
Indexing with alignment and broadcasting #974
Copy link
Copy link
Closed
Labels
Milestone
Description
I think we can bring all of NumPy's advanced indexing to xarray in a very consistent way, with only very minor breaks in backwards compatibility.
For boolean indexing:
da[key]wherekeyis a boolean labelled array (with any number of dimensions) is made equivalent toda.where(key.reindex_like(ds), drop=True). This matches the existing behavior ifkeyis a 1D boolean array. For multi-dimensional arrays, even though the result is now multi-dimensional, this coupled with automatic skipping of NaNs means thatda[key].mean()gives the same result as in NumPy.da[key] = valuewherekeyis a boolean labelled array can be made equivalent toda = da.where(*align(key.reindex_like(da), value.reindex_like(da)))(that is, the three argument form ofwhere).da[key_0, ..., key_n]where all ofkey_iare boolean arrays gets handled in the usual way. It is anIndexingErrorto supply multiple labelled keys if any of them are not already aligned with as the corresponding index coordinates (and share the same dimension name). If they want alignment, we suggest users simply writeda[key_0 & ... & key_n].
For vectorized indexing (by integer or index value):
da[key_0, ..., key_n]where all ofkey_iare integer labelled arrays with any number of dimensions gets handled like NumPy, except instead of broadcasting numpy-style we do broadcasting xarray-style:- If any of
key_iare unlabelled, 1D arrays (e.g., numpy arrays), we convert them into anxarray.Variablealong the respective dimension. 0D arrays remain scalars. This ensures that the result of broadcasting them (in the next step) will be consistent with our current "outer indexing" behavior. Unlabelled higher dimensional arrays triggers anIndexingError. - We ensure all keys have the same dimensions/coordinates by mapping it to
da[*broadcast(key_0, ..., key_n)](note that broadcast now includes automatic alignment). - The result's dimensions and coordinates are copied from the broadcast keys.
- The result's values are taken by mapping each set of integer locations specified by the broadcast version of
key_ito the integer position on the correspondingith axis onda.
- If any of
- Labeled indexing like
ds.loc[key_0, ...., key_n]works exactly as above, except instead of doing integer lookup, we lookup label values in the corresponding index instead. - Indexing with
.iseland.sel/.reindexworks like the two previous cases, except we lookup axes by dimension name instead of axis position. - I haven't fully thought through the implications for assignment (
da[key] = valueorda.loc[key] = value), but I think it works in a straightforwardly similar fashion.
All of these methods should also work for indexing on Dataset by looping over Dataset variables in the usual way.
This framework neatly subsumes most of the major limitations with xarray's existing indexing:
- Boolean indexing on multi-dimensional arrays works in an intuitive way, for both selection and assignment.
- No more need for specialized methods (
sel_points/isel_points) for pointwise indexing. If you want to select along the diagonal of an array, you simply need to supply indexers that use a new dimension. Instead ofarr.sel_points(lat=stations.lat, lon=stations.lon, dim='station'), you would simply writearr.sel(lat=stations.lat, lon=stations.lon)-- thestationdimension is taken automatically from the indexer. - Other use cases for NumPy's advanced indexing that currently are impossible in xarray also automatically work. For example, nearest neighbor interpolation to a completely different grid is now as simple as
ds.reindex(lon=grid.lon, lat=grid.lat, method='nearest', tolerance=0.5)ords.reindex_like(grid, method='nearest', tolerance=0.5).
Questions to consider:
- How does this interact with @benbovy's enhancements for MultiIndex indexing? (Multi-index indexing #802 and Multi-index levels as coordinates #947)
- How do we handle mixed slice and array indexing? In NumPy, this is a major source of confusion, because slicing is done before broadcasting and the order of slices in the result is handled separately from broadcast indices. I think we may be able to resolve this by mapping slices in this case to 1D arrays along their respective axes, and using our normal broadcasting rules.
- Should we deprecate non-boolean indexing with
[]and.loc[]and non-labelled arrays when some but not all dimensions are provided? Instead, we would require explicitly indexing like[key, ...](yes, writing...), which indicates "all trailing axes" like NumPy. This behavior has been suggested for new indexers in NumPy because it precludes a class of bugs where the array has an unexpected number of dimensions. On the other hand, it's not so necessary for us when we have explicit indexing by dimension name with.sel.
xref these comments from @MaximilianR and myself
Note: I would certainly welcome help making this happen from a contributor other than myself, though you should probably wait until I finish #964, first, which lays important groundwork.
Activity
fujiisoup commentedon Jul 5, 2017
What is the current status of these proposed major changes?
Is there any starting basis for them?
I have recently some free time.
shoyer commentedon Jul 5, 2017
I haven't had the time to start working on this -- help would be very gratefully appreciated!
In rough order, I would suggest:
fujiisoup commentedon Jul 10, 2017
@shoyer Thanks for the suggestions.
I don't think I fully understand how boolean indexing works.
Could you show me some example use-cases?
I started from the vectorizing indexing by updating
Variable.__getitem__method,in my folked repo (test).
shoyer commentedon Jul 10, 2017
@fujiisoup My first two bullets for boolean indexing are actually new functionality, so we wouldn't need that for a first pass here. It actually would probably be better to save it for a second PR.
My third bullet on boolean indexing is basically just saying that
da[key_1, ..., key_n]andds.sel(x=key)should be handled in a consistent way with the new indexing behavior when given a boolean array. I don't think this will require any special adjustments -- it should fall out pretty immediately once you get the rest working.shoyer commentedon Jul 10, 2017
I opened a pull request with your branch in #1473 so I can comment/view your changes. I hope that's OK!
shoyer commentedon Feb 4, 2018
Vectorized indexing was implemented by #1473
I've opened a new issue #1887 for multi-dimensional boolean indexing.