diff --git a/easygems/healpix/__init__.py b/easygems/healpix/__init__.py index 919775b..9d6ddab 100644 --- a/easygems/healpix/__init__.py +++ b/easygems/healpix/__init__.py @@ -25,7 +25,8 @@ def get_nest(dx): def get_nside(dx): try: grid_mapping = dx.cf["grid_mapping"] - except AttributeError: + except (KeyError, AttributeError): + # Catch no CF (KeyError) and no grid mapping (AttributeError) if dx.squeeze().ndim > 1: raise ValueError( "Cannot infer the HEALPix resolution from a multidimensional dataset.\n" diff --git a/tests/test_healpix_coords.py b/tests/test_healpix_coords.py index 432a8a7..df495a0 100644 --- a/tests/test_healpix_coords.py +++ b/tests/test_healpix_coords.py @@ -25,10 +25,9 @@ def raw_ds(request): return xr.Dataset( coords={ - crs_name: ( - ("crs",), - [0], - { + crs_name: xr.DataArray( + name=crs_name, + attrs={ "grid_mapping_name": "healpix", **map_parameters, }, @@ -37,6 +36,16 @@ def raw_ds(request): ) +@pytest.fixture +def ds(raw_ds): + ds = raw_ds.assign_coords( + cell=(("cell",), np.arange(12), {"standard_name": "healpix_index"}) + ) + ds = ds.assign(variable=(("cell",), np.ones(12), {"grid_mapping": "healpix"})) + + return ds + + def test_attach_coords_fixes_crs(raw_ds): ds = attach_coords(raw_ds) @@ -76,6 +85,16 @@ def test_get_nside(raw_ds): assert get_nside(np.arange(12)) == 1 +def test_get_nside_dataarray(ds): + da = ds["variable"] + + # Should work with... + assert get_nside(da) == 1 + + # .. and without CRS on data arrays + assert get_nside(da.drop_vars(da.cf["grid_mapping"].name)) == 1 + + def test_get_nest(raw_ds): assert get_nest(raw_ds)