From 5494320d0e71c550a73863e3d8cee3cbab007dc5 Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Mon, 4 May 2026 14:56:02 +0200 Subject: [PATCH 1/3] Add test test_nemo_to_sgrid_with_depth --- src/parcels/_datasets/remote.py | 11 +++++- tests/test_convert.py | 59 ++++++++++++++++++++++++++++++--- 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/src/parcels/_datasets/remote.py b/src/parcels/_datasets/remote.py index 1e99647d0..eb2206591 100644 --- a/src/parcels/_datasets/remote.py +++ b/src/parcels/_datasets/remote.py @@ -112,7 +112,12 @@ def _get_data_home() -> Path: "data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip", "data-zarr/Benchmarks_FESOM2-baroclinic-gyre/grid.zip", ] - + [] + + [ + "data-zarr/Benchmarks_MOi_data_metadata-only/U.zip", + "data-zarr/Benchmarks_MOi_data_metadata-only/V.zip", + "data-zarr/Benchmarks_MOi_data_metadata-only/W.zip", + "data-zarr/Benchmarks_MOi_data_metadata-only/mesh.zip", + ] ) _ODIE = pooch.create( @@ -231,6 +236,10 @@ class _Purpose(enum.Enum): ] + [ ("Benchmarks_FESOM2-baroclinic-gyre/data", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip', zarr_format=2), _Purpose.TESTING)), ("Benchmarks_FESOM2-baroclinic-gyre/grid", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/grid.zip', zarr_format=2),_Purpose.TESTING)), + ("Benchmarks_MOi_data_metadata-only/U", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/U.zip"), _Purpose.TESTING)), + ("Benchmarks_MOi_data_metadata-only/V", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/V.zip"), _Purpose.TESTING)), + ("Benchmarks_MOi_data_metadata-only/W", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/W.zip"), _Purpose.TESTING)), + ("Benchmarks_MOi_data_metadata-only/mesh", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/mesh.zip"), _Purpose.TESTING)), ]) # fmt: on diff --git a/tests/test_convert.py b/tests/test_convert.py index 9a3859312..385a67655 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -12,11 +12,18 @@ from parcels.interpolators._xinterpolators import _get_offsets_dictionary -def test_nemo_to_sgrid(): - U = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/U") - V = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/V") - coords = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/mesh_mask") - +@pytest.mark.parametrize( + "U, V, coords", + [ + pytest.param( + parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/U"), + parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/V"), + parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/mesh_mask"), + id="NemoCurvilinear_data_zonal", + ), + ], +) +def test_nemo_to_sgrid_2d(U, V, coords): # noqa: N803 ds = convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords) assert ds["grid"].attrs == { @@ -43,6 +50,48 @@ def test_nemo_to_sgrid(): parcels.FieldSet.from_sgrid_conventions(ds) + +@pytest.mark.parametrize( + "U, V, depth, coords", + [ + ( + open_remote_dataset("Benchmarks_MOi_data_metadata-only/U")[["vozocrtx"]].rename_vars({"vozocrtx": "U"}), + open_remote_dataset("Benchmarks_MOi_data_metadata-only/V")[["vomecrty"]].rename_vars({"vomecrty": "V"}), + open_remote_dataset("Benchmarks_MOi_data_metadata-only/W")["depthw"], + open_remote_dataset("Benchmarks_MOi_data_metadata-only/mesh")[["glamf", "gphif"]].isel(t=0), + ), + ], +) +def test_nemo_to_sgrid_with_depth(U, V, depth, coords): # noqa: N803 + coords["depthw"] = depth + ds = parcels.convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords) + + assert ds["grid"].attrs == { + "cf_role": "grid_topology", + "topology_dimension": 2, + "node_dimensions": "x y", + "face_dimensions": "x_center:x (padding:low) y_center:y (padding:low)", + "node_coordinates": "lon lat", + "vertical_dimensions": "depth_center:depth (padding:high)", + } + + meta = sgrid.parse_grid_attrs(ds["grid"].attrs) + + # Assuming that node_dimension1 and node_dimension2 correspond to X and Y respectively + # check that U and V are properly defined on the staggered grid + assert { + meta.get_value_by_id("node_dimension1"), # X edge + meta.get_value_by_id("face_dimension2"), # Y center + }.issubset(set(ds["U"].dims)) + assert { + meta.get_value_by_id("face_dimension1"), # X center + meta.get_value_by_id("node_dimension2"), # Y edge + }.issubset(set(ds["V"].dims)) + # pytest.mark.param(open_remote_dataset(""), open_remote_dataset(""), open_remote_dataset(""), id=""), + + parcels.FieldSet.from_sgrid_conventions(ds) + + def test_convert_nemo_offsets(): U = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/U") V = parcels.tutorial.open_dataset("NemoCurvilinear_data_zonal/V") From 65a9c1937d03da0b7109e8ede0d83c1e462a4342 Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Mon, 4 May 2026 16:31:06 +0200 Subject: [PATCH 2/3] Mypy typing to convert.py --- src/parcels/convert.py | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/src/parcels/convert.py b/src/parcels/convert.py index 4551e0ab3..1ba544b1f 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -22,15 +22,18 @@ from parcels._core.utils import sgrid from parcels._logger import logger +from typing import cast if typing.TYPE_CHECKING: import uxarray as ux + from parcels._typing import XgcmAxisDirection + -_NEMO_EXPECTED_COORDS = [ +_NEMO_EXPECTED_COORDS: list[str] = [ "glamf", "gphif", ] # "depthw" # TODO: Depthw needs to be available if the data has a depth dim. Refactor the whole convert module, this can surely all be handled better. -_NEMO_DIMENSION_COORD_NAMES = [ +_NEMO_DIMENSION_COORD_NAMES: list[str] = [ "x", "y", "time", @@ -44,7 +47,7 @@ "gphif", ] -_NEMO_AXIS_VARNAMES = { +_NEMO_AXIS_VARNAMES: dict[str, XgcmAxisDirection] = { "x": "X", "x_center": "X", "y": "Y", @@ -54,7 +57,7 @@ "time": "T", } -_NEMO_VARNAMES_MAPPING = { +_NEMO_VARNAMES_MAPPING: dict[str, str] = { "time_counter": "time", "depthw": "depth", "uo": "U", @@ -62,9 +65,9 @@ "wo": "W", } -_MITGCM_EXPECTED_COORDS = ["XG", "YG", "Zl"] +_MITGCM_EXPECTED_COORDS: list[str] = ["XG", "YG", "Zl"] -_MITGCM_AXIS_VARNAMES = { +_MITGCM_AXIS_VARNAMES: dict[str, XgcmAxisDirection] = { "XC": "X", "XG": "X", "Xp1": "X", @@ -79,22 +82,22 @@ "time": "T", } -_MITGCM_VARNAMES_MAPPING = { +_MITGCM_VARNAMES_MAPPING: dict[str, str] = { "XG": "lon", "YG": "lat", "Zl": "depth", } -_COPERNICUS_MARINE_AXIS_VARNAMES = { +_COPERNICUS_MARINE_AXIS_VARNAMES: dict[XgcmAxisDirection, str] = { "X": "lon", "Y": "lat", "Z": "depth", "T": "time", } -_CROCO_EXPECTED_COORDS = ["x_rho", "y_rho", "s_w", "time"] +_CROCO_EXPECTED_COORDS: list[str] = ["x_rho", "y_rho", "s_w", "time"] -_CROCO_VARNAMES_MAPPING = { +_CROCO_VARNAMES_MAPPING: dict[str, str] = { "x_rho": "lon", "y_rho": "lat", "s_w": "depth", @@ -111,7 +114,7 @@ def _pick_expected_coords(coords: xr.Dataset, expected_coord_names: list[str]) - return xr.Dataset(coords_to_use) -def _maybe_bring_other_depths_to_depth(ds): +def _maybe_bring_other_depths_to_depth(ds: xr.Dataset): for var in ds.data_vars: for old_depth, target in [ ("depthu", "depth_center"), @@ -129,7 +132,7 @@ def _maybe_bring_other_depths_to_depth(ds): return ds -def _maybe_rename_coords(ds, axis_varnames): +def _maybe_rename_coords(ds: xr.Dataset, axis_varnames: dict[XgcmAxisDirection, str]): try: for axis, [coord] in ds.cf.axes.items(): ds = ds.rename({coord: axis_varnames[axis]}) @@ -138,31 +141,33 @@ def _maybe_rename_coords(ds, axis_varnames): return ds -def _maybe_rename_variables(ds, varnames_mapping): +def _maybe_rename_variables(ds: xr.Dataset, varnames_mapping: dict[str, str]): rename_dict = {old: new for old, new in varnames_mapping.items() if (old in ds.data_vars) or (old in ds.coords)} if rename_dict: ds = ds.rename(rename_dict) return ds -def _assign_dims_as_coords(ds, dimension_names): +def _assign_dims_as_coords(ds: xr.Dataset, dimension_names: list[str]): for axis in dimension_names: if axis in ds.dims and axis not in ds.coords: ds = ds.assign_coords({axis: np.arange(ds.sizes[axis])}) return ds -def _drop_unused_dimensions_and_coords(ds, dimension_and_coord_names): +def _drop_unused_dimensions_and_coords(ds: xr.Dataset, dimension_and_coord_names: list[str]): for dim in ds.dims: if dim not in dimension_and_coord_names: + dim = cast(str, dim) ds = ds.drop_dims(dim, errors="ignore") for coord in ds.coords: + coord = cast(str, coord) if coord not in dimension_and_coord_names: ds = ds.drop_vars(coord, errors="ignore") return ds -def _set_coords(ds, dimension_names): +def _set_coords(ds: xr.Dataset, dimension_names): for varname in dimension_names: if varname in ds and varname not in ds.coords: ds = ds.set_coords([varname]) @@ -176,7 +181,7 @@ def _maybe_remove_depth_from_lonlat(ds): return ds -def _set_axis_attrs(ds, dim_axis): +def _set_axis_attrs(ds: xr.Dataset, dim_axis: dict[str, XgcmAxisDirection]): for dim, axis in dim_axis.items(): if dim in ds: ds[dim].attrs["axis"] = axis From e6aa98e83c8dce52c80919d785f540416ab1f9ec Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Mon, 4 May 2026 16:47:02 +0200 Subject: [PATCH 3/3] Add _Status for coords in convert.py Also make depthw an optional coord for NEMO --- src/parcels/convert.py | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/src/parcels/convert.py b/src/parcels/convert.py index 1ba544b1f..bffe5c944 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -12,6 +12,7 @@ from __future__ import annotations +import enum import typing import warnings from typing import cast @@ -22,16 +23,22 @@ from parcels._core.utils import sgrid from parcels._logger import logger -from typing import cast if typing.TYPE_CHECKING: import uxarray as ux + from parcels._typing import XgcmAxisDirection -_NEMO_EXPECTED_COORDS: list[str] = [ - "glamf", - "gphif", -] # "depthw" # TODO: Depthw needs to be available if the data has a depth dim. Refactor the whole convert module, this can surely all be handled better. +class _Status(enum.Enum): + REQUIRED = enum.auto() + OPTIONAL = enum.auto() + + +_NEMO_EXPECTED_COORDS: list[tuple[str, _Status]] = [ + ("glamf", _Status.REQUIRED), + ("gphif", _Status.REQUIRED), + ("depthw", _Status.OPTIONAL), +] _NEMO_DIMENSION_COORD_NAMES: list[str] = [ "x", @@ -65,7 +72,7 @@ "wo": "W", } -_MITGCM_EXPECTED_COORDS: list[str] = ["XG", "YG", "Zl"] +_MITGCM_EXPECTED_COORDS: list[tuple[str, _Status]] = [(name, _Status.REQUIRED) for name in ["XG", "YG", "Zl"]] _MITGCM_AXIS_VARNAMES: dict[str, XgcmAxisDirection] = { "XC": "X", @@ -95,7 +102,9 @@ "T": "time", } -_CROCO_EXPECTED_COORDS: list[str] = ["x_rho", "y_rho", "s_w", "time"] +_CROCO_EXPECTED_COORDS: list[tuple[str, _Status]] = [ + (name, _Status.REQUIRED) for name in ["x_rho", "y_rho", "s_w", "time"] +] _CROCO_VARNAMES_MAPPING: dict[str, str] = { "x_rho": "lon", @@ -104,13 +113,14 @@ } -def _pick_expected_coords(coords: xr.Dataset, expected_coord_names: list[str]) -> xr.Dataset: +def _pick_expected_coords(coords: xr.Dataset, expected_coord_names: list[tuple[str, _Status]]) -> xr.Dataset: coords_to_use = {} - for name in expected_coord_names: + for name, status in expected_coord_names: if name in coords: coords_to_use[name] = coords[name] else: - raise ValueError(f"Expected coordinate '{name}' not found in provided coords dataset.") + if status == _Status.REQUIRED: + raise ValueError(f"Expected coordinate '{name}' not found in provided coords dataset.") return xr.Dataset(coords_to_use)