Find HEALPix index

easygems/healpix/__init__.py

@@ -45,6 +45,19 @@ def get_npix(dx):
     return healpix.nside2npix(get_nside(dx))
 
 
+def get_index(dx):
+    for c in dx.coords.values():
+        if c.attrs.get("standard_name") == "healpix_index":
+            return c.values
+
+    if "cell" in dx.dims:
+        return dx.cell.values
+    elif "values" in dx.dims:
+        return dx["values"].values
+    elif "value" in dx.dims:
+        return dx.value.values
+
+
 def get_extent_mask(dx, extent):
     lon = dx.lon
     lat = dx.lat
@@ -104,12 +117,7 @@ def guess_crs(ds: xr.Dataset):
         stacklevel=4,
     )
 
-    if "cell" in ds.dims:
-        pix = ds.cell
-    elif "values" in ds.dims:
-        pix = ds.values
-    elif "value" in ds.dims:
-        pix = ds.value
+    pix = get_index(ds)
 
     crs = xr.DataArray(
         name="crs",

tests/test_healpix_coords.py

@@ -1,7 +1,7 @@
 from itertools import product
 
 import pytest
-from easygems.healpix import attach_coords, get_nest, get_nside
+from easygems.healpix import attach_coords, get_index, get_nest, get_nside
 
 import cf_xarray as cf_xarray
 import numpy as np
@@ -78,3 +78,26 @@ def test_get_nside(raw_ds):
 
 def test_get_nest(raw_ds):
     assert get_nest(raw_ds)
+
+
+@pytest.mark.parametrize("known_name", ["cell", "value", "values"])
+def test_get_index_byname(raw_ds, known_name):
+    """Test if we can find the HEALPix index by looking for known short names."""
+    ds = raw_ds.assign_coords({known_name: np.arange(12)})
+
+    assert np.array_equal(get_index(ds), ds[known_name].values)
+
+
+def test_get_index_bycf(raw_ds):
+    """Test if we can find the HEALPix index using CF Conventions."""
+    ds = raw_ds.assign_coords(
+        {
+            "unknown_name": (
+                ("unknown_name",),
+                np.arange(12),
+                {"standard_name": "healpix_index"},
+            )
+        }
+    )
+
+    assert np.array_equal(get_index(ds), ds["unknown_name"].values)