diff --git a/easygems/healpix/__init__.py b/easygems/healpix/__init__.py index c228628..919775b 100644 --- a/easygems/healpix/__init__.py +++ b/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", diff --git a/tests/test_healpix_coords.py b/tests/test_healpix_coords.py index ff52939..432a8a7 100644 --- a/tests/test_healpix_coords.py +++ b/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)