replace healpy with cdshealpix (#64)

* experimentation on `healpix-geo`

* replace the custom neighbours search with `cdshealpix`

* adapt the neighbours search tests

* skip the "ring" scheme kernel tests

* replace `nside2resol` with a reimplementation

* reuse the `compute_ring` function from the `numpy` implementation

* use `grid_info` to convert points to cell ids

* fix imports

* replace `healpy` with `cdshealpix` and `healpix-geo`

(`healpix-geo` may be a temporary dependency, we only need it to get the
new functionality in `cds-healpix-rust`)

* make sure multi-dimensional cell ids still work

* remove a obsolete TODO

* depend on `cdshealpix` and `healpix-geo`

* remove the direct dependency on `numba`

* reflect the new deps in the env files

* use the released version of `xdggs`

* enable testing on windows

(`cdshealpix` supports it, and so does `healpix-geo`)

* run CI for `python=3.13`

* re-enable the ring scheme tests

* install `xdggs` and `cdshealpix` from PyPI

(`cdshealpix` is not available for macos arm on conda-forge, yet)

* don't test on `python=3.12`

* trim the environment a bit

* install base matplotlib

* don't try to install `conda`

* install `xdggs` from conda-forge

* Revert "install `xdggs` from conda-forge"

This reverts commit cb00dfa.

* new location of `sphinx`' intersphinx file

.github/workflows/ci.yaml

@@ -51,8 +51,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ["ubuntu-latest", "macos-latest"] # no windows, because healpy doesn't support it
-        python-version: ["3.11", "3.12"]
+        os: ["ubuntu-latest", "macos-latest", "windows-latest"]
+        python-version: ["3.11", "3.13"]
 
     steps:
       - uses: actions/checkout@v4
@@ -68,17 +68,11 @@ jobs:
           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 healpix-convolution
         run: |
           python -m pip install --no-deps -e .
 
-      - name: Version info
-        run: |
-          conda info -a
-          conda list
-
       - name: Import healpix-convolution
         run: |
           python -c 'import healpix_convolution'

ci/requirements/docs.yaml

@@ -9,10 +9,9 @@ dependencies:
   - sphinx-remove-toctrees
   - jupyter-sphinx
   - ipython
-  - numba
   - numpy
   - dask
-  - healpy>=1.17
+  - cdshealpix
   - geopandas
   - pandas
   - sparse>=0.15
@@ -26,4 +25,5 @@ dependencies:
   - cartopy_offlinedata
   - pip
   - pip:
+      - healpix-geo
       - -e ../..

ci/requirements/environment.yaml

@@ -7,18 +7,14 @@ dependencies:
   - pytest-reportlog
   - pytest-cov
   - hypothesis
-  - numba
   - numpy
   - dask
-  - healpy>=1.17
-  - geopandas
-  - pandas
   - sparse
   - xarray
-  - folium
-  - shapely
   - opt_einsum
-  - matplotlib
+  - matplotlib-base
   - pip
   - pip:
-      - git+https://github.com/xarray-contrib/xdggs
+      - xdggs
+      - cdshealpix
+      - healpix-geo

docs/conf.py

@@ -95,7 +95,7 @@
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3/", None),
-    "sphinx": ("https://www.sphinx-doc.org/en/stable/", None),
+    "sphinx": ("https://www.sphinx-doc.org/en/master/", None),
     "numpy": ("https://numpy.org/doc/stable", None),
     "xarray": ("https://docs.xarray.dev/en/latest", None),
     "sparse": ("https://sparse.pydata.org/en/latest", None),

experimentation/neighbours.ipynb

@@ -8,6 +8,7 @@
    "outputs": [],
    "source": [
     "import dask.array as da\n",
+    "import healpix_geo\n",
     "import numpy as np\n",
     "\n",
     "import healpix_convolution.neighbours as nb"
@@ -19,6 +20,87 @@
    "id": "1",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "level = 10\n",
+    "ring = 4\n",
+    "cell_ids = np.arange(12 * 4**level)\n",
+    "cell_ids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_bytes = 8\n",
+    "12 * 4**level * (2 * ring + 1) ** 2 * n_bytes / 1e9"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "%load_ext pyinstrument"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "hp_neighbours = nb.neighbours(\n",
+    "    cell_ids, resolution=level, indexing_scheme=\"nested\", ring=ring\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "cds_neighbours = healpix_geo.nested.neighbours_in_kth_ring(cell_ids, level, ring=ring)"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "cds_counts = (cds_neighbours == -1).sum(axis=1)\n",
+    "hp_counts = (hp_neighbours == -1).sum(axis=1)\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "fig, ax = plt.subplots()\n",
+    "ax.plot(cds_counts)\n",
+    "ax.plot(hp_counts)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.testing.assert_equal(np.sort(cds_neighbours, axis=1), np.sort(hp_neighbours, axis=1))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "from distributed import Client\n",
     "\n",
@@ -29,7 +111,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -41,7 +123,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -55,7 +137,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -66,7 +148,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [

healpix_convolution/distances.py

@@ -1,4 +1,3 @@
-import healpy as hp
 import numpy as np
 
 try:
@@ -10,9 +9,18 @@
     dask_array_type = ()
 
 
-def cell_ids2vectors(cell_ids, nside, nest):
-    flattened = cell_ids.flatten()
-    vecs = np.stack(hp.pix2vec(nside, flattened, nest=nest), axis=-1)
+def spherical_to_euclidean(lon, lat):
+    x = np.cos(lon) * np.cos(lat)
+    y = np.sin(lon) * np.cos(lat)
+    z = np.sin(lat)
+
+    return x, y, z
+
+
+def cell_ids2vectors(cell_ids, grid_info):
+    lon, lat = grid_info.cell_ids2geographic(cell_ids.flatten())
+    vecs = np.stack(spherical_to_euclidean(np.deg2rad(lon), np.deg2rad(lat)), axis=-1)
+
     return np.reshape(vecs, cell_ids.shape + (3,))
 
 
@@ -26,12 +34,11 @@ def angle_between_vectors(a, b, axis):
     return np.arccos(argument)
 
 
-def _distances(a, b, axis, nside, nest):
-    vec_a = cell_ids2vectors(a, nside, nest)
+def _distances(a, b, axis, grid_info):
+    vec_a = cell_ids2vectors(a, grid_info)
 
-    # TODO: contains `-1`, which `pix2vec` doesn't like
     mask = b != -1
-    vec_b_ = cell_ids2vectors(np.where(mask, b, 0), nside, nest)
+    vec_b_ = cell_ids2vectors(np.where(mask, b, 0), grid_info)
     vec_b = np.where(mask[..., None], vec_b_, np.nan)
 
     return angle_between_vectors(vec_a, vec_b, axis=axis)
@@ -58,20 +65,14 @@ def angular_distances(neighbours, *, grid_info, axis=None):
     if axis is None:
         axis = -1
 
-    nest = grid_info.nest
-    nside = grid_info.nside
-
     if isinstance(neighbours, dask_array_type):
         return da.map_blocks(
             _distances,
             neighbours[:, :1],
             neighbours,
             axis=axis,
-            nside=nside,
-            nest=nest,
+            grid_info=grid_info,
             chunks=neighbours.chunks,
         )
     else:
-        return _distances(
-            neighbours[:, :1], neighbours, axis=axis, nside=nside, nest=nest
-        )
+        return _distances(neighbours[:, :1], neighbours, axis=axis, grid_info=grid_info)

healpix_convolution/kernels/gaussian.py

@@ -1,4 +1,3 @@
-import healpy as hp
 import numpy as np
 import xdggs
 
@@ -7,11 +6,15 @@
 from healpix_convolution.neighbours import neighbours
 
 
+def healpix_resolution(level):
+    return 2 * np.pi / np.sqrt(12 * 4**level)
+
+
 def compute_ring(resolution, sigma, truncate, kernel_size):
     if kernel_size is not None:
-        ring = int(kernel_size / 2)
+        ring = int(kernel_size // 2)
     else:
-        cell_distance = hp.nside2resol(2**resolution, arcmin=False)
+        cell_distance = healpix_resolution(resolution)
         ring = int((truncate * sigma / cell_distance) // 2)
 
     return ring if ring >= 1 else 1