add support for querying by geometry

examples/example_healpy.ipynb

@@ -17,7 +17,8 @@
    ],
    "source": [
     "import xarray as xr\n",
-    "import xdggs"
+    "import xdggs\n",
+    "import shapely"
    ]
   },
   {
@@ -2210,7 +2211,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
+   "id": "21254214-e3b8-489c-9386-2f3b00fcb072",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "geom = shapely.box(-6.0, 48.0, -5.5, 48.5)\n",
+    "geom"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99e656f8-7be0-4b35-9492-a4cef6efbe5e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_idx.dggs.query(geom, inclusive=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "86959d2d-9316-4711-a765-c33b60315adf",
    "metadata": {},
    "outputs": [

xdggs/accessor.py

@@ -1,4 +1,6 @@
+import numpy as np
 import numpy.typing as npt
+import shapely
 import xarray as xr
 
 from xdggs.index import DGGSIndex
@@ -53,6 +55,12 @@ def sel_latlon(
         cell_indexers = {self._name: self.index._latlon2cellid(latitude, longitude)}
         return self._obj.sel(cell_indexers)
 
+    def query(self, geom: shapely.Geometry, **options) -> xr.Dataset | xr.DataArray:
+        cell_ids = self._index._geom2cellid(geom, options)
+        mask = np.isin(cell_ids, self._index._pd_index.index.values)
+        cell_indexers = {self._name: cell_ids[mask]}
+        return self._obj.sel(cell_indexers)
+
     def assign_latlon_coords(self) -> xr.Dataset | xr.DataArray:
         """Return a new Dataset or DataArray with new "latitude" and "longitude"
         coordinates representing the grid cell centers."""

xdggs/healpix.py

@@ -3,6 +3,7 @@
 
 import healpy
 import numpy as np
+import shapely
 import xarray as xr
 from xarray.indexes import PandasIndex
 
@@ -51,6 +52,14 @@ def _cellid2latlon(self, cell_ids: Any) -> tuple[np.ndarray, np.ndarray]:
         lon, lat = healpy.pix2ang(self._nside, cell_ids, nest=self._nest, lonlat=True)
         return lat, -lon
 
+    def _geom2cellid(self, geom: shapely.Geometry, options: dict[str, Any]):
+        coords = shapely.get_coordinates(geom)
+
+        lon, lat = coords[:-1, 0], coords[:-1, 1]
+        vertices = healpy.ang2vec(lon, lat, lonlat=True)
+
+        return healpy.query_polygon(self._nside, vertices, nest=self._nest, **options)
+
     def _repr_inline_(self, max_width: int):
         return (
             f"HealpixIndex(nside={self._nside}, nest={self._nest}, rot_latlon={self._rot_latlon!r})"

xdggs/index.py

@@ -64,6 +64,9 @@ def _cellid2latlon(self, cell_ids: Any) -> tuple[np.ndarray, np.ndarray]:
         """convert cell ids to latitude / longitude (cell centers)."""
         raise NotImplementedError()
 
+    def _geom2cellid(self, geom, options):
+        raise NotImplementedError()
+
     @property
     def cell_centers(self) -> tuple[np.ndarray, np.ndarray]:
         return self._cellid2latlon(self._pd_index.index.values)