xarray-contrib / xoak

xarray extension that provides tree-based indexes used for selecting irregular, n-dimensional data.
https://xoak.readthedocs.io
MIT License
57 stars 4 forks source link

indexing at a single point #37

Open dcherian opened 3 years ago

dcherian commented 3 years ago
import numpy as np
import xarray as xr
import xoak

shape = (100, 100)
lat = np.random.uniform(-90, 90, size=shape)
lon = np.random.uniform(-180, 180, size=shape)

field = lat + lon
ds_mesh = xr.Dataset(
    coords={'lat': (('x', 'y'), lat), 'lon': (('x', 'y'), lon)},
    data_vars={'field': (('x', 'y'), field)},
).expand_dims(z=10)

ds_mesh.xoak.set_index(("lat", "lon"), "scipy_kdtree")

This does not work

ds_mesh.xoak.sel(lat=50, lon=50)

This workaround does what I want:

tempdim = "__xoak_points__"

ds_mesh.xoak.sel(
    lon=xr.Variable(tempdim, [50]), lat=xr.Variable(tempdim, [50])
).squeeze(tempdim)

Would you be open to adding this workaround?

benbovy commented 3 years ago

Sorry for the late reply @dcherian!

Yes, I think it would be nice to support selection of a single point with something more convenient than the workaround you've shown.