Add HEALPix Support

ci/asv.yml

@@ -11,6 +11,7 @@ dependencies:
   - holoviews
   - matplotlib-base
   - matplotlib-inline
+  - healpix
   - netcdf4
   - numba
   - numpy<2.1

ci/docs.yml

@@ -12,6 +12,7 @@ dependencies:
     - pytest
     - pytest-cov
     - scipy
+    - healpix
     - xarray
     - pip
     - ipykernel

ci/environment.yml

@@ -11,6 +11,7 @@ dependencies:
   - geoviews
   - h5netcdf
   - hdf5
+  - healpix
   - holoviews
   - hvplot
   - matplotlib-base

docs/user-guide/healpix.ipynb

@@ -0,0 +1,111 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d2efda6d5dc392c6",
+   "metadata": {},
+   "source": "# Working with HEALPix Grids"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "19631a950304b3ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import uxarray as ux\n",
+    "import cartopy.crs as ccrs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb2d4492bf05887c",
+   "metadata": {},
+   "source": "## Representing HEALPix in the UGRID Conventions"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "93f2df22a6decb93",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ux.Grid.from_healpix(resolution_level=3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "428a2bf61174e989",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ux.Grid.from_healpix(resolution_level=3, pixels_only=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c4aa75312d6af2f7",
+   "metadata": {},
+   "source": "## Plotting"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "34c9de64665bcd77",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "(\n",
+    "    ux.Grid.from_healpix(resolution_level=2).plot(\n",
+    "        periodic_elements=\"split\", projection=ccrs.Orthographic()\n",
+    "    )\n",
+    "    + ux.Grid.from_healpix(resolution_level=3).plot(\n",
+    "        periodic_elements=\"split\", projection=ccrs.Orthographic()\n",
+    "    )\n",
+    "    + ux.Grid.from_healpix(resolution_level=4).plot(\n",
+    "        periodic_elements=\"split\", projection=ccrs.Orthographic()\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c71f78a05975445c",
+   "metadata": {},
+   "source": "## Remapping"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "696b7f1b92e35058",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# source_uxds = ux.open_dataset(\"../../\")\n",
+    "# destination_grid = ux.Grid.from_healpix(resolution_level=5)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/userguide.rst

@@ -88,6 +88,9 @@ Supplementary Guides
 
 These user guides provide additional details about specific features in UXarray.
 
+`Working with HEALPix Grids <user-guide/healpix.ipynb>`_
+ Use UXarray with HEALPix
+
 `Compatibility with HoloViz Tools <user-guide/holoviz.ipynb>`_
  Use UXarray with HoloViz tools
 
@@ -117,5 +120,6 @@ These user guides provide additional details about specific features in UXarray.
    user-guide/dual-mesh.ipynb
    user-guide/structured.ipynb
    user-guide/from-points.ipynb
+   user-guide/healpix.ipynb
    user-guide/holoviz.ipynb
    user-guide/from_file.ipynb

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
   "geopandas",
   "xarray",
   "hvplot",
+  "healpix",
   "polars",
 ]
 # minimal dependencies end

uxarray/core/dataarray.py

@@ -153,6 +153,10 @@ def data_mapping(self):
         else:
             return None
 
+    def to_healpix(self):
+        # TODO:
+        pass
+
     def to_geodataframe(
         self,
         periodic_elements: Optional[str] = "exclude",

uxarray/grid/grid.py

@@ -30,6 +30,7 @@
 from uxarray.io._geos import _read_geos_cs
 from uxarray.io._icon import _read_icon
 from uxarray.io._fesom2 import _read_fesom2_asci, _read_fesom2_netcdf
+from uxarray.io._healpix import _pixels_to_ugrid, _populate_healpix_boundaries
 from uxarray.io._structured import _read_structured_grid
 from uxarray.io._voronoi import _spherical_voronoi_from_points
 from uxarray.io._delaunay import (
@@ -176,12 +177,14 @@ def __init__(
         inverse_indices: Optional[xr.Dataset] = None,
     ):
         # check if inputted dataset is a minimum representable 2D UGRID unstructured grid
-        if not _validate_minimum_ugrid(grid_ds):
-            raise ValueError(
-                "Grid unable to be represented in the UGRID conventions. Representing an unstructured grid requires "
-                "at least the following variables: ['node_lon',"
-                "'node_lat', and 'face_node_connectivity']"
-            )
+        # TODO:
+        if source_grid_spec != "HEALPix":
+            if not _validate_minimum_ugrid(grid_ds):
+                raise ValueError(
+                    "Grid unable to be represented in the UGRID conventions. Representing an unstructured grid requires "
+                    "at least the following variables: ['node_lon',"
+                    "'node_lat', and 'face_node_connectivity']"
+                )
 
         # grid spec not provided, check if grid_ds is a minimum representable UGRID dataset
         if source_grid_spec is None:
@@ -196,12 +199,13 @@ def __init__(
         # mapping of ugrid dimensions and variables to source dataset's conventions
         self._source_dims_dict = source_dims_dict
 
-        # source grid specification (i.e. UGRID, MPAS, SCRIP, etc.)
-        self.source_grid_spec = source_grid_spec
-
         # internal xarray dataset for storing grid variables
         self._ds = grid_ds
 
+        # source grid specification (i.e. UGRID, MPAS, SCRIP, etc.)
+        self.source_grid_spec = source_grid_spec
+        self._ds = self._ds.assign_attrs({"source_grid_spec": source_grid_spec})
+
         # initialize attributes
         self._antimeridian_face_indices = None
         self._ds.assign_attrs({"source_grid_spec": self.source_grid_spec})
@@ -579,6 +583,15 @@ def from_face_vertices(
 
         return cls(grid_ds, source_grid_spec="Face Vertices")
 
+    @classmethod
+    def from_healpix(cls, resolution_level, nest=True, pixels_only=True):
+        grid_ds = _pixels_to_ugrid(resolution_level, nest)
+
+        if not pixels_only:
+            _populate_healpix_boundaries(grid_ds)
+
+        return cls.from_dataset(grid_ds, source_grid_spec="HEALPix")
+
     def validate(self, check_duplicates=True):
         """Validates the current ``Grid``, checking for Duplicate Nodes,
         Present Connectivity, and Non-Zero Face Areas.
@@ -878,8 +891,11 @@ def node_lon(self) -> xr.DataArray:
         Dimensions: ``(n_node, )``
         """
         if "node_lon" not in self._ds:
-            _set_desired_longitude_range(self._ds)
-            _populate_node_latlon(self)
+            if self.source_grid_spec == "HEALPix":
+                _populate_healpix_boundaries(self._ds)
+            else:
+                _set_desired_longitude_range(self._ds)
+                _populate_node_latlon(self)
         return self._ds["node_lon"]
 
     @node_lon.setter
@@ -895,8 +911,11 @@ def node_lat(self) -> xr.DataArray:
         Dimensions: ``(n_node, )``
         """
         if "node_lat" not in self._ds:
-            _set_desired_longitude_range(self._ds)
-            _populate_node_latlon(self)
+            if self.source_grid_spec == "HEALPix":
+                _populate_healpix_boundaries(self._ds)
+            else:
+                _set_desired_longitude_range(self._ds)
+                _populate_node_latlon(self)
         return self._ds["node_lat"]
 
     @node_lat.setter
@@ -1131,6 +1150,12 @@ def face_node_connectivity(self) -> xr.DataArray:
         Nodes are in counter-clockwise order.
         """
 
+        if (
+            "face_node_connectivity" not in self._ds
+            and self.source_grid_spec == "HEALPix"
+        ):
+            _populate_healpix_boundaries(self)
+
         if self._ds["face_node_connectivity"].ndim == 1:
             face_node_connectivity_1d = self._ds["face_node_connectivity"].values
             face_node_connectivity_2d = np.expand_dims(

uxarray/io/_healpix.py

@@ -0,0 +1,102 @@
+import healpix as hp
+import xarray as xr
+import numpy as np
+import polars as pl
+from typing import Any, Tuple
+
+import uxarray.conventions.ugrid as ugrid
+from uxarray.grid.coordinates import _set_desired_longitude_range
+from uxarray.constants import INT_DTYPE
+
+
+def pix2corner_ang(
+    nside: int, ipix: Any, nest: bool = False, lonlat: bool = False
+) -> Tuple[Any, Any, Any, Any]:
+    if nest:
+        fu = hp._chp.nest2ang_uv
+    else:
+        fu = hp._chp.ring2ang_uv
+    right = fu(nside, ipix, 1, 0)
+    top = fu(nside, ipix, 1, 1)
+    left = fu(nside, ipix, 0, 1)
+    bottom = fu(nside, ipix, 0, 0)
+    corners = [right, top, left, bottom]
+    if lonlat:
+        corners = [hp.lonlat_from_thetaphi(*x) for x in corners]
+    corners = np.array(corners)
+    if len(corners.shape) == 3:
+        corners = corners.transpose((2, 0, 1))
+    return np.nan_to_num(corners, copy=False)
+
+
+def _pixels_to_ugrid(resolution_level, nest):
+    ds = xr.Dataset()
+
+    nside = hp.order2nside(resolution_level)
+    npix = hp.nside2npix(nside)
+
+    hp_lon, hp_lat = hp.pix2ang(
+        nside=nside, ipix=np.arange(npix), lonlat=True, nest=nest
+    )
+    hp_lon = (hp_lon + 180) % 360 - 180
+
+    ds["face_lon"] = xr.DataArray(hp_lon, dims=["n_face"], attrs=ugrid.FACE_LON_ATTRS)
+    ds["face_lat"] = xr.DataArray(hp_lat, dims=["n_face"], attrs=ugrid.FACE_LAT_ATTRS)
+
+    ds = ds.assign_attrs({"resolution_level": resolution_level})
+    ds = ds.assign_attrs({"n_side": nside})
+    ds = ds.assign_attrs({"n_pix": npix})
+    ds = ds.assign_attrs({"nest": nest})
+
+    return ds
+
+
+def _populate_healpix_boundaries(ds):
+    # Get corners of all the pixels at once
+
+    n_side = ds.attrs["n_side"]
+    n_pix = ds.attrs["n_pix"]
+    nest = ds.attrs["nest"]  # <-- Retrieve 'nest' from ds
+
+    # Pass 'nest' to pix2corner_ang!
+    corners = pix2corner_ang(n_side, np.arange(n_pix), nest=nest, lonlat=True)
+
+    # Flatten the cell corner data to a 2D array of shape (n_cell * 4, 2)
+    nodes = corners.reshape(-1, 2)  # Shape: (n_cell * 4, 2)
+    nodes_df = pl.DataFrame(nodes)
+    unique_nodes_df = nodes_df.unique()
+
+    # Add a unique index to `unique_nodes_df`
+    unique_nodes_with_index = unique_nodes_df.with_row_index("unique_index")
+
+    # Perform a left join on all columns
+    merged = nodes_df.join(
+        unique_nodes_with_index, on=list(nodes_df.columns), how="left"
+    )
+
+    # Extract the inverse indices
+    inverse_indices_df = merged.select("unique_index")
+
+    unique_nodes_df_np = unique_nodes_df.to_numpy()
+    inverse_indices_df_np = inverse_indices_df.to_numpy()
+
+    # Extract node longitude and latitude arrays
+    node_lon = unique_nodes_df_np[:, 0]
+    node_lat = unique_nodes_df_np[:, 1]
+
+    # Reshape inverse indices to get face-node connectivity
+    n_cell = corners.shape[0]
+    face_node_connectivity = inverse_indices_df_np.reshape(n_cell, 4).astype(
+        dtype=INT_DTYPE
+    )
+
+    ds["node_lon"] = xr.DataArray(node_lon, dims=["n_node"], attrs=ugrid.NODE_LON_ATTRS)
+    ds["node_lat"] = xr.DataArray(node_lat, dims=["n_node"], attrs=ugrid.NODE_LAT_ATTRS)
+
+    ds["face_node_connectivity"] = xr.DataArray(
+        face_node_connectivity,
+        dims=ugrid.FACE_NODE_CONNECTIVITY_DIMS,
+        attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS,
+    )
+
+    _set_desired_longitude_range(ds)