diff --git a/examples/example_healpy.ipynb b/examples/example_healpy.ipynb index c3e6b19..f5c51ee 100644 --- a/examples/example_healpy.ipynb +++ b/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": [ diff --git a/xdggs/accessor.py b/xdggs/accessor.py index 2b39f4f..85f7241 100644 --- a/xdggs/accessor.py +++ b/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.""" diff --git a/xdggs/healpix.py b/xdggs/healpix.py index de9d791..cd4e168 100644 --- a/xdggs/healpix.py +++ b/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})" diff --git a/xdggs/index.py b/xdggs/index.py index 11b0804..31aea6d 100644 --- a/xdggs/index.py +++ b/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)