Add testing of SGRID xgcm ingestion helpers

Author: VeckoTheGecko

src/parcels/_core/sgrid.py

@@ -313,7 +313,7 @@ class SGridParsingException(Exception):
     pass
 
 
-def parse_grid(attrs: dict[str, Hashable]) -> Grid2DMetadata | Grid3DMetadata:
+def parse_grid_attrs(attrs: dict[str, Hashable]) -> Grid2DMetadata | Grid3DMetadata:
     grid: Grid2DMetadata | Grid3DMetadata
     try:
         grid = Grid2DMetadata.from_attrs(attrs)
@@ -341,7 +341,7 @@ def parse_sgrid(ds: xr.Dataset):
     try:
         grid_topology = get_grid_topology(ds)
         assert grid_topology is not None, "No grid_topology variable found in dataset"
-        grid = parse_grid(grid_topology.attrs)
+        grid = parse_grid_attrs(grid_topology.attrs)
 
     except Exception as e:
         raise SGridParsingException(f"Error parsing {grid_topology=!r}") from e

tests/strategies/sgrid.py

@@ -7,20 +7,15 @@
 cover SGrid to the extent to which Parcels is concerned.
 """
 
-import string
-
+import xarray.testing.strategies as xr_st
 from hypothesis import strategies as st
 
 from parcels._core import sgrid
 
-ALLOWED_DIM_LETTERS = (
-    string.ascii_letters + string.digits + "_"
-)  # We can make this more aligned with SGrid by adjusting our regex - but this is good for now
-
 padding = st.sampled_from(sgrid.Padding)
-dimension_name = st.text(
-    min_size=1, alphabet=st.characters(categories=(), whitelist_characters=ALLOWED_DIM_LETTERS)
-).filter(lambda s: " " not in s)  # assuming for now spaces are allowed in dimension names in SGrid convention
+dimension_name = xr_st.names().filter(
+    lambda s: " " not in s
+)  # assuming for now spaces are allowed in dimension names in SGrid convention
 dim_dim_padding = (
     st.tuples(dimension_name, dimension_name, padding)
     .filter(lambda t: t[0] != t[1])

tests/test_sgrid.py

@@ -1,10 +1,129 @@
+import numpy as np
 import pytest
+import xarray as xr
+import xgcm
 from hypothesis import example, given
 
 from parcels._core import sgrid
 from tests.strategies import sgrid as sgrid_strategies
 
 
+def get_unique_dim_names(grid: sgrid.Grid2DMetadata | sgrid.Grid3DMetadata) -> set[str]:
+    dims = set()
+    dims.update(set(grid.node_dimensions))
+
+    for value in [
+        grid.node_dimensions,
+        grid.face_dimensions if isinstance(grid, sgrid.Grid2DMetadata) else grid.volume_dimensions,
+        grid.vertical_dimensions if isinstance(grid, sgrid.Grid2DMetadata) else None,
+    ]:
+        if value is None:
+            continue
+        for item in value:
+            if isinstance(item, sgrid.DimDimPadding):
+                dims.add(item.dim1)
+                dims.add(item.dim2)
+            else:
+                assert isinstance(item, str)
+                dims.add(item)
+    return dims
+
+
+def dummy_sgrid_ds(grid: sgrid.Grid2DMetadata | sgrid.Grid3DMetadata) -> xr.Dataset:
+    if isinstance(grid, sgrid.Grid2DMetadata):
+        return dummy_sgrid_2d_ds(grid)
+    elif isinstance(grid, sgrid.Grid3DMetadata):
+        return dummy_sgrid_3d_ds(grid)
+    else:
+        raise NotImplementedError(f"Cannot create dummy SGrid dataset for grid type {type(grid)}")
+
+
+def dummy_sgrid_2d_ds(grid: sgrid.Grid2DMetadata) -> xr.Dataset:
+    ds = dummy_comodo_3d_ds()
+
+    renamings = {}
+    if grid.vertical_dimensions is None:
+        ds = ds.isel(ZC=0, ZG=0)
+    else:
+        renamings.update({"ZC": grid.vertical_dimensions[0].dim2, "ZG": grid.vertical_dimensions[0].dim1})
+
+    for old, new in zip(["XG", "YG"], grid.node_dimensions, strict=True):
+        renamings[old] = new
+
+    for old, dim_dim_padding in zip(["XC", "YC"], grid.face_dimensions, strict=True):
+        renamings[old] = dim_dim_padding.dim1
+
+    ds = ds.rename_dims(renamings)
+
+    ds["grid"] = xr.DataArray(1, attrs=grid.to_attrs())
+    ds.attrs["convention"] = "SGRID"
+    return ds
+
+
+def dummy_sgrid_3d_ds(grid: sgrid.Grid3DMetadata) -> xr.Dataset:
+    ds = dummy_comodo_3d_ds()
+
+    renamings = {}
+    for old, new in zip(["XG", "YG", "ZG"], grid.node_dimensions, strict=True):
+        renamings[old] = new
+
+    for old, dim_dim_padding in zip(["XC", "YC", "ZC"], grid.volume_dimensions, strict=True):
+        renamings[old] = dim_dim_padding.dim1
+
+    ds = ds.rename_dims(renamings)
+
+    ds["grid"] = xr.DataArray(1, attrs=grid.to_attrs())
+    ds.attrs["convention"] = "SGRID"
+    return ds
+
+
+def dummy_comodo_3d_ds() -> xr.Dataset:
+    T, Z, Y, X = 7, 6, 5, 4
+    TIME = xr.date_range("2000", "2001", T)
+    return xr.Dataset(
+        {
+            "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
+            "U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
+        },
+        coords={
+            # "XG": (
+            #     ["XG"],
+            #     2 * np.pi / X * np.arange(0, X),
+            #     {"axis": "X", "c_grid_axis_shift": -0.5},
+            # ),
+            # "XC": (["XC"], 2 * np.pi / X * (np.arange(0, X) + 0.5), {"axis": "X"}),
+            # "YG": (
+            #     ["YG"],
+            #     2 * np.pi / (Y) * np.arange(0, Y),
+            #     {"axis": "Y", "c_grid_axis_shift": -0.5},
+            # ),
+            # "YC": (
+            #     ["YC"],
+            #     2 * np.pi / (Y) * (np.arange(0, Y) + 0.5),
+            #     {"axis": "Y"},
+            # ),
+            # "ZG": (
+            #     ["ZG"],
+            #     np.arange(Z),
+            #     {"axis": "Z", "c_grid_axis_shift": -0.5},
+            # ),
+            # "ZC": (
+            #     ["ZC"],
+            #     np.arange(Z) + 0.5,
+            #     {"axis": "Z"},
+            # ),
+            # "lon": (["XG"], 2 * np.pi / X * np.arange(0, X)),
+            # "lat": (["YG"], 2 * np.pi / (Y) * np.arange(0, Y)),
+            # "depth": (["ZG"], np.arange(Z)),
+            "time": (["time"], TIME, {"axis": "T"}),
+        },
+    )
+
+
 @example(
     edge_node_padding=(
         sgrid.DimDimPadding("edge1", "node1", sgrid.Padding.NONE),
@@ -69,3 +188,49 @@ def test_Grid3DMetadata_roundtrip(grid: sgrid.Grid3DMetadata):
     attrs = grid.to_attrs()
     parsed = sgrid.Grid3DMetadata.from_attrs(attrs)
     assert parsed == grid
+
+
+@given(sgrid_strategies.grid_metadata)
+def test_parse_grid_attrs(grid: sgrid.SGridMetadataProtocol):
+    attrs = grid.to_attrs()
+    parsed = sgrid.parse_grid_attrs(attrs)
+    assert parsed == grid
+
+
+@given(sgrid_strategies.grid2Dmetadata())
+def test_parse_sgrid_2d(grid_metadata: sgrid.Grid2DMetadata):
+    """Test the ingestion of datasets in XGCM to ensure that it matches the SGRID metadata provided"""
+    ds = dummy_sgrid_2d_ds(grid_metadata)
+
+    ds, xgcm_kwargs = sgrid.parse_sgrid(ds)
+    grid = xgcm.Grid(ds, autoparse_metadata=False, **xgcm_kwargs)
+
+    for ddp, axis in zip(grid_metadata.face_dimensions, ["X", "Y"], strict=True):
+        dim_node, dim_edge, padding = ddp.dim1, ddp.dim2, ddp.padding
+        coords = grid.axes[axis].coords
+        assert coords["center"] == dim_edge
+        assert coords[sgrid.SGRID_PADDING_TO_XGCM_POSITION[padding]] == dim_node
+
+    if grid_metadata.vertical_dimensions is None:
+        assert "Z" not in grid.axes
+    else:
+        ddp = grid_metadata.vertical_dimensions[0]
+        dim_node, dim_edge, padding = ddp.dim1, ddp.dim2, ddp.padding
+        coords = grid.axes["Z"].coords
+        assert coords["center"] == dim_edge
+        assert coords[sgrid.SGRID_PADDING_TO_XGCM_POSITION[padding]] == dim_node
+
+
+@given(sgrid_strategies.grid3Dmetadata())
+def test_parse_sgrid_3d(grid_metadata: sgrid.Grid3DMetadata):
+    """Test the ingestion of datasets in XGCM to ensure that it matches the SGRID metadata provided"""
+    ds = dummy_sgrid_3d_ds(grid_metadata)
+
+    ds, xgcm_kwargs = sgrid.parse_sgrid(ds)
+    grid = xgcm.Grid(ds, autoparse_metadata=False, **xgcm_kwargs)
+
+    for ddp, axis in zip(grid_metadata.volume_dimensions, ["X", "Y", "Z"], strict=True):
+        dim_node, dim_edge, padding = ddp.dim1, ddp.dim2, ddp.padding
+        coords = grid.axes[axis].coords
+        assert coords["center"] == dim_edge
+        assert coords[sgrid.SGRID_PADDING_TO_XGCM_POSITION[padding]] == dim_node