From 0d751cbbf74f405cb6854fcf79ffa95085215c09 Mon Sep 17 00:00:00 2001 From: Joe Schoonover Date: Sat, 31 Jan 2026 16:39:39 -0500 Subject: [PATCH] Switch fesom and icon ingestion over to convert methodology; update tests --- src/parcels/_core/fieldset.py | 60 +---------- src/parcels/convert.py | 189 +++++++++++++++++++++++++++++++++ tests/test_fieldset.py | 16 +-- tests/test_uxarray_fieldset.py | 27 ++--- 4 files changed, 207 insertions(+), 85 deletions(-) diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index 4dfbffe17..fdb47f8e9 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -183,6 +183,9 @@ def gridset(self) -> list[BaseGrid]: @classmethod def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"): """Create a FieldSet from a Parcels compliant uxarray.UxDataset. + + This is the primary ingestion method in Parcels for structured grid datasets. + The main requirements for a uxDataset are naming conventions for vertical grid dimensions & coordinates zf - Name for coordinate and dimension for vertical positions at layer interfaces @@ -225,63 +228,6 @@ def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"): return cls(list(fields.values())) - @classmethod - def from_fesom2(cls, ds: ux.UxDataset, mesh: str = "spherical"): - """Create a FieldSet from a FESOM2 uxarray.UxDataset. - - Parameters - ---------- - ds : uxarray.UxDataset - uxarray.UxDataset as obtained from the uxarray package. - - Returns - ------- - FieldSet - FieldSet object containing the fields from the dataset that can be used for a Parcels simulation. - """ - ds = ds.copy() - ds_dims = list(ds.dims) - if not all(dim in ds_dims for dim in ["time", "nz", "nz1"]): - raise ValueError( - f"Dataset missing one of the required dimensions 'time', 'nz', or 'nz1' for FESOM data. Found dimensions {ds_dims}" - ) - ds = ds.rename( - { - "nz": "zf", # Vertical Interface - "nz1": "zc", # Vertical Center - } - ).set_index(zf="zf", zc="zc") - - return FieldSet.from_ugrid_conventions(ds, mesh=mesh) - - @classmethod - def from_icon(cls, ds: ux.UxDataset, mesh: str = "spherical"): - """Create a FieldSet from a ICON uxarray.UxDataset. - - Parameters - ---------- - ds : uxarray.UxDataset - uxarray.UxDataset as obtained from the uxarray package. - - Returns - ------- - FieldSet - FieldSet object containing the fields from the dataset that can be used for a Parcels simulation. - """ - ds = ds.copy() - ds_dims = list(ds.dims) - if not all(dim in ds_dims for dim in ["time", "depth", "depth_2"]): - raise ValueError( - f"Dataset missing one of the required dimensions 'time', 'depth', or 'depth_2' for ICON data. Found dimensions {ds_dims}" - ) - ds = ds.rename( - { - "depth_2": "zf", # Vertical Interface - "depth": "zc", # Vertical Center - } - ).set_index(zf="zf", zc="zc") - return FieldSet.from_ugrid_conventions(ds, mesh=mesh) - @classmethod def from_sgrid_conventions( cls, ds: xr.Dataset, mesh: Mesh | None = None diff --git a/src/parcels/convert.py b/src/parcels/convert.py index dacf22d58..e8f2b84b9 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -529,3 +529,192 @@ def copernicusmarine_to_sgrid( ) return ds + + +# Known vertical dimension mappings by model +_FESOM2_VERTICAL_DIMS = {"interface": "nz", "center": "nz1"} +_ICON_VERTICAL_DIMS = {"interface": "depth_2", "center": "depth"} + + +def _detect_vertical_coordinates( + ds: ux.UxDataset, + known_mappings: dict[str, str] | None = None, +) -> tuple[str, str]: + """Detect vertical coordinate dimensions for faces (zf) and centers (zc). + + Detection strategy (with fallback): + 1. Use known_mappings if provided and dimensions exist + 2. Look for CF convention axis='Z' metadata + 3. Find dimension pairs where sizes differ by exactly 1 + + Parameters + ---------- + ds : ux.UxDataset + UxDataset to analyze for vertical coordinates. + known_mappings : dict[str, str] | None + Optional mapping with keys "interface" and "center" specifying + the dimension names for layer interfaces (zf) and centers (zc). + + Returns + ------- + tuple[str, str] + Tuple of (interface_dim_name, center_dim_name). + + Raises + ------ + ValueError + If vertical coordinates cannot be detected. + """ + ds_dims = set(ds.dims) + + # Strategy 1: Use known mappings if provided and dimensions exist + if known_mappings is not None: + interface_dim = known_mappings.get("interface") + center_dim = known_mappings.get("center") + if interface_dim in ds_dims and center_dim in ds_dims: + logger.info( + f"Using known vertical dimension mapping: {interface_dim!r} (interfaces) and {center_dim!r} (centers)." + ) + return interface_dim, center_dim + logger.debug(f"Known mappings {known_mappings} not found in dataset dimensions {ds_dims}. Trying CF metadata.") + + # Strategy 2: Look for CF convention axis='Z' metadata + z_dims = [] + for dim in ds_dims: + if dim in ds.coords: + coord = ds.coords[dim] + if coord.attrs.get("axis") == "Z": + z_dims.append(dim) + elif coord.attrs.get("positive") in ("up", "down"): + z_dims.append(dim) + elif "depth" in coord.attrs.get("standard_name", "").lower(): + z_dims.append(dim) + + if len(z_dims) == 2: + # Sort by size - interface has n+1 values, center has n + z_dims_sorted = sorted(z_dims, key=lambda d: ds.sizes[d], reverse=True) + interface_dim, center_dim = z_dims_sorted + if ds.sizes[interface_dim] == ds.sizes[center_dim] + 1: + logger.info( + f"Detected vertical dimensions from CF metadata: {interface_dim!r} (interfaces) and {center_dim!r} (centers)." + ) + return interface_dim, center_dim + + # Strategy 3: Find dimension pairs where sizes differ by exactly 1 + # Skip known non-vertical dimensions + skip_dims = {"time", "n_face", "n_node", "n_edge", "n_max_face_nodes"} + candidate_dims = [d for d in ds_dims if d not in skip_dims] + + for dim1 in candidate_dims: + for dim2 in candidate_dims: + if dim1 != dim2: + size1, size2 = ds.sizes[dim1], ds.sizes[dim2] + if size1 == size2 + 1: + logger.info( + f"Auto-detected vertical dimensions by size difference: {dim1!r} (interfaces, size={size1}) " + f"and {dim2!r} (centers, size={size2})." + ) + return dim1, dim2 + + raise ValueError( + f"Could not detect vertical coordinate dimensions in dataset with dims {list(ds_dims)}. " + "Please ensure the dataset has vertical layer interface and center dimensions, " + "or rename them manually to 'zf' (interfaces) and 'zc' (centers)." + ) + + +def _rename_vertical_dims( + ds: ux.UxDataset, + interface_dim: str, + center_dim: str, +) -> ux.UxDataset: + """Rename vertical dimensions to zf (interfaces) and zc (centers). + + Parameters + ---------- + ds : ux.UxDataset + UxDataset with vertical dimensions to rename. + interface_dim : str + Current name of the interface dimension. + center_dim : str + Current name of the center dimension. + + Returns + ------- + ux.UxDataset + Dataset with renamed dimensions and indexed coordinates. + """ + rename_dict = {} + if interface_dim != "zf": + rename_dict[interface_dim] = "zf" + if center_dim != "zc": + rename_dict[center_dim] = "zc" + + if rename_dict: + logger.info(f"Renaming vertical dimensions: {rename_dict}") + ds = ds.rename(rename_dict) + + ds = ds.set_index(zf="zf", zc="zc") + return ds + + +def fesom_to_ugrid(ds: ux.UxDataset) -> ux.UxDataset: + """Convert FESOM2 UxDataset to Parcels UGRID-compliant format. + + Renames vertical dimensions: + - nz -> zf (vertical layer faces/interfaces) + - nz1 -> zc (vertical layer centers) + + Parameters + ---------- + ds : ux.UxDataset + FESOM2 UxDataset as obtained from uxarray. + + Returns + ------- + ux.UxDataset + UGRID-compliant dataset ready for FieldSet.from_ugrid_conventions(). + + Examples + -------- + >>> import uxarray as ux + >>> from parcels import FieldSet + >>> from parcels.convert import fesom_to_ugrid + >>> ds = ux.open_mfdataset(grid_path, data_path) + >>> ds_ugrid = fesom_to_ugrid(ds) + >>> fieldset = FieldSet.from_ugrid_conventions(ds_ugrid, mesh="flat") + """ + ds = ds.copy() + interface_dim, center_dim = _detect_vertical_coordinates(ds, _FESOM2_VERTICAL_DIMS) + return _rename_vertical_dims(ds, interface_dim, center_dim) + + +def icon_to_ugrid(ds: ux.UxDataset) -> ux.UxDataset: + """Convert ICON UxDataset to Parcels UGRID-compliant format. + + Renames vertical dimensions: + - depth_2 -> zf (vertical layer faces/interfaces) + - depth -> zc (vertical layer centers) + + Parameters + ---------- + ds : ux.UxDataset + ICON UxDataset as obtained from uxarray. + + Returns + ------- + ux.UxDataset + UGRID-compliant dataset ready for FieldSet.from_ugrid_conventions(). + + Examples + -------- + >>> import uxarray as ux + >>> from parcels import FieldSet + >>> from parcels.convert import icon_to_ugrid + >>> ds = ux.open_mfdataset(grid_path, data_path) + >>> ds_ugrid = icon_to_ugrid(ds) + >>> fieldset = FieldSet.from_ugrid_conventions(ds_ugrid, mesh="flat") + """ + ds = ds.copy() + interface_dim, center_dim = _detect_vertical_coordinates(ds, _ICON_VERTICAL_DIMS) + return _rename_vertical_dims(ds, interface_dim, center_dim) diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index cdec96601..8056d4617 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -6,7 +6,7 @@ import pytest import xarray as xr -from parcels import Field, ParticleFile, ParticleSet, VectorField, XGrid +from parcels import Field, ParticleFile, ParticleSet, VectorField, XGrid, convert from parcels._core.fieldset import CalendarError, FieldSet, _datetime_to_msg from parcels._datasets.structured.generic import T as T_structured from parcels._datasets.structured.generic import datasets as datasets_structured @@ -243,33 +243,33 @@ def test_fieldset_add_field_after_pset(): def test_fieldset_from_icon(): - ds = datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"] - fieldset = FieldSet.from_icon(ds) + ds = convert.icon_to_ugrid(datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"]) + fieldset = FieldSet.from_ugrid_conventions(ds) assert "U" in fieldset.fields assert "V" in fieldset.fields assert "UVW" in fieldset.fields def test_fieldset_from_fesom2(): - ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"] - fieldset = FieldSet.from_fesom2(ds) + ds = convert.fesom_to_ugrid(datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"]) + fieldset = FieldSet.from_ugrid_conventions(ds) assert "U" in fieldset.fields assert "V" in fieldset.fields assert "UVW" in fieldset.fields def test_fieldset_from_fesom2_missingUV(): - ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"] + ds = convert.fesom_to_ugrid(datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"]) # Intentionally create a dataset that is missing the U field localds = ds.rename({"U": "notU"}) with pytest.raises(ValueError) as info: - _ = FieldSet.from_fesom2(localds) + _ = FieldSet.from_ugrid_conventions(localds) assert "Dataset has only one of the two variables 'U' and 'V'" in str(info) # Intentionally create a dataset that is missing the V field localds = ds.rename({"V": "notV"}) with pytest.raises(ValueError) as info: - _ = FieldSet.from_fesom2(localds) + _ = FieldSet.from_ugrid_conventions(localds) assert "Dataset has only one of the two variables 'U' and 'V'" in str(info) diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py index 8d68c58ac..001c411c6 100644 --- a/tests/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -12,6 +12,7 @@ download_example_dataset, ) from parcels._datasets.unstructured.generic import datasets as datasets_unstructured +from parcels.convert import fesom_to_ugrid, icon_to_ugrid from parcels.interpolators import ( Ux_Velocity, UxConstantFaceConstantZC, @@ -29,12 +30,7 @@ def ds_fesom_channel() -> ux.UxDataset: f"{fesom_path}/w.fesom_channel.nc", ] ds = ux.open_mfdataset(grid_path, data_path).rename_vars({"u": "U", "v": "V", "w": "W"}) - ds = ds.rename( - { - "nz": "zf", # Vertical Interface - "nz1": "zc", # Vertical Center - } - ).set_index(zf="zf", zc="zc") + ds = fesom_to_ugrid(ds) return ds @@ -121,12 +117,7 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): Since the underlying data is constant, we can check that the values are as expected. """ ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"] - ds = ds.rename( - { - "nz": "zf", # Vertical Interface - "nz1": "zc", # Vertical Center - } - ).set_index(zf="zf", zc="zc") + ds = fesom_to_ugrid(ds) grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh="flat") UVW = VectorField( name="UVW", @@ -174,12 +165,7 @@ def test_fesom2_square_delaunay_antimeridian_eval(): Since the underlying data is constant, we can check that the values are as expected. """ ds = datasets_unstructured["fesom2_square_delaunay_antimeridian"] - ds = ds.rename( - { - "nz": "zf", # Vertical Interface - "nz1": "zc", # Vertical Center - } - ).set_index(zf="zf", zc="zc") + ds = fesom_to_ugrid(ds) P = Field( name="p", data=ds.p, @@ -196,7 +182,8 @@ def test_fesom2_square_delaunay_antimeridian_eval(): def test_icon_evals(): ds = datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"].copy(deep=True) - fieldset = FieldSet.from_icon(ds, mesh="flat") + ds = icon_to_ugrid(ds) + fieldset = FieldSet.from_ugrid_conventions(ds, mesh="flat") # Query points, are chosen to be just a fraction off from the center of a cell for testing # This generic dataset has an effective lateral grid-spacing of 3 degrees and vertical grid @@ -204,7 +191,7 @@ def test_icon_evals(): # within the cell and make for easy exactness checking of constant and linear interpolation xc = ds.uxgrid.face_lon.values yc = ds.uxgrid.face_lat.values - zc = 0.0 * xc + ds.depth.values[1] # Make zc the same length as xc + zc = 0.0 * xc + ds.zc.values[1] # Make zc the same length as xc tq = 0.0 * xc xq = xc + 0.1