|
6 | 6 | "source": [ |
7 | 7 | "# Custom Indexes\n", |
8 | 8 | "\n", |
9 | | - "While Xoak provides some built-in index adapters, it is easy to adapt and register new indexes. " |
| 9 | + "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." |
10 | 10 | ] |
11 | 11 | }, |
12 | 12 | { |
|
16 | 16 | "outputs": [], |
17 | 17 | "source": [ |
18 | 18 | "import numpy as np\n", |
19 | | - "import xarray as xr\n", |
20 | | - "import xoak" |
21 | | - ] |
22 | | - }, |
23 | | - { |
24 | | - "cell_type": "markdown", |
25 | | - "metadata": {}, |
26 | | - "source": [ |
27 | | - "An instance of `xoak.IndexRegistry` by default contains a collection of Xoak built-in index adapters:" |
28 | | - ] |
29 | | - }, |
30 | | - { |
31 | | - "cell_type": "code", |
32 | | - "execution_count": null, |
33 | | - "metadata": {}, |
34 | | - "outputs": [], |
35 | | - "source": [ |
36 | | - "ireg = xoak.IndexRegistry()\n", |
37 | | - "\n", |
38 | | - "ireg" |
| 19 | + "import xarray as xr" |
39 | 20 | ] |
40 | 21 | }, |
41 | 22 | { |
|
44 | 25 | "source": [ |
45 | 26 | "## Example: add a brute-force \"index\"\n", |
46 | 27 | "\n", |
47 | | - "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", |
48 | | - "\n", |
49 | | - "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. " |
| 28 | + "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. " |
50 | 29 | ] |
51 | 30 | }, |
52 | 31 | { |
|
55 | 34 | "metadata": {}, |
56 | 35 | "outputs": [], |
57 | 36 | "source": [ |
| 37 | + "from collections.abc import Mapping\n", |
| 38 | + "from typing import Any\n", |
| 39 | + "\n", |
58 | 40 | "from sklearn.metrics.pairwise import pairwise_distances_argmin_min\n", |
| 41 | + "from xarray.indexes.nd_point_index import TreeAdapter\n", |
59 | 42 | "\n", |
60 | 43 | "\n", |
61 | | - "@ireg.register('brute_force')\n", |
62 | | - "class BruteForceIndex(xoak.IndexAdapter):\n", |
| 44 | + "class BruteForceTreeAdapter(TreeAdapter):\n", |
63 | 45 | " \"\"\"Brute-force nearest neighbor lookup.\"\"\"\n", |
64 | 46 | " \n", |
65 | | - " def build(self, points):\n", |
66 | | - " # there is no index to build here, just return the points\n", |
67 | | - " return points\n", |
68 | | - " \n", |
69 | | - " def query(self, index, points):\n", |
70 | | - " positions, distances = pairwise_distances_argmin_min(points, index)\n", |
71 | | - " return distances, positions\n" |
72 | | - ] |
73 | | - }, |
74 | | - { |
75 | | - "cell_type": "markdown", |
76 | | - "metadata": {}, |
77 | | - "source": [ |
78 | | - "This new index now appears in the registry:" |
79 | | - ] |
80 | | - }, |
81 | | - { |
82 | | - "cell_type": "code", |
83 | | - "execution_count": null, |
84 | | - "metadata": {}, |
85 | | - "outputs": [], |
86 | | - "source": [ |
87 | | - "ireg" |
| 47 | + " def __init__(self, points: np.ndarray, options: Mapping[str, Any]):\n", |
| 48 | + " self._index_points = points\n", |
| 49 | + "\n", |
| 50 | + " def query(self, points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n", |
| 51 | + " positions, distances = pairwise_distances_argmin_min(points, self._index_points)\n", |
| 52 | + " return distances, positions\n", |
| 53 | + "\n", |
| 54 | + " def equals(self, other: \"BruteForceTreeAdapter\") -> bool:\n", |
| 55 | + " return np.array_equal(self._index_points, other._index_points)\n" |
88 | 56 | ] |
89 | 57 | }, |
90 | 58 | { |
91 | 59 | "cell_type": "markdown", |
92 | 60 | "metadata": {}, |
93 | 61 | "source": [ |
94 | | - "Let's use this index in the basic example below:" |
| 62 | + "Let's use this adapter in the basic example below:" |
95 | 63 | ] |
96 | 64 | }, |
97 | 65 | { |
|
113 | 81 | ")\n", |
114 | 82 | "\n", |
115 | 83 | "# set the brute-force index (doesn't really build any index in this case)\n", |
116 | | - "ds_mesh.xoak.set_index(['meshx', 'meshy'], ireg.brute_force)\n", |
| 84 | + "ds_mesh = ds_mesh.set_xindex(\n", |
| 85 | + " ['meshx', 'meshy'],\n", |
| 86 | + " xr.indexes.NDPointIndex,\n", |
| 87 | + " tree_adapter_cls=BruteForceTreeAdapter,\n", |
| 88 | + ")\n", |
117 | 89 | "\n", |
118 | 90 | "# create trajectory points\n", |
119 | 91 | "ds_trajectory = xr.Dataset({\n", |
|
122 | 94 | "})\n", |
123 | 95 | "\n", |
124 | 96 | "# select mesh points\n", |
125 | | - "ds_selection = ds_mesh.xoak.sel(\n", |
| 97 | + "ds_selection = ds_mesh.sel(\n", |
126 | 98 | " meshx=ds_trajectory.trajx,\n", |
127 | | - " meshy=ds_trajectory.trajy\n", |
| 99 | + " meshy=ds_trajectory.trajy,\n", |
| 100 | + " method=\"nearest\",\n", |
128 | 101 | ")\n", |
129 | 102 | "\n", |
130 | 103 | "# plot results\n", |
131 | 104 | "ds_trajectory.plot.scatter(x='trajx', y='trajy', c='k', alpha=0.7);\n", |
132 | 105 | "ds_selection.plot.scatter(x='meshx', y='meshy', hue='field', alpha=0.9);" |
133 | 106 | ] |
134 | | - }, |
135 | | - { |
136 | | - "cell_type": "code", |
137 | | - "execution_count": null, |
138 | | - "metadata": {}, |
139 | | - "outputs": [], |
140 | | - "source": [] |
141 | 107 | } |
142 | 108 | ], |
143 | 109 | "metadata": { |
144 | 110 | "kernelspec": { |
145 | | - "display_name": "Python [conda env:xoak_dev]", |
| 111 | + "display_name": "Python 3 (ipykernel)", |
146 | 112 | "language": "python", |
147 | | - "name": "conda-env-xoak_dev-py" |
| 113 | + "name": "python3" |
148 | 114 | }, |
149 | 115 | "language_info": { |
150 | 116 | "codemirror_mode": { |
|
156 | 122 | "name": "python", |
157 | 123 | "nbconvert_exporter": "python", |
158 | 124 | "pygments_lexer": "ipython3", |
159 | | - "version": "3.8.6" |
| 125 | + "version": "3.13.5" |
160 | 126 | } |
161 | 127 | }, |
162 | 128 | "nbformat": 4, |
|
0 commit comments