diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 18d6a9aba..d4b81f83d 100644 --- a/src/parcels/_core/model.py +++ b/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) diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py index 08ab4d482..04c6fe473 100644 --- a/src/parcels/_datasets/structured/generated.py +++ b/src/parcels/_datasets/structured/generated.py @@ -10,7 +10,8 @@ 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( @@ -18,12 +19,12 @@ def simple_UV_dataset(n_times=360, nz=2, ny=30, nx=4, maxdepth=1, mesh_str="sphe 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, diff --git a/src/parcels/_datasets/structured/strategies.py b/src/parcels/_datasets/structured/strategies.py index 73680f096..10bc51fb3 100644 --- a/src/parcels/_datasets/structured/strategies.py +++ b/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) diff --git a/src/parcels/_sgrid/accessor.py b/src/parcels/_sgrid/accessor.py index 04d050d74..e220da728 100644 --- a/src/parcels/_sgrid/accessor.py +++ b/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 diff --git a/src/parcels/convert.py b/src/parcels/convert.py index d28af6186..ec82ada43 100644 --- a/src/parcels/convert.py +++ b/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), diff --git a/tests/sgrid/test_accessor.py b/tests/sgrid/test_accessor.py index 1f3423ad0..fdfe45de3 100644 --- a/tests/sgrid/test_accessor.py +++ b/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,12 +130,105 @@ 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( @@ -144,11 +236,10 @@ def test_assert_metadata_ds_consistency_failures(ds, dim): ) 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( diff --git a/tests/test_xgrid.py b/tests/test_xgrid.py index 38151cd1c..f487afaaa 100644 --- a/tests/test_xgrid.py +++ b/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)