Add method to import from fesom2

src/parcels/_core/fieldset.py

@@ -6,13 +6,15 @@
 
 import cf_xarray  # noqa: F401
 import numpy as np
+import uxarray as ux
 import xarray as xr
 import xgcm
 
 from parcels._core.converters import Geographic, GeographicPolar
 from parcels._core.field import Field, VectorField
 from parcels._core.utils.time import get_datetime_type_calendar
 from parcels._core.utils.time import is_compatible as datetime_is_compatible
+from parcels._core.uxgrid import UxGrid
 from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS, XGrid
 from parcels._logger import logger
 from parcels._typing import Mesh
@@ -251,6 +253,50 @@ def from_copernicusmarine(ds: xr.Dataset):
 
         return FieldSet(list(fields.values()))
 
+    def from_fesom2(ds: ux.UxDataset):
+        """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'. Found dimensions {ds_dims}"
+            )
+        grid = UxGrid(ds.uxgrid, z=ds.coords["nz"])
+        ds = _discover_fesom2_U_and_V(ds)
+
+        fields = {}
+        if "U" in ds.data_vars and "V" in ds.data_vars:
+            fields["U"] = Field("U", ds["U"], grid)
+            fields["V"] = Field("V", ds["V"])
+            fields["U"].units = GeographicPolar()
+            fields["V"].units = Geographic()
+
+            if "W" in ds.data_vars:
+                ds["W"] -= ds[
+                    "W"
+                ]  # Negate W to convert from up positive to down positive (as that's the direction of positive z)
+                fields["W"] = Field("W", ds["W"], grid)
+                fields["UVW"] = VectorField("UVW", fields["U"], fields["V"], fields["W"])
+            else:
+                fields["UV"] = VectorField("UV", fields["U"], fields["V"])
+
+        for varname in set(ds.data_vars) - set(fields.keys()):
+            fields[varname] = Field(varname, ds[varname], grid)
+
+        return FieldSet(list(fields.values()))
+
 
 class CalendarError(Exception):  # TODO: Move to a parcels errors module
     """Exception raised when the calendar of a field is not compatible with the rest of the Fields. The user should ensure that they only add fields to a FieldSet that have compatible CFtime calendars."""
@@ -359,7 +405,42 @@ def _discover_copernicusmarine_U_and_V(ds: xr.Dataset) -> xr.Dataset:
     return ds
 
 
-def _ds_rename_using_standard_names(ds: xr.Dataset, name_dict: dict[str, str]) -> xr.Dataset:
+def _discover_fesom2_U_and_V(ds: ux.UxDataset) -> ux.UxDataset:
+    # Assumes that the dataset has U and V data
+
+    cf_UV_standard_name_fallbacks = [("unod", "vnod"), ("u", "v")]
+    cf_W_standard_name_fallbacks = ["w"]
+
+    if "W" not in ds:
+        for cf_standard_name_W in cf_W_standard_name_fallbacks:
+            if cf_standard_name_W in ds.cf.standard_names:
+                ds = _ds_rename_using_standard_names(ds, {cf_standard_name_W: "W"})
+                break
+
+    if "U" in ds and "V" in ds:
+        return ds  # U and V already present
+    elif "U" in ds or "V" in ds:
+        raise ValueError(
+            "Dataset has only one of the two variables 'U' and 'V'. Please rename the appropriate variable in your dataset to have both 'U' and 'V' for Parcels simulation."
+        )
+
+    for cf_standard_name_U, cf_standard_name_V in cf_UV_standard_name_fallbacks:
+        if cf_standard_name_U in ds.cf.standard_names:
+            if cf_standard_name_V not in ds.cf.standard_names:
+                raise ValueError(
+                    f"Dataset has variable with CF standard name {cf_standard_name_U!r}, "
+                    f"but not the matching variable with CF standard name {cf_standard_name_V!r}. "
+                    "Please rename the appropriate variables in your dataset to have both 'U' and 'V' for Parcels simulation."
+                )
+        else:
+            continue
+
+        ds = _ds_rename_using_standard_names(ds, {cf_standard_name_U: "U", cf_standard_name_V: "V"})
+        break
+    return ds
+
+
+def _ds_rename_using_standard_names(ds: xr.Dataset | ux.UxDataset, name_dict: dict[str, str]) -> xr.Dataset:
     for standard_name, rename_to in name_dict.items():
         name = ds.cf[standard_name].name
         ds = ds.rename({name: rename_to})

tests/test_fieldset.py

@@ -11,6 +11,7 @@
 from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models
 from parcels._datasets.structured.generic import T as T_structured
 from parcels._datasets.structured.generic import datasets as datasets_structured
+from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from tests import utils
 
 ds = datasets_structured["ds_2d_left"]
@@ -270,3 +271,11 @@ def test_fieldset_from_copernicusmarine_with_W(caplog):
     assert "UV" not in fieldset.fields
     assert "UVW" in fieldset.fields
     assert "renamed it to 'W'" in caplog.text
+
+
+def test_fieldset_from_fesom2():
+    ds = datasets_unstructured["stommel_gyre_delaunay"]
+    fieldset = FieldSet.from_fesom2(ds)
+    assert "U" in fieldset.fields
+    assert "V" in fieldset.fields
+    assert "UVW" in fieldset.fields