Add HEALPix grid mapping support to GridMapping (#614)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: dcherian

cf_xarray/accessor.py

@@ -589,7 +589,46 @@ def _create_grid_mapping(
     cf_name = var.attrs.get("grid_mapping_name", var_name)
 
     # Create CRS from the grid mapping variable
-    crs = pyproj.CRS.from_cf(var.attrs)
+    if cf_name == "healpix":
+        # pyproj does not recognize "healpix" as a grid mapping name,
+        # but the grid uses geographic (lat/lon) coordinates on a sphere.
+        # Build a geographic CRS from the earth_radius parameter.
+        earth_radius = var.attrs.get("earth_radius", 6371229.0)
+        crs = pyproj.CRS.from_json_dict(
+            {
+                "$schema": "https://proj.org/schemas/v0.6/projjson.schema.json",
+                "type": "GeographicCRS",
+                "name": "HEALPix Grid",
+                "datum": {
+                    "type": "GeodeticReferenceFrame",
+                    "name": "Unknown",
+                    "ellipsoid": {
+                        "name": "Custom",
+                        "semi_major_axis": earth_radius,
+                        "semi_minor_axis": earth_radius,
+                    },
+                },
+                "coordinate_system": {
+                    "subtype": "ellipsoidal",
+                    "axis": [
+                        {
+                            "name": "Latitude",
+                            "abbreviation": "lat",
+                            "direction": "north",
+                            "unit": "degree",
+                        },
+                        {
+                            "name": "Longitude",
+                            "abbreviation": "lon",
+                            "direction": "east",
+                            "unit": "degree",
+                        },
+                    ],
+                },
+            }
+        )
+    else:
+        crs = pyproj.CRS.from_cf(var.attrs)
 
     # Get associated coordinate variables, fallback to dimension names
     coordinates: list[Hashable] = grid_mapping_dict.get(var_name, [])
@@ -600,7 +639,11 @@ def _create_grid_mapping(
     #     The appropriate values of the standard_name depend on the grid mapping and are given in Appendix F, Grid Mappings.
     # """
     if not coordinates and len(grid_mapping_dict) == 1:
-        if crs.to_cf().get("grid_mapping_name") == "rotated_latitude_longitude":
+        if cf_name == "healpix":
+            # For HEALPix grids, coordinates are typically latitude/longitude
+            # (if present), but pixel indices are the primary indexing mechanism.
+            xname, yname = "longitude", "latitude"
+        elif crs.to_cf().get("grid_mapping_name") == "rotated_latitude_longitude":
             xname, yname = "grid_longitude", "grid_latitude"
         elif crs.is_geographic:
             xname, yname = "longitude", "latitude"

cf_xarray/datasets.py

@@ -784,6 +784,52 @@ def _create_inexact_bounds():
         np.zeros((10, 20)),
         {"standard_name": "projected_y_coordinate"},
     )
+    # HEALPix dataset: refinement_level=1, nside=2, 12*nside^2 = 48 pixels, nested ordering
+    healpix_ds = xr.Dataset(
+        {
+            "tas": xr.DataArray(
+                np.random.default_rng(42).standard_normal((2, 48)).astype(np.float32),
+                dims=["time", "healpix_index"],
+                attrs={
+                    "standard_name": "air_temperature",
+                    "units": "K",
+                    "coordinates": "height",
+                    "grid_mapping": "healpix",
+                    "cell_methods": "time: mean area: mean",
+                },
+            ),
+            "healpix": xr.DataArray(
+                np.int32(0),
+                attrs={
+                    "grid_mapping_name": "healpix",
+                    "earth_radius": 6371000,
+                    "indexing_scheme": "nested",
+                    "refinement_level": 1,
+                },
+            ),
+        },
+        coords={
+            "time": xr.DataArray(
+                [0.0, 1.0],
+                dims=["time"],
+                attrs={
+                    "standard_name": "time",
+                    "calendar": "proleptic_gregorian",
+                    "units": "days since 2025-06-01",
+                },
+            ),
+            "healpix_index": xr.DataArray(
+                np.arange(48),
+                dims=["healpix_index"],
+                attrs={"standard_name": "healpix_index"},
+            ),
+            "height": xr.DataArray(
+                np.float32(2.0),
+                attrs={"standard_name": "height", "units": "m"},
+            ),
+        },
+    )
+
 except ImportError:
     pass
 

cf_xarray/tests/test_accessor.py

@@ -1285,6 +1285,69 @@ def test_bad_grid_mapping_attribute():
     assert grid_mappings == ()  # No valid grid mappings since 'foo' doesn't exist
 
 
+@requires_pyproj
+def test_healpix_grid_mapping():
+    """Test GridMapping integration for HEALPix grids."""
+    from ..datasets import healpix_ds as ds
+
+    # Grid mapping discovery
+    assert ds.cf.grid_mapping_names == {"healpix": ["healpix"]}
+
+    # Grid mapping propagation to data variables
+    da = ds.cf["air_temperature"]
+    assert "healpix" in da.coords
+
+    # grid_mapping_name property
+    assert da.cf.grid_mapping_name == "healpix"
+
+    # .cf.grid_mappings property on Dataset
+    gms = ds.cf.grid_mappings
+    assert len(gms) == 1
+    gm = gms[0]
+    assert gm.name == "healpix"
+    assert gm.crs is not None
+    assert gm.crs.is_geographic
+    assert gm.array.name == "healpix"
+    assert gm.array.shape == ()  # scalar variable
+    assert isinstance(gm.coordinates, tuple)
+
+    # DataArray grid_mappings should also work
+    da_gms = da.cf.grid_mappings
+    assert len(da_gms) == 1
+    assert da_gms[0].name == "healpix"
+    assert da_gms[0].crs.is_geographic
+
+    # Repr should include healpix
+    assert "healpix" in ds.cf.__repr__()
+
+
+@requires_pyproj
+def test_healpix_crs_properties():
+    """Test that the CRS built for HEALPix has correct properties."""
+    from ..datasets import healpix_ds as ds
+
+    gm = ds.cf.grid_mappings[0]
+    crs = gm.crs
+
+    # Must be geographic (lat/lon on a sphere)
+    assert crs.is_geographic
+    assert not crs.is_projected
+
+    # Earth shape parameters from the grid mapping variable
+    assert crs.ellipsoid.semi_major_metre == 6371000
+    assert crs.ellipsoid.semi_minor_metre == 6371000
+
+    # Should not have an EPSG code (custom sphere)
+    assert crs.to_epsg() is None
+
+    # Verify HEALPix-specific attributes are preserved in the grid mapping array
+    gm_attrs = gm.array.attrs
+    assert gm_attrs["grid_mapping_name"] == "healpix"
+    assert gm_attrs["refinement_level"] == 1
+    assert gm_attrs["indexing_scheme"] == "nested"
+    assert gm_attrs["earth_radius"] == 6371000
+
+
 def test_docstring() -> None:
     assert "One of ('X'" in airds.cf.groupby.__doc__
     assert "Time variable accessor e.g. 'T.month'" in airds.cf.groupby.__doc__