From 37b1cbf4e7adbde7930e05f558ac9aca734cc894 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Tue, 3 Feb 2026 11:37:09 -0700 Subject: [PATCH] Add HEALPix grid mapping support to GridMapping Co-Authored-By: Claude Opus 4.5 --- cf_xarray/accessor.py | 47 +++++++++++++++++++++++- cf_xarray/datasets.py | 46 +++++++++++++++++++++++ cf_xarray/tests/test_accessor.py | 63 ++++++++++++++++++++++++++++++++ 3 files changed, 154 insertions(+), 2 deletions(-) diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py index 21feb6ff..3cc7fef2 100644 --- a/cf_xarray/accessor.py +++ b/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" diff --git a/cf_xarray/datasets.py b/cf_xarray/datasets.py index b38d3ff1..a4a4b443 100644 --- a/cf_xarray/datasets.py +++ b/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 diff --git a/cf_xarray/tests/test_accessor.py b/cf_xarray/tests/test_accessor.py index 3b65fbc6..40b1872c 100644 --- a/cf_xarray/tests/test_accessor.py +++ b/cf_xarray/tests/test_accessor.py @@ -1282,6 +1282,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__