Skip to content

Add support in convert.nemo_to_sgrid for depthw coordinate #2626

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
VeckoTheGecko wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add support in convert.nemo_to_sgrid for depthw coordinate #2626

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8dc3bec
Add conversion to fieldset to test_nemo_to_sgrid
VeckoTheGecko May 4, 2026
5494320
Add test test_nemo_to_sgrid_with_depth
VeckoTheGecko May 4, 2026
65a9c19
Mypy typing to convert.py
VeckoTheGecko May 4, 2026
e6aa98e
Add _Status for coords in convert.py
VeckoTheGecko May 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion src/parcels/_datasets/remote.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand Down
61 changes: 38 additions & 23 deletions src/parcels/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

from __future__ import annotations

import enum
import typing
import warnings
from typing import cast
Expand All @@ -25,12 +26,21 @@
if typing.TYPE_CHECKING:
import uxarray as ux

_NEMO_EXPECTED_COORDS = [
"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.
Copy link
Copy Markdown
Contributor Author

@VeckoTheGecko VeckoTheGecko May 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note

Refactor the whole convert module, this can surely all be handled better.

Is still on my TODO list (though quite low down). I have ideas on how we can make the conversion very robust but I want that to come later.

from parcels._typing import XgcmAxisDirection


class _Status(enum.Enum):
REQUIRED = enum.auto()
OPTIONAL = enum.auto()

_NEMO_DIMENSION_COORD_NAMES = [

_NEMO_EXPECTED_COORDS: list[tuple[str, _Status]] = [
("glamf", _Status.REQUIRED),
("gphif", _Status.REQUIRED),
("depthw", _Status.OPTIONAL),
]

_NEMO_DIMENSION_COORD_NAMES: list[str] = [
"x",
"y",
"time",
Expand All @@ -44,7 +54,7 @@
"gphif",
]

_NEMO_AXIS_VARNAMES = {
_NEMO_AXIS_VARNAMES: dict[str, XgcmAxisDirection] = {
"x": "X",
"x_center": "X",
"y": "Y",
Expand All @@ -54,17 +64,17 @@
"time": "T",
}

_NEMO_VARNAMES_MAPPING = {
_NEMO_VARNAMES_MAPPING: dict[str, str] = {
"time_counter": "time",
"depthw": "depth",
"uo": "U",
"vo": "V",
"wo": "W",
}

_MITGCM_EXPECTED_COORDS = ["XG", "YG", "Zl"]
_MITGCM_EXPECTED_COORDS: list[tuple[str, _Status]] = [(name, _Status.REQUIRED) for name in ["XG", "YG", "Zl"]]

_MITGCM_AXIS_VARNAMES = {
_MITGCM_AXIS_VARNAMES: dict[str, XgcmAxisDirection] = {
"XC": "X",
"XG": "X",
"Xp1": "X",
Expand All @@ -79,39 +89,42 @@
"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[tuple[str, _Status]] = [
(name, _Status.REQUIRED) for name in ["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",
}


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)


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"),
Expand All @@ -129,7 +142,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]})
Expand All @@ -138,31 +151,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])
Expand All @@ -176,7 +191,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
Expand Down
60 changes: 55 additions & 5 deletions tests/test_convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 == {
Expand All @@ -41,6 +48,49 @@ def test_nemo_to_sgrid():
meta.get_value_by_id("node_dimension2"), # Y edge
}.issubset(set(ds["V"].dims))

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")
Expand Down
Loading