support querying by geometries

xdggs/accessor.py

@@ -158,6 +158,17 @@ def sel_latlon(
         }
         return self._obj.sel(cell_indexers)
 
+    def query(self, geom):
+        index, indexer = self._index.query(geom)
+
+        coords = xr.Coordinates.from_xindex(index)
+
+        return (
+            self._obj.drop_indexes(self._name)
+            .isel({self._index._dim: indexer})
+            .assign_coords(coords)
+        )
+
     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/grid.py

@@ -50,6 +50,9 @@ def geographic2cell_ids(self, lon, lat):
     def cell_boundaries(self, cell_ids, backend="shapely"):
         raise NotImplementedError()
 
+    def geometry2cell_ids(self, geom, *, containment=None):
+        raise NotImplementedError()
+
     def zoom_to(self, cell_ids, level: int):
         raise NotImplementedError()
 

xdggs/h3.py

@@ -9,16 +9,20 @@
 try:
     from h3ronpy import change_resolution
     from h3ronpy.vector import (
+        ContainmentMode,
         cells_to_coordinates,
         cells_to_wkb_polygons,
         coordinates_to_cells,
+        geometry_to_cells,
     )
 except ImportError:
     from h3ronpy.arrow import change_resolution
     from h3ronpy.arrow.vector import (
+        ContainmentMode,
         cells_to_coordinates,
         cells_to_wkb_polygons,
         coordinates_to_cells,
+        geometry_to_cells,
     )
 
 from xdggs.grid import DGGSInfo, translate_parameters
@@ -197,6 +201,22 @@ def cell_boundaries(self, cell_ids, backend="shapely"):
             raise ValueError(f"invalid backend: {backend!r}")
         return backend_func(wkb)
 
+    def geometry2cell_ids(self, geom, *, containment="contains_centroid"):
+        modes = {
+            "contains_centroid": ContainmentMode.ContainsCentroid,
+            "contains_boundary": ContainmentMode.ContainsBoundary,
+            "covers": ContainmentMode.Covers,
+            "intersects_boundary": ContainmentMode.IntersectsBoundary,
+        }
+        mode = modes.get(containment)
+        if mode is None:
+            raise ValueError(
+                f"invalid mode: {containment}."
+                f" Must be one of [{', '.join(repr(m) for m in modes)}]"
+            )
+
+        return geometry_to_cells(geom, resolution=self.level)
+
     def zoom_to(self, cell_ids, level):
         if level > self.level:
             raise NotImplementedError(

xdggs/healpix.py

@@ -342,6 +342,29 @@ def cell_boundaries(self, cell_ids: Any, backend="shapely") -> np.ndarray:
 
         return backend_func(vertices)
 
+    def geometry2cell_ids(self, geom, *, containment=None):
+        import healpix_geo.nested
+
+        if self.indexing_scheme != "nested":
+            raise ValueError(
+                "rasterizing geometries is only supported for the 'nested' scheme"
+            )
+
+        if geom.geom_type != "Polygon":
+            raise NotImplementedError(
+                f"geometries of type {geom.geom_type!r} are not supported, yet"
+            )
+
+        coords = np.asarray(geom.exterior.coords)
+        lon_ = coords[:, 0]
+        lat_ = coords[:, 1]
+
+        ipix, _, _ = healpix_geo.nested.polygon_search(
+            lon_, lat_, depth=self.level, ellipsoid=self._format_ellipsoid(), flat=True
+        )
+
+        return ipix
+
     def zoom_to(self, cell_ids, level):
         if self.indexing_scheme == "ring":
             raise ValueError(

xdggs/index.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 from xarray.indexes import Index, PandasIndex
 
@@ -83,6 +84,16 @@ def sel(self, labels, method=None, tolerance=None):
             raise ValueError("finding nearest grid cell has no meaning")
         return self._index.sel(labels, method=method, tolerance=tolerance)
 
+    def query(self, geom, *, containment=None):
+        rasterized = self._grid.geometry2cell_ids(geom, containment=containment)
+
+        geometry_index = pd.Index(rasterized)
+        new_index, _, indexer = geometry_index.join(
+            self._pd_index.index, how="inner", return_indexers=True
+        )
+
+        return self._replace(new_index), indexer
+
     def join(self, other: Self, how: JoinOptions = "inner") -> Self:
         if self.grid_info != other.grid_info:
             raise ValueError(