Check SGRID node dimensions and coordinates against the dataset

src/parcels/_core/model.py

@@ -24,7 +24,6 @@
 )
 from parcels._logger import logger
 from parcels._python import NOTSET, NotSetType
-from parcels._sgrid.accessor import assert_metadata_ds_consistency
 from parcels.interpolators import (
     CGrid_Velocity,
     Ux_Velocity,
@@ -262,7 +261,7 @@ def from_sgrid_conventions(
         #     ds["lat"] = ds[node_dimensions[1]]
 
         assert_vertical_dimensionality(ds, allowed_dimensionality=[1, 4])
-        assert_metadata_ds_consistency(ds, ds.sgrid.metadata)
+        ds.sgrid.assert_metadata_ds_consistency()
 
         vector_fields = resolve_vector_fields(ds, vector_fields)
         assert_valid_vector_fields(ds, vector_fields)

src/parcels/_datasets/structured/generated.py

@@ -10,20 +10,21 @@
 def simple_UV_dataset(n_times=360, nz=2, ny=30, nx=4, maxdepth=1, mesh_str="spherical"):
     max_lon = 180.0 if mesh_str == "spherical" else 1e6
     max_lat = 90.0 if mesh_str == "spherical" else 1e6
-    unit = "degrees" if mesh_str == "spherical" else "m"
+    x_unit = "degrees_east" if mesh_str == "spherical" else "m"
+    y_unit = "degrees_north" if mesh_str == "spherical" else "m"
     shape = (n_times, nz, ny, nx)
 
     return xr.Dataset(
         {"U": (["time", "depth", "YG", "XG"], np.zeros(shape)), "V": (["time", "depth", "YG", "XG"], np.zeros(shape))},
         coords={
             "time": (["time"], xr.date_range("2000", "2001", n_times), {"axis": "T"}),
             "depth": (["depth"], np.linspace(0, maxdepth, nz), {"axis": "Z", "units": "m"}),
-            "YC": (["YC"], np.arange(ny) + 0.5, {"axis": "Y", "units": unit}),
-            "YG": (["YG"], np.arange(ny), {"axis": "Y", "units": unit}),
-            "XC": (["XC"], np.arange(nx) + 0.5, {"axis": "X", "units": unit}),
-            "XG": (["XG"], np.arange(nx), {"axis": "X", "units": unit}),
-            "lat": (["YG"], np.linspace(-max_lat, max_lat, ny), {"axis": "Y", "units": unit}),
-            "lon": (["XG"], np.linspace(-max_lon, max_lon, nx), {"axis": "X", "units": unit}),
+            "YC": (["YC"], np.arange(ny) + 0.5, {"axis": "Y", "units": y_unit}),
+            "YG": (["YG"], np.arange(ny), {"axis": "Y", "units": y_unit}),
+            "XC": (["XC"], np.arange(nx) + 0.5, {"axis": "X", "units": x_unit}),
+            "XG": (["XG"], np.arange(nx), {"axis": "X", "units": x_unit}),
+            "lat": (["YG"], np.linspace(-max_lat, max_lat, ny), {"axis": "Y", "units": y_unit}),
+            "lon": (["XG"], np.linspace(-max_lon, max_lon, nx), {"axis": "X", "units": x_unit}),
         },
     ).pipe(
         sgrid._attach_sgrid_metadata,
@@ -302,7 +303,8 @@ def peninsula_dataset(xdim=100, ydim=50, mesh_str="flat", grid_type="A"):
     # Convert from m to lat/lon for spherical meshes
     lon = La / 1852.0 / 60.0 if mesh_str == "spherical" else La
     lat = Wa / 1852.0 / 60.0 if mesh_str == "spherical" else Wa
-    unit = "m" if mesh_str == "flat" else "degree"
+    x_unit = "m" if mesh_str == "flat" else "degrees_east"
+    y_unit = "m" if mesh_str == "flat" else "degrees_north"
 
     return xr.Dataset(
         {
@@ -315,8 +317,8 @@ def peninsula_dataset(xdim=100, ydim=50, mesh_str="flat", grid_type="A"):
             "YG": (["YG"], np.arange(ydim), {"axis": "Y"}),
             "XC": (["XC"], np.arange(xdim) - 0.5, {"axis": "X"}),
             "XG": (["XG"], np.arange(xdim), {"axis": "X"}),
-            "lat": (["YG"], lat, {"axis": "Y", "units": unit}),
-            "lon": (["XG"], lon, {"axis": "X", "units": unit}),
+            "lat": (["YG"], lat, {"axis": "Y", "units": y_unit}),
+            "lon": (["XG"], lon, {"axis": "X", "units": x_unit}),
         },
     ).pipe(
         sgrid._attach_sgrid_metadata,

src/parcels/_datasets/structured/strategies.py

@@ -81,8 +81,10 @@ def sgrid_dataset(draw, grid: sgrid.SGrid2DMetadata | None = None) -> xr.Dataset
         data_vars[f"field_{i}"] = (dims, data)
 
     coords = {
-        coord_name1: (coord1_dims, c1),
-        coord_name2: (coord2_dims, c2),
+        node_dim1: ([node_dim1], np.arange(N), {"axis": "X"}),
+        node_dim2: ([node_dim2], np.arange(M), {"axis": "Y"}),
+        coord_name1: (coord1_dims, c1, {"units": "m"}),
+        coord_name2: (coord2_dims, c2, {"units": "m"}),
     }
 
     ds = xr.Dataset(data_vars=data_vars, coords=coords)

src/parcels/_sgrid/accessor.py

@@ -22,6 +22,16 @@ def metadata(self) -> SGrid2DMetadata:
             raise NotImplementedError("Support for 3D SGRID metadata not supported.")
         return grid
 
+    def assert_metadata_ds_consistency(self) -> None:
+        """Assert that the attached SGRID metadata is consistent with the dataset.
+
+        Raises
+        ------
+        SGridDatasetInconsistency
+            If the dataset is not consistent with its SGRID metadata.
+        """
+        _assert_metadata_ds_consistency(self._ds, self.metadata)
+
     def rename(self, name_dict: dict[str, str]) -> xr.Dataset:
         """Similar to Xarray's rename functionality - but also updates attached SGRID metadata."""
         ds = self._ds.copy()
@@ -87,11 +97,17 @@ def isel(self, indexers: Mapping[str, Any] | None = None, **indexers_kwargs):
         indexers = _complete_isel_indexing(self._ds, indexers, metadata)
 
         ds = self._ds.isel(indexers=indexers)
-        assert_metadata_ds_consistency(ds, metadata)
+        _assert_metadata_ds_consistency(ds, metadata)
         return ds
 
 
-def assert_metadata_ds_consistency(ds: xr.Dataset, metadata: SGrid2DMetadata):
+def _assert_metadata_ds_consistency(ds: xr.Dataset, metadata: SGrid2DMetadata):
+    _assert_face_node_sizes_consistent(ds, metadata)
+    _assert_node_dimensions_have_xy_axes(ds, metadata)
+    _assert_node_coordinates_well_defined(ds, metadata)
+
+
+def _assert_face_node_sizes_consistent(ds: xr.Dataset, metadata: SGrid2DMetadata):
     vertical_dimensions: tuple[FaceNodePadding, ...] = metadata.vertical_dimensions or tuple()
 
     for obj in itertools.chain(metadata.face_dimensions, vertical_dimensions):
@@ -113,10 +129,70 @@ def assert_metadata_ds_consistency(ds: xr.Dataset, metadata: SGrid2DMetadata):
                 f"Due to dataset padding of {padding!r}, expected face dimension {face} to actually be size {expected_n_faces}."
             )
 
-    # TODO: Also check on coordinates
+
+# Accepted CF units for longitude and latitude (https://cfconventions.org/cf-conventions/cf-conventions.html#longitude-coordinate)
+_DEGREE_UNITS = {"degree", "degrees"}
+_DEGREE_UNITS_PER_AXIS = {
+    "X": {"degrees_east", "degree_east", "degree_E", "degrees_E", "degreeE", "degreesE"} | _DEGREE_UNITS,
+    "Y": {"degrees_north", "degree_north", "degree_N", "degrees_N", "degreeN", "degreesN"} | _DEGREE_UNITS,
+}
+
+
+def _assert_node_dimensions_have_xy_axes(ds: xr.Dataset, metadata: SGrid2DMetadata):
+    # SGRID conventions don't explicitly mandate an X then Y ordering, but this is evident from all their examples
+    for dim, expected_axis in zip(metadata.node_dimensions, ("X", "Y"), strict=True):
+        if dim not in ds.variables:  # dimension has no coordinate variable in this dataset
+            continue
+
+        axis = ds[dim].attrs.get("axis")
+        if axis != expected_axis:
+            raise SGridDatasetInconsistency(
+                f"SGRID metadata lists {dim!r} in node_dimensions={metadata.node_dimensions!r}, so expected variable {dim!r} "
+                f"to have attribute axis={expected_axis!r}. Instead got axis={axis!r}. "
+                f'HINT: Make sure node_dimensions in the SGRID metadata is ordered as "X" then "Y", and that the '
+                f"coordinate variables in the dataset have the corresponding 'axis' attributes."
+            )
+
+
+def _assert_node_coordinates_well_defined(ds: xr.Dataset, metadata: SGrid2DMetadata):
+    if metadata.node_coordinates is None:
+        return
+
+    node_dims = metadata.node_dimensions
+    both_node_dims_in_ds = all(d in ds.dims for d in node_dims)
+
+    for coord, node_dim, expected_axis in zip(metadata.node_coordinates, node_dims, ("X", "Y"), strict=True):
+        if coord not in ds.variables:
+            raise SGridDatasetInconsistency(
+                f"SGRID metadata lists {coord!r} in node_coordinates={metadata.node_coordinates!r}, "
+                f"but there is no variable {coord!r} in the dataset."
+            )
+
+        da = ds[coord]
+        if not set(da.dims) <= set(node_dims):
+            raise SGridDatasetInconsistency(
+                f"Node coordinate {coord!r} has dims {da.dims}, but node coordinates can only be defined on the "
+                f"node_dimensions={node_dims!r}."
+            )
+        # A 1D node coordinate (i.e., a rectilinear grid) must lie along its own axis. If one of the node
+        # dimensions has been indexed away (e.g., via isel) a 2D coordinate can legitimately be 1D along the other
+        if both_node_dims_in_ds and da.ndim == 1 and da.dims != (node_dim,):
+            raise SGridDatasetInconsistency(
+                f"Node coordinate {coord!r} is 1D with dims {da.dims}, but as entry {expected_axis!r} in "
+                f"node_coordinates={metadata.node_coordinates!r} it is expected to be defined on node dimension {node_dim!r}. "
+                f'HINT: Make sure node_coordinates in the SGRID metadata is ordered as "X" then "Y".'
+            )
+
+        units = da.attrs.get("units")
+        if isinstance(units, str) and "degree" in units.lower() and units not in _DEGREE_UNITS_PER_AXIS[expected_axis]:
+            raise SGridDatasetInconsistency(
+                f"Node coordinate {coord!r} has units={units!r}, but as entry {expected_axis!r} in "
+                f"node_coordinates={metadata.node_coordinates!r} expected one of {sorted(_DEGREE_UNITS_PER_AXIS[expected_axis])}. "
+                f'HINT: Make sure node_coordinates in the SGRID metadata is ordered as "X" (longitude) then "Y" (latitude).'
+            )
 
 
-class SGridDatasetInconsistency(Exception):
+class SGridDatasetInconsistency(ValueError):
     """Attached metadata is not compatible with Xarray dataset"""
 
     pass

src/parcels/convert.py

@@ -599,7 +599,7 @@ def croco_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.D
         sgrid.SGrid2DMetadata(
             cf_role="grid_topology",
             topology_dimension=2,
-            node_dimensions=("lon", "lat"),
+            node_dimensions=("xi_rho", "eta_rho"),
             node_coordinates=("lon", "lat"),
             face_dimensions=(
                 sgrid.FaceNodePadding("xi_u", "xi_rho", sgrid.Padding.NONE),

tests/sgrid/test_accessor.py

@@ -8,7 +8,7 @@
 import parcels._strategies as pst
 from parcels._datasets.structured.generic import datasets_sgrid
 from parcels._datasets.structured.strategies import sgrid_dataset
-from parcels._sgrid.accessor import SGridDatasetInconsistency, assert_metadata_ds_consistency
+from parcels._sgrid.accessor import SGridDatasetInconsistency
 
 
 @st.composite
@@ -81,8 +81,7 @@ def test_rename_dataset(ds):
 
 @given(sgrid_dataset())
 def test_assert_metadata_ds_consistency(ds):
-    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
-    assert_metadata_ds_consistency(ds, metadata)
+    ds.sgrid.assert_metadata_ds_consistency()
 
 
 @given(ds=sgrid_dataset(), dim=st.sampled_from(["face_dimension1", "face_dimension2", "vertical_dimension"]))
@@ -104,7 +103,7 @@ def test_assert_metadata_ds_consistency_dropped_dim(ds, dim):
     assume(fnp.face in ds.dims)
 
     ds = ds.isel({fnp.face: 0})
-    assert_metadata_ds_consistency(ds, metadata)
+    ds.sgrid.assert_metadata_ds_consistency()
 
 
 @given(ds=sgrid_dataset(), dim=st.sampled_from(["face_dimension1", "face_dimension2", "vertical_dimension"]))
@@ -131,24 +130,116 @@ def test_assert_metadata_ds_consistency_failures(ds, dim):
         SGridDatasetInconsistency,
         match="Node dimension .* has size .*, and face dimension .* has size of .* .* expected face dimension .* to actually be size .*",
     ):
-        assert_metadata_ds_consistency(ds, metadata)
+        ds.sgrid.assert_metadata_ds_consistency()
 
 
 _SYMMETRIC_DATASETS = {k: v for k, v in datasets_sgrid.items() if k in ("ds_2d_padded_high", "ds_2d_padded_low")}
 
 
+@given(ds=sgrid_dataset(), axis=st.sampled_from([None, "Y", "Z"]))
+def test_assert_metadata_ds_consistency_node_dimension_axis(ds, axis):
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    node_dim1 = metadata.node_dimensions[0]
+    if axis is None:
+        del ds[node_dim1].attrs["axis"]
+    else:
+        ds[node_dim1].attrs["axis"] = axis
+
+    with pytest.raises(SGridDatasetInconsistency, match=r"expected variable .* to have attribute axis='X'"):
+        ds.sgrid.assert_metadata_ds_consistency()
+
+
+@given(ds=sgrid_dataset())
+def test_assert_metadata_ds_consistency_node_dimensions_swapped(ds):
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    node_dim1, node_dim2 = metadata.node_dimensions
+    ds[node_dim1].attrs["axis"], ds[node_dim2].attrs["axis"] = "Y", "X"
+
+    with pytest.raises(SGridDatasetInconsistency, match=r"expected variable .* to have attribute axis='X'"):
+        ds.sgrid.assert_metadata_ds_consistency()
+
+
+@given(ds=sgrid_dataset())
+def test_assert_metadata_ds_consistency_node_dimension_without_coordinate(ds):
+    # a dimension without a coordinate variable has no axis attribute to check
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    ds = ds.drop_vars(metadata.node_dimensions[0])
+    ds.sgrid.assert_metadata_ds_consistency()
+
+
+@given(ds=sgrid_dataset())
+def test_assert_metadata_ds_consistency_missing_node_coordinate(ds):
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    assert metadata.node_coordinates is not None
+    ds = ds.drop_vars(metadata.node_coordinates[0])
+
+    with pytest.raises(SGridDatasetInconsistency, match=r"but there is no variable .* in the dataset"):
+        ds.sgrid.assert_metadata_ds_consistency()
+
+
+@given(ds=sgrid_dataset())
+def test_assert_metadata_ds_consistency_node_coordinate_on_face(ds):
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    assert metadata.node_coordinates is not None
+    coord = metadata.node_coordinates[0]
+    face_dim = metadata.face_dimensions[0].face
+    assume(face_dim in ds.dims)
+    ds[coord] = xr.DataArray(np.arange(ds.sizes[face_dim], dtype=float), dims=[face_dim], attrs={"units": "m"})
+
+    with pytest.raises(SGridDatasetInconsistency, match=r"node coordinates can only be defined on the node_dimensions"):
+        ds.sgrid.assert_metadata_ds_consistency()
+
+
+@pytest.mark.parametrize("ds", [pytest.param(ds, id=id_) for id_, ds in _SYMMETRIC_DATASETS.items()])
+def test_assert_metadata_ds_consistency_1d_node_coordinates_swapped(ds):
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    assert metadata.node_coordinates is not None
+    coord1, coord2 = metadata.node_coordinates
+    ds = ds.assign_coords({coord1: ds[coord2].variable, coord2: ds[coord1].variable})
+
+    with pytest.raises(SGridDatasetInconsistency, match=r"is 1D with dims .* expected to be defined on node dimension"):
+        ds.sgrid.assert_metadata_ds_consistency()
+
+
+@pytest.mark.parametrize(
+    "x_units, y_units, valid",
+    [
+        ("degrees_east", "degrees_north", True),
+        ("degree_E", "degreesN", True),
+        ("degrees", "degree", True),
+        ("m", "m", True),
+        ("degrees_north", "degrees_east", False),
+        ("degrees_east", "degrees_east", False),
+        ("degrees_north", "degrees_north", False),
+    ],
+)
+@pytest.mark.parametrize("ds", [pytest.param(ds, id=id_) for id_, ds in _SYMMETRIC_DATASETS.items()])
+def test_assert_metadata_ds_consistency_node_coordinate_units(ds, x_units, y_units, valid):
+    ds = ds.copy()
+    metadata: sgrid.SGrid2DMetadata = ds.sgrid.metadata
+    assert metadata.node_coordinates is not None
+    coord1, coord2 = metadata.node_coordinates
+    ds[coord1].attrs["units"] = x_units
+    ds[coord2].attrs["units"] = y_units
+
+    if valid:
+        ds.sgrid.assert_metadata_ds_consistency()
+    else:
+        with pytest.raises(SGridDatasetInconsistency, match=r"Node coordinate .* has units=.*expected one of"):
+            ds.sgrid.assert_metadata_ds_consistency()
+
+
 @pytest.mark.parametrize("ds", [pytest.param(ds, id=id_) for id_, ds in _SYMMETRIC_DATASETS.items()])
 @pytest.mark.parametrize("indexer", [slice(None, None, 3), [0]])
 @pytest.mark.parametrize(
     "node_dim, face_dim", [("node_dimension1", "face_dimension1"), ("node_dimension2", "face_dimension2")]
 )
 def test_isel(ds, indexer, node_dim, face_dim):
     # Covers HIGH/LOW-specific indexer types (lists, step slices) not exercised by property tests.
-    metadata = ds.sgrid.metadata
 
     ds_trimmed = ds.sgrid.isel({node_dim: indexer})
 
-    assert_metadata_ds_consistency(ds_trimmed, metadata)
+    ds_trimmed.sgrid.assert_metadata_ds_consistency()
 
     # Assert that other dims haven't been affected
     for dim, size_before in ds.sizes.items():
@@ -206,7 +297,7 @@ def test_isel_p1_consistency_invariant(ds, s):
 
     result = ds.sgrid.isel({node_dim: s})
 
-    assert_metadata_ds_consistency(result, metadata)
+    result.sgrid.assert_metadata_ds_consistency()
 
 
 @given(

tests/test_xgrid.py

@@ -21,6 +21,7 @@
     _transpose_xfield_data_to_tzyx,
 )
 from parcels._datasets.structured.generic import X, Y, Z, datasets, datasets_sgrid
+from parcels._sgrid.accessor import SGridDatasetInconsistency
 from tests import utils
 
 GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"])
@@ -111,8 +112,8 @@ def test_invalid_lon_lat():
     ds["lat"].attrs["units"] = "meters"
 
     with pytest.raises(
-        ValueError,
-        match=r".*is defined on the faces of the grid, but must be defined on the F nodes\.",
+        SGridDatasetInconsistency,
+        match=r"Node coordinate .* has dims .*, but node coordinates can only be defined on the node_dimensions=.*",
     ):
         FieldSet.from_sgrid_conventions(ds)