Integration with Xarray NDPointIndex

[benbovy]

I took some steps in advance and implemented here the integration with Xarray's NDPointIndex (pydata/xarray#10478) that will be available soon (likely in the next Xarray release 2025.7).

Here is how look the introduction examples by directly reusing Xarray's API:

>>> ds_mesh = ds_mesh.set_xindex(
...     ['lat', 'lon'],
...     xr.indexes.NDPointIndex,
...     tree_adapter_cls=xoak.SklearnGeoBallTreeAdapter
... )
>>> ds_mesh
<xarray.Dataset> Size: 240kB
Dimensions:  (x: 100, y: 100)
Coordinates:
  * lat      (x, y) float64 80kB -75.24 -22.48 -53.27 ... -44.86 -24.44 59.81
  * lon      (x, y) float64 80kB -116.4 -125.4 20.51 119.2 ... 136.0 98.4 -79.56
Dimensions without coordinates: x, y
Data variables:
    field    (x, y) float64 80kB -191.6 -147.9 -32.75 ... 91.16 73.97 -19.75
Indexes:
  ┌ lat      NDPointIndex (SklearnGeoBallTreeAdapter)
  └ lon

>>> ds_mesh.sel(
...     lat=ds_trajectory.latitude,
...     lon=ds_trajectory.longitude,
...     method="nearest",
... )
<xarray.Dataset> Size: 720B
Dimensions:  (trajectory: 30)
Coordinates:
    lat      (trajectory) float64 240B -9.39 -9.903 -4.245 ... 36.67 38.23 41.48
    lon      (trajectory) float64 240B -149.4 -140.4 -129.3 ... 138.4 149.8
Dimensions without coordinates: trajectory
Data variables:
    field    (trajectory) float64 240B -158.8 -150.3 -133.6 ... 176.7 191.3

We could probably start deprecating xoak's API. For simplicity Xarray's NDPointIndex doesn't support dask/chunked coordinates and indexers like Xoak currently does. Would it be OK of deprecating that as well? I feel that this experimental feature has never worked really well. We could certainly try supporting it later in Xarray in a more robust way.

This would also close #19, #20 and #37.

[willirath]

We could probably start deprecating xoak's API. For simplicity Xarray's NDPointIndex doesn't support dask/chunked coordinates and indexers like Xoak currently does. Would it be OK of deprecating that as well? I feel that this experimental feature has never worked really well. We could certainly try supporting it later in Xarray in a more robust way.

I agree. Even for the largest meshes, we're not nearly at a point, where the coordinates could not be handled in memory, I think. So feel free to drop dask support here!

[benbovy]

This should be good to go in I think. I updated the tests and the documentation and I added warnings for the deprecated API.

Unless someone has any objection, I plan to merge this soon and release v0.2.0 right away!

[dcherian]

My only request is that you add a page to https://xarray-indexes.readthedocs.io/ :) and write a short blog post

[benbovy]

That seems doable :). Actually one of the NDPointIndex adapters implemented in Xoak is already in https://xarray-indexes.readthedocs.io/blocks/ndpoint.html but I can update the note there after releasing Xoak.

[dcherian]

nice work, Benoit!

[dcherian]

Can you make a pangeo discourse forum post about this, or perhaps you can do a quick blogpost and then make the forum post? I'm sure many people there are interested.