Update NEMO ingestion code (#2549)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: VeckoTheGecko

src/parcels/convert.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import typing
+import warnings
 from typing import cast
 
 import numpy as np
@@ -24,16 +25,32 @@
 if typing.TYPE_CHECKING:
     import uxarray as ux
 
-_NEMO_EXPECTED_COORDS = ["glamf", "gphif"]
-
-_NEMO_DIMENSION_COORD_NAMES = ["x", "y", "time", "x", "x_center", "y", "y_center", "depth", "glamf", "gphif"]
+_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.
+
+_NEMO_DIMENSION_COORD_NAMES = [
+    "x",
+    "y",
+    "time",
+    "x",
+    "x_center",
+    "y",
+    "y_center",
+    "depth",
+    "depth_center",
+    "glamf",
+    "gphif",
+]
 
 _NEMO_AXIS_VARNAMES = {
     "x": "X",
     "x_center": "X",
     "y": "Y",
     "y_center": "Y",
     "depth": "Z",
+    "depth_center": "Z",
     "time": "T",
 }
 
@@ -95,16 +112,18 @@ def _pick_expected_coords(coords: xr.Dataset, expected_coord_names: list[str]) -
 
 
 def _maybe_bring_other_depths_to_depth(ds):
-    if "depth" in ds.coords:
-        for var in ds.data_vars:
-            for old_depth in ["depthu", "depthv", "deptht", "depthw"]:
-                if old_depth in ds[var].dims:
-                    ds[var] = ds[var].assign_coords(**{old_depth: ds["depth"].values}).rename({old_depth: "depth"})
-    return ds
-
+    for var in ds.data_vars:
+        for old_depth, target in [
+            ("depthu", "depth_center"),
+            ("depthv", "depth_center"),
+            ("deptht", "depth_center"),
+            ("depthw", "depth"),
+        ]:
+            if old_depth in ds[var].dims:
+                ds[var] = ds[var].rename({old_depth: target})
 
-def _maybe_create_depth_dim(ds):
     if "depth" not in ds.dims:
+        warnings.warn("No depth dimension found in your dataset. Assuming no depth (i.e., surface data).", stacklevel=1)
         ds = ds.expand_dims({"depth": [0]})
         ds["depth"] = xr.DataArray([0], dims=["depth"])
     return ds
@@ -286,17 +305,18 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da
         if coords.sizes["time"] != 1:
             raise ValueError("Time dimension in coords must be length 1 (i.e., no time-varying grid).")
         coords = coords.isel(time=0).drop("time")
-    if len(coords.dims) == 3:
+
+    if (
+        len(coords.dims) == 3
+    ):  #! This should really be looking at the dimensionality of the lons and lats arrays. Currently having 2D lon lat and 1D depth triggers this `if` clause
         for dim, len_ in coords.sizes.items():
             if len_ == 1:
                 # TODO: log statement about selecting along z dim of 1
                 coords = coords.isel({dim: 0})
-    if len(coords.dims) != 2:
-        raise ValueError("Expected coordsinates to be 2 dimensional")
-
+    # if len(coords.dims) != 2: #! This should really be looking at the dimensionality of the lons and lats arrays. Currently having 2D lon lat and 1D depth triggers this `if` clause
+    #     raise ValueError("Expected coordinates to be 2 dimensional")
     ds = xr.merge(list(fields.values()) + [coords])
     ds = _maybe_rename_variables(ds, _NEMO_VARNAMES_MAPPING)
-    ds = _maybe_create_depth_dim(ds)
     ds = _maybe_bring_other_depths_to_depth(ds)
     ds = _drop_unused_dimensions_and_coords(ds, _NEMO_DIMENSION_COORD_NAMES)
     ds = _assign_dims_as_coords(ds, _NEMO_DIMENSION_COORD_NAMES)
@@ -332,7 +352,7 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da
                 sgrid.FaceNodePadding("x_center", "x", sgrid.Padding.LOW),
                 sgrid.FaceNodePadding("y_center", "y", sgrid.Padding.LOW),
             ),
-            vertical_dimensions=(sgrid.FaceNodePadding("z_center", "depth", sgrid.Padding.HIGH),),
+            vertical_dimensions=(sgrid.FaceNodePadding("depth_center", "depth", sgrid.Padding.HIGH),),
         ).to_attrs(),
     )
 
@@ -525,7 +545,7 @@ def copernicusmarine_to_sgrid(
                 sgrid.FaceNodePadding("x_center", "lon", sgrid.Padding.LOW),
                 sgrid.FaceNodePadding("y_center", "lat", sgrid.Padding.LOW),
             ),
-            vertical_dimensions=(sgrid.FaceNodePadding("z_center", "depth", sgrid.Padding.LOW),),
+            vertical_dimensions=(sgrid.FaceNodePadding("depth_center", "depth", sgrid.Padding.LOW),),
         ).to_attrs(),
     )
 

tests/test_convert.py

@@ -23,7 +23,7 @@ def test_nemo_to_sgrid():
         "node_dimensions": "x y",
         "face_dimensions": "x_center:x (padding:low) y_center:y (padding:low)",
         "node_coordinates": "lon lat",
-        "vertical_dimensions": "z_center:depth (padding:high)",
+        "vertical_dimensions": "depth_center:depth (padding:high)",
     }
 
     meta = sgrid.parse_grid_attrs(ds["grid"].attrs)