diff --git a/xdggs/accessor.py b/xdggs/accessor.py index d3195935..1d87fadf 100644 --- a/xdggs/accessor.py +++ b/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.""" diff --git a/xdggs/grid.py b/xdggs/grid.py index 8aef6e2d..8f5afdc6 100644 --- a/xdggs/grid.py +++ b/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() diff --git a/xdggs/h3.py b/xdggs/h3.py index 3a780654..d08395b7 100644 --- a/xdggs/h3.py +++ b/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( diff --git a/xdggs/healpix.py b/xdggs/healpix.py index dc6bde0d..b41b1367 100644 --- a/xdggs/healpix.py +++ b/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( diff --git a/xdggs/index.py b/xdggs/index.py index 358ee2bc..cb28e83b 100644 --- a/xdggs/index.py +++ b/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(