Skip to content

Integration with Xarray NDPointIndex #44

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
benbovy merged 15 commits into from
Dec 11, 2025
Merged

Integration with Xarray NDPointIndex #44

benbovy merged 15 commits into from
Dec 11, 2025

Conversation

@benbovy
Copy link
Member

@benbovy benbovy commented Jul 3, 2025 •
edited
Loading

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.

benbovy added 2 commits July 3, 2025 13:30
@benbovy
update pre-commit config
9fc0687 
Temporarily remove flake8 (will likely use ruff later).
@benbovy
Add TreeAdapter classes for xarray's NDPointIndex
4614225
@benbovy benbovy marked this pull request as draft July 3, 2025 11:54
@benbovy benbovy requested a review from willirath July 3, 2025 11:55
@benbovy benbovy mentioned this pull request Jul 3, 2025
Merged
9 tasks
@willirath
Copy link
Contributor

willirath commented Jul 4, 2025

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 benbovy mentioned this pull request Nov 17, 2025
Merged
@benbovy
Copy link
Member Author

benbovy commented Dec 8, 2025

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
Copy link

dcherian commented Dec 8, 2025

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

@benbovy
Copy link
Member Author

benbovy commented Dec 8, 2025

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.

@benbovy benbovy marked this pull request as ready for review December 11, 2025 07:54
@benbovy benbovy merged commit 1e2f3d0 into xarray-contrib:master Dec 11, 2025
10 checks passed
This was referenced Dec 11, 2025
Closed
Closed
@benbovy benbovy deleted the integration-xarray-ndpointindex branch December 11, 2025 09:07
@dcherian
Copy link

dcherian commented Dec 11, 2025

nice work, Benoit!

@dcherian
Copy link

dcherian commented Dec 11, 2025 •
edited
Loading

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@willirath willirath willirath approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Allow setting multiple indexes per DataArray / Dataset

3 participants

@benbovy @willirath @dcherian