support querying by geometries

xdggs/accessor.py

@@ -119,6 +119,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

@@ -49,6 +49,9 @@ def geographic2cell_ids(self, lon, lat):
     def cell_boundaries(self, cell_ids, backend="shapely"):
         raise NotImplementedError()
 
+    def rasterize_geometry(self, geom, *, mode=None):
+        raise NotImplementedError()
+
 
 def translate_parameters(mapping, translations):
     def translate(name, value):

xdggs/h3.py

@@ -13,15 +13,19 @@
 
 try:
     from h3ronpy.vector import (
+        ContainmentMode,
         cells_to_coordinates,
         cells_to_wkb_polygons,
         coordinates_to_cells,
+        geometry_to_cells,
     )
 except ImportError:
     from h3ronpy.arrow.vector import (
+        ContainmentMode,
         cells_to_coordinates,
         cells_to_wkb_polygons,
         coordinates_to_cells,
+        geometry_to_cells,
     )
 from xarray.indexes import PandasIndex
 
@@ -201,6 +205,22 @@ def cell_boundaries(self, cell_ids, backend="shapely"):
             raise ValueError(f"invalid backend: {backend!r}")
         return backend_func(wkb)
 
+    def rasterize_geometry(self, geom, *, mode="contains_centroid"):
+        modes = {
+            "contains_centroid": ContainmentMode.ContainsCentroid,
+            "contains_boundary": ContainmentMode.ContainsBoundary,
+            "covers": ContainmentMode.Covers,
+            "intersects_boundary": ContainmentMode.IntersectsBoundary,
+        }
+        containment_mode = modes.get(mode)
+        if containment_mode is None:
+            raise ValueError(
+                f"invalid mode: {mode}."
+                f" Must be one of [{', '.join(repr(m) for m in modes)}]"
+            )
+
+        return geometry_to_cells(geom, resolution=self.level)
+
 
 @register_dggs("h3")
 class H3Index(DGGSIndex):

xdggs/healpix.py

@@ -308,6 +308,29 @@ def cell_boundaries(self, cell_ids: Any, backend="shapely") -> np.ndarray:
 
         return backend_func(vertices)
 
+    def rasterize_geometry(self, geom, *, mode=None):
+        from astropy.coordinates import Latitude, Longitude
+
+        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_ = Longitude(coords[:, 0], unit="deg")
+        lat_ = Latitude(coords[:, 1], unit="deg")
+
+        ipix, _, _ = cdshealpix.nested.polygon_search(
+            lon_, lat_, depth=self.level, flat=True
+        )
+
+        return ipix
+
 
 @register_dggs("healpix")
 class HealpixIndex(DGGSIndex):

xdggs/index.py

@@ -2,6 +2,7 @@
 from typing import Any, Union
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 from xarray.indexes import Index, PandasIndex
 
@@ -89,6 +90,16 @@ def sel(self, labels, method=None, tolerance=None):
             raise ValueError("finding nearest grid cell has no meaning")
         return self._pd_index.sel(labels, method=method, tolerance=tolerance)
 
+    def query(self, geom):
+        rasterized = self._grid.rasterize_geometry(geom)
+
+        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 _replace(self, new_pd_index: PandasIndex):
         raise NotImplementedError()