hierarchical operations on cell ids: parents

ci/environment.yml

@@ -16,3 +16,6 @@ dependencies:
   - arro3-core
   - cdshealpix
   - h3ronpy
+  - pip
+  - pip:
+      - healpix-geo

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",

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`
 

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):

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):

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):

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

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)

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"],