hierarchical operations on cell ids: parents (#62)

* index and accessor implementations for `parents`

* return the parents as a `DataArray`

* implement the healpix version of the `parents` computation

* same for h3

* switch to level

* depend on `healpix-geo`

* define `zoom_to`

* docstrings and typing

* implement `zoom_to` for healpix

* partially implement `zoom_to` for H3

Fetching the children will need some thought, the children are returned
to a flattened list while we need a 2D array (with masks)

* check the implementation for healpix

* install `healpix-geo` from PyPI

* check `h3`'s implementation of `zoom_to`

* change to `NotImplementedError`

* actually fill in the expected cell ids

* bump healpix-geo

* install `healpix-geo` from conda-forge

* install `healpix-geo` from github for nightly

* bump lock file

Author: keewis

ci/environment.yml

@@ -0,0 +1,21 @@
+channels:
+  - conda-forge
+dependencies:
+  - xarray
+  - h5netcdf
+  - netcdf4
+  - pooch
+  - matplotlib-base
+  - shapely
+  - pytest
+  - hypothesis
+  - ruff
+  - typing-extensions
+  - geoarrow-pyarrow
+  - lonboard
+  - arro3-core
+  - cdshealpix
+  - h3ronpy
+  - pip
+  - pip:
+      - healpix-geo

pixi.lock

@@ -25,6 +25,7 @@ requires-python = ">=3.11"
 dependencies = [
   "xarray",
   "cdshealpix",
+  "healpix-geo>=0.0.6",
   "h3ronpy",
   "lonboard>=0.9.3",
   "pyproj>=3.3",
@@ -138,6 +139,7 @@ pyproj = ">=3.3"
 matplotlib-base = ">=3.10.5"
 arro3-core = ">=0.4.0"
 pooch = ">=1.8.2"
+healpix-geo = ">=0.0.6"
 
 [tool.pixi.feature.tests.dependencies]
 pytest = "*"
@@ -165,6 +167,7 @@ pyarrow = "*"
 h3ronpy = { git = "git+https://github.com/keewis/h3ronpy.git", subdirectory = "h3ronpy", rev = "version" }
 cdshealpix = { git = "git+https://github.com/cds-astro/cds-healpix-python.git" }
 astropy = { git = "git+https://github.com/astropy/astropy.git" }
+healpix-geo = { git = "git+https://github.com/eopf-dggs/healpix-geo.git" }
 
 [tool.pixi.feature.nightly.dependencies]
 lonboard = ">=0.11.1"

pyproject.toml

@@ -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/accessor.py

@@ -39,6 +39,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/grid.py

@@ -7,12 +7,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,
@@ -196,6 +198,14 @@ def cell_boundaries(self, cell_ids, backend="shapely"):
             raise ValueError(f"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/h3.py

@@ -303,6 +303,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/healpix.py

@@ -98,6 +98,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/index.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_h3.py

@@ -366,6 +366,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"],