-
- Notifications
You must be signed in to change notification settings - Fork 1.2k
Open
Description
Is your feature request related to a problem?
Currently NDPointIndex (e.g. with KDTree) only allows advanced indexing.
Can we make it support slicing too (with appropriate semantics)?
Describe the solution you'd like
I just wrote this up
import xarray as xr ds = xr.tutorial.open_dataset("ROMS_example").set_xindex(("lat_rho", "lon_rho"), xr.indexes.NDPointIndex,) # intended slicer slicers = dict(lat_rho=slice(28, 29), lon_rho=slice(-91, -89))To get this to work, I did this:
import itertools index = ds.xindexes["lat_rho"] edges = tuple((slicer.start, slicer.stop) for slicer in slicers.values()) vectorized_sel = { name: xr.DataArray(dims=("pts",), data=data) for name, data in zip( slicers.keys(), map(np.asarray, zip(*itertools.product(*edges))) ) } idxrs = index.sel(vectorized_sel, method="nearest").dim_indexers new_slicers = { name: slice(array.min().item(), array.max().item()) for name, array in idxrs.items() } subset =ds.sel(new_slicers)Output
As you can see: this effectively defines slicing as an operation that "returns a continuous slice of data such that every point within the bounding box defined by the slice objects is returned." Does this make sense?
ds.salt.isel(s_rho=-1, ocean_time=0).plot(x="lon_rho", y="lat_rho") ds.salt.isel(s_rho=-1, ocean_time=0, **new_slicers).plot( x="lon_rho", y="lat_rho", cmap="Blues", add_colorbar=False ) x0, x1 = slicers["lon_rho"].start, slicers["lon_rho"].stop y0, y1 = slicers["lat_rho"].start, slicers["lat_rho"].stop import matplotlib.pyplot as plt plt.plot([x0, x1, x1, x0, x0], [y0, y0, y1, y1, y0], color="r")
Describe alternatives you've considered
We could mask out data outside the box, but that then becomes a copy instead of a view.
Additional context
No response
TomNicholas and aladinor
Metadata
Metadata
Assignees
Type
Projects
Status
To do