Uxgrid edge data interpolators

src/parcels/_core/model.py

@@ -28,6 +28,8 @@
 from parcels.interpolators import (
     CGrid_Velocity,
     Ux_Velocity,
+    UxConstantEdgeConstantZC,
+    UxConstantEdgeLinearZF,
     UxConstantFaceConstantZC,
     UxConstantFaceLinearZF,
     UxLinearNodeConstantZC,
@@ -517,10 +519,14 @@ def _select_uxinterpolator(da: ux.UxDataArray):
         "zc,n_face": UxConstantFaceConstantZC,
         # (zc,n_node): node/corner laterally, layer centers vertically — barycentric lateral & piecewise constant vertical
         "zc,n_node": UxLinearNodeConstantZC,
+        # (zc,n_edge): edge-center laterally, layer centres vertically — piecewise constant
+        "zc,n_edge": UxConstantEdgeConstantZC,
         # (zf,n_node): node/corner laterally, layer interfaces vertically — barycentric lateral & linear vertical
         "zf,n_node": UxLinearNodeLinearZF,
         # (zf,n_face): face-center laterally, layer interfaces vertically — piecewise constant lateral & linear vertical
         "zf,n_face": UxConstantFaceLinearZF,
+        # (zf,n_edge): edge-center laterally, layer interfaces vertically — piecewise constant lateral & linear vertical
+        "zf,n_edge": UxConstantEdgeLinearZF,
     }
     # Extract only spatial dimensions, neglecting time
     da_spatial_dims = tuple(d for d in da.dims if d not in ("time",))
@@ -536,7 +542,7 @@ def _select_uxinterpolator(da: ux.UxDataArray):
     for d in da_spatial_dims:
         if d in ("zf", "zc"):
             vdim = d
-        if d in ("n_face", "n_node"):
+        if d in ("n_face", "n_node", "n_edge"):
             ldim = d
     # Map to supported interpolators
     if vdim and ldim:

src/parcels/interpolators/__init__.py

@@ -1,5 +1,7 @@
 from ._uxinterpolators import (
     Ux_Velocity,
+    UxConstantEdgeConstantZC,
+    UxConstantEdgeLinearZF,
     UxConstantFaceConstantZC,
     UxConstantFaceLinearZF,
     UxLinearNodeConstantZC,
@@ -32,6 +34,8 @@
     "Ux_Velocity",
     "UxConstantFaceConstantZC",
     "UxConstantFaceLinearZF",
+    "UxConstantEdgeConstantZC",
+    "UxConstantEdgeLinearZF",
     "UxLinearNodeConstantZC",
     "UxLinearNodeLinearZF",
 ]

src/parcels/interpolators/_uxinterpolators.py

@@ -155,6 +155,74 @@ def _zsample(z_index):
         return value.compute() if is_dask_collection(value) else value
 
 
+class UxConstantEdgeConstantZC(ScalarInterpolator):
+    """
+    Piecewise constant interpolation kernel for edge registered data that is vertically centered (on zc points).
+    Face values are taken to be the mean of the surrounding edges.
+    """
+
+    def interp(
+        self,
+        particle_positions: dict[str, float | np.ndarray],
+        grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]],
+        field: Field,
+    ):
+        # Broadcast the per-axis indices to a common (npart,) shape (``ti`` may be scalar for time-constant fields)
+        ti, zi, fi = np.broadcast_arrays(
+            grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"]
+        )
+        edge_ids = field.grid.uxgrid.face_edge_connectivity[fi, :].values
+
+        tdim, zdim, edim = field.data.dims
+        selection_dict = {
+            tdim: xr.DataArray(ti, dims="points"),
+            zdim: xr.DataArray(zi, dims="points"),
+            edim: xr.DataArray(edge_ids, dims=("points", "edges")),
+        }
+        edge_data = field.data.isel(selection_dict, ignore_grid=True)
+        value = edge_data.mean("edges").data  # Face value is the mean of its surrounding edges
+        return value.compute() if is_dask_collection(value) else value
+
+
+class UxConstantEdgeLinearZF(ScalarInterpolator):
+    """
+    Piecewise constant interpolation (lateral) with linear vertical interpolation kernel for edge registered data
+    that is located at vertical interface levels (on zf points).
+    Face values are taken to be the mean of the surrounding edges.
+    """
+
+    def interp(
+        self,
+        particle_positions: dict[str, float | np.ndarray],
+        grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]],
+        field: Field,
+    ):
+        # Broadcast the per-axis indices to a common (npart,) shape (``ti`` may be scalar for time-constant fields)
+        ti, zi, fi = np.broadcast_arrays(
+            grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"]
+        )
+        zeta = grid_positions["Z"]["bcoord"]
+        edge_ids = field.grid.uxgrid.face_edge_connectivity[fi, :].values
+
+        tdim, zdim, edim = field.data.dims
+
+        def _zsample(z_index):
+            """Mean of the edge values surrounding each face at a single vertical interface level."""
+            selection_dict = {
+                tdim: xr.DataArray(ti, dims="points"),
+                zdim: xr.DataArray(z_index, dims="points"),
+                edim: xr.DataArray(edge_ids, dims=("points", "edges")),
+            }
+            edge_data = field.data.isel(selection_dict, ignore_grid=True)
+            return edge_data.mean("edges").data
+
+        # zi is the vertical layer index, bounded by the interface levels zi and zi + 1
+        fzk = _zsample(zi)
+        fzkp1 = _zsample(zi + 1)
+        value = fzk * (1 - zeta) + fzkp1 * zeta
+        return value.compute() if is_dask_collection(value) else value
+
+
 class Ux_Velocity(VectorInterpolator):  # noqa: N801
     """Interpolation kernel for Vectorfields of velocity on a UxGrid."""
 

