Integration with Xarray NDPointIndex

.github/workflows/test.yaml

@@ -19,7 +19,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest", "macos-latest", "windows-latest"]
-        python-version: ["3.10", "3.13"]
+        python-version: ["3.11", "3.13"]
     steps:
       - name: Checkout
         uses: actions/checkout@v6

doc/api.rst

@@ -5,11 +5,32 @@ API Reference
 
 This page provides an auto-generated summary of Xoak's API.
 
+.. currentmodule:: xoak
+
+Xarray NDPointIndex tree adapters
+---------------------------------
+
+The following classes may be used with :py:class:`xarray.indexes.NDPointIndex`, they can
+be passed as ``tree_adapter_cls`` option value via :py:meth:`xarray.Dataset.set_xindex` or
+:py:meth:`xarray.DataArray.set_xindex`.
+
+.. autosummary::
+   :toctree: _api_generated/
+
+    S2PointTreeAdapter
+    SklearnBallTreeAdapter
+    SklearnGeoBallTreeAdapter
+    SklearnKDTreeAdapter
+
 .. currentmodule:: xarray
 
 Dataset.xoak
 ------------
 
+.. warning::
+
+   This API is deprecated and will be removed in a future version of Xoak.
+
 This accessor extends :py:class:`xarray.Dataset` with all the methods and
 properties listed below. Proper use of this accessor should be like:
 
@@ -40,6 +61,10 @@ properties listed below. Proper use of this accessor should be like:
 DataArray.xoak
 --------------
 
+.. warning::
+
+   This API is deprecated and will be removed in a future version of Xoak.
+
 The accessor above is also registered for :py:class:`xarray.DataArray`.
 
 **Properties**
@@ -58,39 +83,3 @@ The accessor above is also registered for :py:class:`xarray.DataArray`.
 
     DataArray.xoak.set_index
     DataArray.xoak.sel
-
-Indexes
--------
-
-.. currentmodule:: xoak
-
-.. autosummary::
-   :toctree: _api_generated/
-
-    IndexAdapter
-    IndexRegistry
-
-**Xoak's built-in index adapters**
-
-.. currentmodule:: xoak.index.scipy_adapters
-
-.. autosummary::
-   :toctree: _api_generated/
-
-    ScipyKDTreeAdapter
-
-.. currentmodule:: xoak.index.sklearn_adapters
-
-.. autosummary::
-   :toctree: _api_generated/
-
-    SklearnKDTreeAdapter
-    SklearnBallTreeAdapter
-    SklearnGeoBallTreeAdapter
-
-.. currentmodule:: xoak.index.s2_adapters
-
-.. autosummary::
-   :toctree: _api_generated/
-
-    S2PointIndexAdapter

doc/environment.yml

@@ -10,7 +10,7 @@ dependencies:
   - numpy
   - pys2index
   - scikit-learn
-  - sphinx
+  - sphinx=8.2.3
   - sphinx-autosummary-accessors
   - pydata-sphinx-theme=0.15.4
   - sphinx-book-theme=1.1.4

doc/examples/custom_indexes.ipynb

@@ -6,7 +6,7 @@
    "source": [
     "# Custom Indexes\n",
     "\n",
-    "While Xoak provides some built-in index adapters, it is easy to adapt and register new indexes. "
+    "Xoak provides some built-in adapters for [xarray.indexes.NDPointIndex](https://docs.xarray.dev/en/stable/generated/xarray.indexes.NDPointIndex.html) ; it is easy to create custom ones."
    ]
   },
   {
@@ -16,26 +16,7 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "import xarray as xr\n",
-    "import xoak"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "An instance of `xoak.IndexRegistry` by default contains a collection of Xoak built-in index adapters:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ireg = xoak.IndexRegistry()\n",
-    "\n",
-    "ireg"
+    "import xarray as xr"
    ]
   },
   {
@@ -44,9 +25,7 @@
    "source": [
     "## Example: add a brute-force \"index\"\n",
     "\n",
-    "Every Xoak supported index is a subclass of `xoak.IndexAdapter` that must implement the `build` and `query` methods. The `IndexRegistry.register` decorator may be used to register a new index adpater.\n",
-    "\n",
-    "Let's create and register a new adapter, which simply performs brute-force nearest-neighbor lookup by computing the pairwise distances between all index and query points and finding the minimum distance. "
+    "This example adapter simply performs brute-force nearest-neighbor lookup by computing the pairwise distances between all index and query points and finding the minimum distance. "
    ]
   },
   {
@@ -55,43 +34,32 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from collections.abc import Mapping\n",
+    "from typing import Any\n",
+    "\n",
     "from sklearn.metrics.pairwise import pairwise_distances_argmin_min\n",
+    "from xarray.indexes.nd_point_index import TreeAdapter\n",
     "\n",
     "\n",
-    "@ireg.register('brute_force')\n",
-    "class BruteForceIndex(xoak.IndexAdapter):\n",
+    "class BruteForceTreeAdapter(TreeAdapter):\n",
     "    \"\"\"Brute-force nearest neighbor lookup.\"\"\"\n",
     "    \n",
-    "    def build(self, points):\n",
-    "        # there is no index to build here, just return the points\n",
-    "        return points\n",
-    "    \n",
-    "    def query(self, index, points):\n",
-    "        positions, distances = pairwise_distances_argmin_min(points, index)\n",
-    "        return distances, positions\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This new index now appears in the registry:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ireg"
+    "    def __init__(self, points: np.ndarray, options: Mapping[str, Any]):\n",
+    "        self._index_points = points\n",
+    "\n",
+    "    def query(self, points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n",
+    "        positions, distances = pairwise_distances_argmin_min(points, self._index_points)\n",
+    "        return distances, positions\n",
+    "\n",
+    "    def equals(self, other: \"BruteForceTreeAdapter\") -> bool:\n",
+    "        return np.array_equal(self._index_points, other._index_points)\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Let's use this index in the basic example below:"
+    "Let's use this adapter in the basic example below:"
    ]
   },
   {
@@ -113,7 +81,11 @@
     ")\n",
     "\n",
     "# set the brute-force index (doesn't really build any index in this case)\n",
-    "ds_mesh.xoak.set_index(['meshx', 'meshy'], ireg.brute_force)\n",
+    "ds_mesh = ds_mesh.set_xindex(\n",
+    "    ['meshx', 'meshy'],\n",
+    "    xr.indexes.NDPointIndex,\n",
+    "    tree_adapter_cls=BruteForceTreeAdapter,\n",
+    ")\n",
     "\n",
     "# create trajectory points\n",
     "ds_trajectory = xr.Dataset({\n",
@@ -122,29 +94,23 @@
     "})\n",
     "\n",
     "# select mesh points\n",
-    "ds_selection = ds_mesh.xoak.sel(\n",
+    "ds_selection = ds_mesh.sel(\n",
     "    meshx=ds_trajectory.trajx,\n",
-    "    meshy=ds_trajectory.trajy\n",
+    "    meshy=ds_trajectory.trajy,\n",
+    "    method=\"nearest\",\n",
     ")\n",
     "\n",
     "# plot results\n",
     "ds_trajectory.plot.scatter(x='trajx', y='trajy', c='k', alpha=0.7);\n",
     "ds_selection.plot.scatter(x='meshx', y='meshy', hue='field', alpha=0.9);"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python [conda env:xoak_dev]",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "conda-env-xoak_dev-py"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -156,7 +122,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.6"
+   "version": "3.13.5"
   }
  },
  "nbformat": 4,