replace healpy with cdshealpix (#92)

* start replacing `healpy` with `cdshealpix` * work around a dtype casting bug in cdshealpix * reorder the cell vertices (`cdshealpix` starts with a different pixel than `healpy` – otherwise they're the same) * depend on `cdshealpix` instead of `healpy` * add `cdshealpix-python` to intersphinx * also install `cdshealpix` instead of `healpy` in the upstream-dev ci * change the backend library in the docstring of `from_dict` * test with more dtypes * run CI on windows * don't install `conda` since we don't need it * install `cdshealpix` from PyPI * only enforce colors for `pytest` * override `FORCE_COLOR` for `setup-micromamba` specifically * pin `micromamba` instead * use `assert_allclose` to compare
xarray-contrib · Nov 22, 2024 · ce81a17 · ce81a17
1 parent 61c5ac2
commit ce81a17
Show file tree

Hide file tree

Showing 10 changed files with 85 additions and 61 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -45,7 +45,7 @@ jobs:
       fail-fast: false
 
       matrix:
-        os: ["ubuntu-latest", "macos-latest"]
+        os: ["ubuntu-latest", "macos-latest", "windows-latest"]
         python-version: ["3.10", "3.11", "3.12"]
 
     steps:
@@ -61,13 +61,13 @@ jobs:
       - name: Setup micromamba
         uses: mamba-org/setup-micromamba@v2
         with:
+          micromamba-version: "1.5.10-0"
           environment-file: ${{ env.CONDA_ENV_FILE }}
           environment-name: xdggs-tests
           cache-environment: true
           cache-environment-key: "${{runner.os}}-${{runner.arch}}-py${{matrix.python-version}}-${{env.TODAY}}-${{hashFiles(env.CONDA_ENV_FILE)}}"
           create-args: >-
             python=${{matrix.python-version}}
-            conda
 
       - name: Install xdggs
         run: |

diff --git a/ci/docs.yml b/ci/docs.yml
@@ -12,7 +12,6 @@ dependencies:
   - sphinx-inline-tabs
   - sphinxcontrib-bibtex
   - xarray
-  - healpy
   - h5netcdf
   - netcdf4
   - pooch
@@ -28,4 +27,5 @@ dependencies:
   - pip
   - pip:
       - h3ronpy
+      - cdshealpix
       - -e ..
diff --git a/ci/environment.yml b/ci/environment.yml
@@ -2,7 +2,6 @@ channels:
   - conda-forge
 dependencies:
   - xarray
-  - healpy
   - h5netcdf
   - netcdf4
   - pooch
@@ -18,3 +17,4 @@ dependencies:
   - pip
   - pip:
       - h3ronpy
+      - cdshealpix
diff --git a/ci/install-upstream-wheels.sh b/ci/install-upstream-wheels.sh
@@ -10,7 +10,7 @@ fi
 $conda remove -y --force \
     xarray \
     pandas \
-    healpy
+    cdshealpix
 python -m pip uninstall -y h3ronpy
 
 # install from scientific-python wheels
@@ -36,23 +36,5 @@ python -m pip install \
 python -m pip install --no-deps --upgrade \
     git+https://github.com/Unidata/cftime \
     git+https://github.com/astropy/astropy \
-    "git+https://github.com/nmandery/h3ronpy#subdirectory=h3ronpy"
-
-# install healpy from github
-# need to run `auditwheel` to include the shared libs
-python -m pip install auditwheel
-
-# build and repair the wheel
-# - manual clone
-mkdir deps
-git clone --filter=blob:none --quiet https://github.com/healpy/healpy deps/healpy
-pushd deps/healpy
-git submodule update --init --recursive -q
-# - build and repair the wheel
-mkdir -p built_wheel repaired_wheel
-python -m pip wheel --no-deps . --wheel-dir built_wheel
-auditwheel repair --plat linux_x86_64 -w repaired_wheel built_wheel/healpy-*.whl
-# - install the repaired wheel
-python -m pip install --upgrade --no-deps repaired_wheel/healpy-*.whl
-# - clean up
-popd; rm -rf deps
+    "git+https://github.com/nmandery/h3ronpy#subdirectory=h3ronpy" \
+    git+https://github.com/cds-astro/cds-healpix-python
diff --git a/docs/conf.py b/docs/conf.py
@@ -97,6 +97,7 @@
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable", None),
     "lonboard": ("https://developmentseed.org/lonboard/latest", None),
     "healpy": ("https://healpy.readthedocs.io/en/latest", None),
+    "cdshealpix-python": ("https://cds-astro.github.io/cds-healpix-python", None),
     "shapely": ("https://shapely.readthedocs.io/en/stable", None),
 }
 

diff --git a/docs/tutorials/healpix.ipynb b/docs/tutorials/healpix.ipynb
@@ -152,19 +152,29 @@
     "derived_ds"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "11",
+   "metadata": {},
+   "source": [
+    "```{note}\n",
+    "We need to use {py:func}`xarray.testing.assert_allclose` to compare cell centers because the cell center coordinates were computed using a [different library](https://github.com/healpy/healpy) with a slightly different implementation, resulting in small floating point differences.\n",
+    "```"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "11",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
-    "xr.testing.assert_equal(derived_ds, original_ds)"
+    "xr.testing.assert_allclose(derived_ds, original_ds)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "12",
+   "id": "13",
    "metadata": {},
    "source": [
     "### Cell boundary polygons\n",
@@ -175,7 +185,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -185,7 +195,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "14",
+   "id": "15",
    "metadata": {},
    "source": [
     "## Plotting\n",
@@ -200,7 +210,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "15",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -219,7 +229,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.11.6"
   }
  },
  "nbformat": 4,

diff --git a/environment.yml b/environment.yml
@@ -3,7 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - xarray
-  - healpy
+  - cdshealpix
   - ruff
   - pytest
   - h5netcdf

diff --git a/pyproject.toml b/pyproject.toml
@@ -38,7 +38,7 @@ classifiers = [
 requires-python = ">=3.10"
 dependencies = [
   "xarray",
-  "healpy",
+  "cdshealpix",
   "h3ronpy",
   "typing-extensions",
   "lonboard>=0.9.3",
@@ -90,7 +90,8 @@ fixable = ["I", "TID252"]
 known-first-party = ["xdggs"]
 known-third-party = [
   "xarray",
-  "healpy",
+  "cdshealpix",
+  "astropy",
   "h3ronpy",
 ]
 

diff --git a/xdggs/healpix.py b/xdggs/healpix.py
@@ -8,7 +8,8 @@
 except ImportError:  # pragma: no cover
     from typing_extensions import Self
 
-import healpy
+import cdshealpix.nested
+import cdshealpix.ring
 import numpy as np
 import xarray as xr
 from xarray.indexes import PandasIndex
@@ -105,7 +106,7 @@ class HealpixInfo(DGGSInfo):
 
         .. warning::
             Note that ``"unique"`` is currently not supported as the underlying library
-            (:doc:`healpy <healpy:index>`) does not support it.
+            (:doc:`cdshealpix <cdshealpix-python:index>`) does not support it.
     """
 
     level: int
@@ -210,9 +211,17 @@ def cell_ids2geographic(self, cell_ids):
         lat : array-like
             The latitude coordinate values of the grid cells in degree
         """
-        lon, lat = healpy.pix2ang(self.nside, cell_ids, nest=self.nest, lonlat=True)
+        converters = {
+            "nested": cdshealpix.nested.healpix_to_lonlat,
+            "ring": lambda cell_ids, level: cdshealpix.ring.healpix_to_lonlat(
+                cell_ids, nside=2**level
+            ),
+        }
+        converter = converters[self.indexing_scheme]
+
+        lon, lat = converter(cell_ids, self.level)
 
-        return lon, lat
+        return np.asarray(lon.to("degree")), np.asarray(lat.to("degree"))
 
     def geographic2cell_ids(self, lon, lat):
         """
@@ -233,7 +242,20 @@ def geographic2cell_ids(self, lon, lat):
         cell_ids : array-like
             Array-like containing the cell ids.
         """
-        return healpy.ang2pix(self.nside, lon, lat, lonlat=True, nest=self.nest)
+        from astropy.coordinates import Latitude, Longitude
+
+        converters = {
+            "nested": cdshealpix.nested.lonlat_to_healpix,
+            "ring": lambda lon, lat, level: cdshealpix.ring.lonlat_to_healpix(
+                lon, lat, nside=2**level
+            ),
+        }
+        converter = converters[self.indexing_scheme]
+
+        longitude = Longitude(lon, unit="degree")
+        latitude = Latitude(lat, unit="degree")
+
+        return converter(longitude, latitude, self.level)
 
     def cell_boundaries(self, cell_ids: Any, backend="shapely") -> np.ndarray:
         """
@@ -255,11 +277,19 @@ def cell_boundaries(self, cell_ids: Any, backend="shapely") -> np.ndarray:
             - ``"shapely"``: return a array of :py:class:`shapely.Polygon` objects
             - ``"geoarrow"``: return a ``geoarrow`` array
         """
-        boundary_vectors = healpy.boundaries(
-            self.nside, cell_ids, step=1, nest=self.nest
-        )
+        converters = {
+            "nested": cdshealpix.nested.vertices,
+            "ring": lambda cell_ids, level, **kwargs: cdshealpix.ring.vertices(
+                cell_ids, nside=2**level, **kwargs
+            ),
+        }
+        converter = converters[self.indexing_scheme]
+
+        lon_, lat_ = converter(cell_ids, self.level, step=1)
+
+        lon = np.asarray(lon_.to("degree"))
+        lat = np.asarray(lat_.to("degree"))
 
-        lon, lat = healpy.vec2ang(np.moveaxis(boundary_vectors, 1, -1), lonlat=True)
         lon_reshaped = np.reshape(lon, (-1, 4))
         lat_reshaped = np.reshape(lat, (-1, 4))
 

diff --git a/xdggs/tests/test_healpix.py b/xdggs/tests/test_healpix.py
@@ -55,8 +55,7 @@ def grid_mappings(cls):
 
     def cell_ids(max_value=None, dtypes=None):
         if dtypes is None:
-            # healpy can't deal with `uint32` or less (it segfaults occasionally)
-            dtypes = st.sampled_from(["uint64"])
+            dtypes = st.sampled_from(["int32", "int64", "uint32", "uint64"])
         shapes = npst.array_shapes(min_dims=1, max_dims=1)
 
         return npst.arrays(
@@ -89,9 +88,7 @@ def grid_and_cell_ids(
         cell_levels = st.shared(levels, key="common-levels")
         grid_levels = st.shared(levels, key="common-levels")
         cell_ids_ = cell_levels.flatmap(
-            lambda level: cls.cell_ids(
-                max_value=12 * 2 ** (level * 2) - 1, dtypes=dtypes
-            )
+            lambda level: cls.cell_ids(max_value=12 * 4**level - 1, dtypes=dtypes)
         )
         grids_ = cls.grids(
             levels=grid_levels,
@@ -222,10 +219,10 @@ def test_roundtrip(self, level, indexing_scheme):
                 np.array([2]),
                 np.array(
                     [
-                        [-135.0, 90.0],
-                        [-180.0, 41.8103149],
                         [-135.0, 0.0],
                         [-90.0, 41.8103149],
+                        [-135.0, 90.0],
+                        [-180.0, 41.8103149],
                     ]
                 ),
             ),
@@ -235,16 +232,16 @@ def test_roundtrip(self, level, indexing_scheme):
                 np.array(
                     [
                         [
-                            [0.0, 66.44353569],
-                            [0.0, 54.3409123],
                             [22.5, 41.8103149],
                             [30.0, 54.3409123],
+                            [0.0, 66.44353569],
+                            [0.0, 54.3409123],
                         ],
                         [
-                            [-45.0, 41.8103149],
-                            [-56.25, 30.0],
                             [-45.0, 19.47122063],
                             [-33.75, 30.0],
+                            [-45.0, 41.8103149],
+                            [-56.25, 30.0],
                         ],
                     ]
                 ),
@@ -255,16 +252,16 @@ def test_roundtrip(self, level, indexing_scheme):
                 np.array(
                     [
                         [
-                            [-5.625, 4.78019185],
-                            [-11.25, 0.0],
                             [-5.625, -4.78019185],
                             [0.0, 0.0],
+                            [-5.625, 4.78019185],
+                            [-11.25, 0.0],
                         ],
                         [
-                            [73.125, 35.68533471],
-                            [67.5, 30.0],
                             [73.125, 24.62431835],
                             [78.75, 30.0],
+                            [73.125, 35.68533471],
+                            [67.5, 30.0],
                         ],
                     ]
                 ),
@@ -274,10 +271,10 @@ def test_roundtrip(self, level, indexing_scheme):
                 np.array([79]),
                 np.array(
                     [
-                        [0.0, 41.8103149],
-                        [-11.25, 30],
                         [0.0, 19.47122063],
                         [11.25, 30],
+                        [0.0, 41.8103149],
+                        [-11.25, 30],
                     ]
                 ),
             ),
@@ -300,6 +297,9 @@ def test_cell_boundaries(self, params, cell_ids, backend, expected_coords):
 
     @given(
         *strategies.grid_and_cell_ids(
+            # a dtype casting bug in the valid range check of `cdshealpix`
+            # causes this test to fail for large levels
+            levels=st.integers(min_value=0, max_value=10),
             indexing_schemes=st.sampled_from(["nested", "ring"]),
             dtypes=st.sampled_from(["int64"]),
         )