diff --git a/ci/environment.yml b/ci/environment.yml index d33d87d0..d45bdc07 100644 --- a/ci/environment.yml +++ b/ci/environment.yml @@ -16,3 +16,6 @@ dependencies: - arro3-core - cdshealpix - h3ronpy + - pip + - pip: + - healpix-geo diff --git a/pyproject.toml b/pyproject.toml index 94f0f46d..bb384f28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ requires-python = ">=3.10" dependencies = [ "xarray", "cdshealpix", + "healpix-geo>=0.0.3", "h3ronpy", "typing-extensions", "lonboard>=0.9.3", diff --git a/xdggs/accessor.py b/xdggs/accessor.py index 1852c981..d0656c35 100644 --- a/xdggs/accessor.py +++ b/xdggs/accessor.py @@ -174,6 +174,31 @@ def cell_boundaries(self): boundaries, coords={self._name: self.cell_ids}, dims=self.cell_ids.dims ) + def zoom_to(self, level: int): + """Change the refinement level of the cell ids to `level`. + + Parameters + ---------- + level : int + The refinement level to change to. Can be smaller than the dataset's + level to compute parents, or bigger to fetch the children. In the + latter case, the array will have an additional `"children"` + dimension. + + Returns + ------- + zoomed : xr.DataArray + The children or parents of the current cells. + """ + zoomed = self.index.zoom_to(level=level) + + if zoomed.ndim == 1: + dims = self.cell_ids.dims + else: + dims = [*self.cell_ids.dims, "children"] + + return xr.DataArray(zoomed, coords={self._name: self.cell_ids}, dims=dims) + def explore(self, *, cmap="viridis", center=None, alpha=None, coords=None): """interactively explore the data using `lonboard` diff --git a/xdggs/grid.py b/xdggs/grid.py index eab18437..abac079e 100644 --- a/xdggs/grid.py +++ b/xdggs/grid.py @@ -49,6 +49,9 @@ def geographic2cell_ids(self, lon, lat): def cell_boundaries(self, cell_ids, backend="shapely"): raise NotImplementedError() + def zoom_to(self, cell_ids, level: int): + raise NotImplementedError() + def translate_parameters(mapping, translations): def translate(name, value): diff --git a/xdggs/h3.py b/xdggs/h3.py index 0a5b103c..e379ef35 100644 --- a/xdggs/h3.py +++ b/xdggs/h3.py @@ -12,12 +12,14 @@ import xarray as xr try: + from h3ronpy import change_resolution from h3ronpy.vector import ( cells_to_coordinates, cells_to_wkb_polygons, coordinates_to_cells, ) except ImportError: + from h3ronpy.arrow import change_resolution from h3ronpy.arrow.vector import ( cells_to_coordinates, cells_to_wkb_polygons, @@ -201,6 +203,14 @@ def cell_boundaries(self, cell_ids, backend="shapely"): raise ValueError("invalid backend: {backend!r}") return backend_func(wkb) + def zoom_to(self, cell_ids, level): + if level > self.level: + raise NotImplementedError( + "extracting children is not supported for H3, yet." + ) + + return np.asarray(change_resolution(cell_ids, level)) + @register_dggs("h3") class H3Index(DGGSIndex): diff --git a/xdggs/healpix.py b/xdggs/healpix.py index 2da42bf5..642fecb1 100644 --- a/xdggs/healpix.py +++ b/xdggs/healpix.py @@ -308,6 +308,17 @@ def cell_boundaries(self, cell_ids: Any, backend="shapely") -> np.ndarray: return backend_func(vertices) + def zoom_to(self, cell_ids, level): + if self.indexing_scheme == "ring": + raise ValueError( + "Scaling does not make sense for the 'ring' scheme." + " Please convert to a nested scheme first." + ) + + from healpix_geo.nested import zoom_to + + return zoom_to(cell_ids, self.level, level) + @register_dggs("healpix") class HealpixIndex(DGGSIndex): diff --git a/xdggs/index.py b/xdggs/index.py index b0f71f28..9f63d2de 100644 --- a/xdggs/index.py +++ b/xdggs/index.py @@ -95,6 +95,9 @@ def cell_centers(self) -> tuple[np.ndarray, np.ndarray]: def cell_boundaries(self) -> np.ndarray: return self.grid_info.cell_boundaries(self._pd_index.index.values) + def zoom_to(self, level: int) -> np.ndarray: + return self._grid.zoom_to(self._pd_index.index.values, level=level) + @property def grid_info(self) -> DGGSInfo: return self._grid diff --git a/xdggs/tests/test_h3.py b/xdggs/tests/test_h3.py index e1c2629d..c9802830 100644 --- a/xdggs/tests/test_h3.py +++ b/xdggs/tests/test_h3.py @@ -209,6 +209,38 @@ def test_cell_boundaries(self, level, cell_ids, backend, expected_coords): shapely.testing.assert_geometries_equal(converter(actual), expected) + @pytest.mark.parametrize( + ["level", "cell_ids", "new_level", "expected"], + ( + pytest.param( + 3, + np.array([0x832833FFFFFFFFF, 0x832834FFFFFFFFF, 0x832835FFFFFFFFF]), + 2, + np.array([0x822837FFFFFFFFF, 0x822837FFFFFFFFF, 0x822837FFFFFFFFF]), + id="parents", + ), + pytest.param( + 3, + np.array([0x832833FFFFFFFFF, 0x832834FFFFFFFFF, 0x832835FFFFFFFFF]), + 1, + np.array([0x81283FFFFFFFFFF, 0x81283FFFFFFFFFF, 0x81283FFFFFFFFFF]), + id="grandparents", + ), + ), + ) + def test_zoom_to(self, level, cell_ids, new_level, expected): + grid = h3.H3Info(level=level) + + actual = grid.zoom_to(cell_ids, level=new_level) + np.testing.assert_equal(actual, expected) + + def test_zoom_to_children_not_implemented(self): + grid = h3.H3Info(level=3) + cell_ids = np.arange(3) + + with pytest.raises(NotImplementedError): + grid.zoom_to(cell_ids, level=4) + @pytest.mark.parametrize("level", levels) @pytest.mark.parametrize("dim", dims) diff --git a/xdggs/tests/test_healpix.py b/xdggs/tests/test_healpix.py index bc75eef7..3e880fe6 100644 --- a/xdggs/tests/test_healpix.py +++ b/xdggs/tests/test_healpix.py @@ -371,6 +371,48 @@ def test_geographic2cell_ids( np.testing.assert_equal(actual, expected) + @pytest.mark.parametrize( + ["level", "cell_ids", "new_level", "expected"], + ( + pytest.param( + 1, + np.array([0, 4, 8, 12, 16]), + 0, + np.array([0, 1, 2, 3, 4]), + id="level1-parents", + ), + pytest.param( + 1, + np.array([0, 1, 2, 3]), + 2, + np.array( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]] + ), + id="level1-children", + ), + pytest.param( + 1, + np.array([0, 4]), + 3, + np.stack([np.arange(16), 4 * 4**2 + np.arange(16)]), + id="level1-grandchildren", + ), + ), + ) + def test_zoom_to(self, level, cell_ids, new_level, expected): + grid = healpix.HealpixInfo(level=level, indexing_scheme="nested") + + actual = grid.zoom_to(cell_ids, level=new_level) + + np.testing.assert_equal(actual, expected) + + def test_zoom_to_ring(self): + cell_ids = np.array([1, 2, 3]) + grid = healpix.HealpixInfo(level=1, indexing_scheme="ring") + + with pytest.raises(ValueError, match="Scaling does not make sense.*'ring'.*"): + grid.zoom_to(cell_ids, level=0) + @pytest.mark.parametrize( ["mapping", "expected"],