Merge pull request #2499 from Parcels-code/feature/ux_from_conventions

feature/ux_from_conventions

Author: fluidnumericsJoe

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

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)

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)
 
 

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,15 +182,16 @@ 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
     # spacing of 100m - shifting by 1/10 of a degree laterally and 10m vertically should keep us
     # 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