tests/test_uxarray_fieldset.py

@@ -13,10 +13,13 @@
 )
 from parcels._core.mesh import FlatMesh, SphericalMesh
 from parcels._core.model import UnstructuredModelData
-from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
+from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset, simple_small_delaunay
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.interpolators import (
+    UxConstantEdgeConstantZC,
+    UxConstantEdgeLinearZF,
     UxConstantFaceConstantZC,
+    UxConstantFaceLinearZF,
     UxLinearNodeConstantZC,
     UxLinearNodeLinearZF,
 )
@@ -156,6 +159,53 @@ def test_icon_evals():
     assert np.allclose(fieldset.p.eval(t=tq, z=zq, y=yq, x=xq), zc * (xq + yq))
 
 
+def test_linear_edge_velocity_evals():
+    ds = simple_small_delaunay()
+    xe = ds.uxgrid.edge_lon.values
+    ye = ds.uxgrid.edge_lat.values
+    for name, zdim, value in (("U", "zc", xe), ("V", "zc", ye), ("W", "zf", xe + ye)):
+        ds[name] = ux.UxDataArray(
+            np.full((ds.sizes["time"], ds.sizes[zdim], ds.uxgrid.n_edge), value),
+            dims=["time", zdim, "n_edge"],
+            coords={"time": ds.time, zdim: ds[zdim]},
+            uxgrid=ds.uxgrid,
+        )
+    fieldset = FieldSet.from_ugrid_conventions(ds)
+
+    xc = ds.uxgrid.face_lon.values
+    yc = ds.uxgrid.face_lat.values
+    zq = np.full(xc.shape, 300.0)
+    tq = np.zeros(xc.shape)
+
+    assert np.allclose(fieldset.U.eval(t=tq, z=zq, y=yc, x=xc), xc)
+    assert np.allclose(fieldset.V.eval(t=tq, z=zq, y=yc, x=xc), yc)
+    assert np.allclose(fieldset.W.eval(t=tq, z=zq, y=yc, x=xc), xc + yc)
+
+
+def test_from_ugrid_conventions_selects_correct_interpolator():
+    """Fields get the default interpolator that matches their vertical and lateral registration."""
+    ds = convert.icon_to_ugrid(datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"].copy(deep=True))
+    expected_interpolators = {
+        ("zc", "n_face"): UxConstantFaceConstantZC,
+        ("zf", "n_face"): UxConstantFaceLinearZF,
+        ("zc", "n_node"): UxLinearNodeConstantZC,
+        ("zf", "n_node"): UxLinearNodeLinearZF,
+        ("zc", "n_edge"): UxConstantEdgeConstantZC,
+        ("zf", "n_edge"): UxConstantEdgeLinearZF,
+    }
+    for vertical_dim, lateral_dim in expected_interpolators:
+        ds[f"{lateral_dim}_{vertical_dim}"] = ux.UxDataArray(
+            np.ones((ds.sizes["time"], ds.sizes[vertical_dim], getattr(ds.uxgrid, lateral_dim))),
+            dims=["time", vertical_dim, lateral_dim],
+            coords={"time": ds.time, vertical_dim: ds[vertical_dim]},
+            uxgrid=ds.uxgrid,
+        )
+    fieldset = FieldSet.from_ugrid_conventions(ds)
+
+    for (vertical_dim, lateral_dim), expected_interpolator in expected_interpolators.items():
+        assert isinstance(getattr(fieldset, f"{lateral_dim}_{vertical_dim}").interp_method, expected_interpolator)
+
+
 _NESTEDGRIDS_NODES = np.array(
     [
         [10.0, 15.0],