From 1e3c1e760f6fc313677277ac24f104d94b096264 Mon Sep 17 00:00:00 2001 From: Wyatt Sieminski Date: Fri, 23 Oct 2026 19:22:54 +0000 Subject: [PATCH 1/4] Added uxinterpolators for data on n_edge --- src/parcels/_core/model.py | 8 ++- src/parcels/interpolators/__init__.py | 2 + src/parcels/interpolators/_uxinterpolators.py | 68 +++++++++++++++++++ 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index cb8594cff2..37eba789f0 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -30,6 +30,8 @@ Ux_Velocity, UxConstantFaceConstantZC, UxConstantFaceLinearZF, + UxConstantEdgeConstantZC, + UxConstantEdgeLinearZF, UxLinearNodeConstantZC, UxLinearNodeLinearZF, XLinear, @@ -500,10 +502,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",)) @@ -519,7 +525,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: diff --git a/src/parcels/interpolators/__init__.py b/src/parcels/interpolators/__init__.py index 4aa06526ed..bed31f8f76 100644 --- a/src/parcels/interpolators/__init__.py +++ b/src/parcels/interpolators/__init__.py @@ -2,6 +2,8 @@ Ux_Velocity, UxConstantFaceConstantZC, UxConstantFaceLinearZF, + UxConstantEdgeConstantZC, + UxConstantEdgeLinearZF, UxLinearNodeConstantZC, UxLinearNodeLinearZF, ) diff --git a/src/parcels/interpolators/_uxinterpolators.py b/src/parcels/interpolators/_uxinterpolators.py index 2666b57e91..6db1000eb7 100644 --- a/src/parcels/interpolators/_uxinterpolators.py +++ b/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.""" From 1ff4756b9dcb8dcd534301a5b6309ef39d3d54be Mon Sep 17 00:00:00 2001 From: Wyatt Sieminski Date: Fri, 23 Oct 2026 19:45:11 +0000 Subject: [PATCH 2/4] Added regression testing for default interpolator selection --- tests/test_uxarray_fieldset.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py index 01d765fd7e..cebee79aa9 100644 --- a/tests/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -16,7 +16,11 @@ from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset from parcels._datasets.unstructured.generic import datasets as datasets_unstructured from parcels.interpolators import ( + UxConstantEdgeConstantZC, + UxConstantEdgeLinearZF, UxConstantFaceConstantZC, + UxConstantFaceLinearZF, + UxLinearNodeConstantZC, UxLinearNodeLinearZF, ) from tests.utils import create_uxgrid_from_triangulation @@ -156,6 +160,30 @@ def test_icon_evals(): assert np.allclose(fieldset.p.eval(t=tq, z=zq, y=yq, x=xq), zc * (xq + yq)) +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], From f0bef6855aadc111e9604c4856a003f0ff6f9156 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:54:53 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/parcels/_core/model.py | 4 ++-- src/parcels/interpolators/__init__.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 37eba789f0..bc50017a68 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -28,10 +28,10 @@ from parcels.interpolators import ( CGrid_Velocity, Ux_Velocity, - UxConstantFaceConstantZC, - UxConstantFaceLinearZF, UxConstantEdgeConstantZC, UxConstantEdgeLinearZF, + UxConstantFaceConstantZC, + UxConstantFaceLinearZF, UxLinearNodeConstantZC, UxLinearNodeLinearZF, XLinear, diff --git a/src/parcels/interpolators/__init__.py b/src/parcels/interpolators/__init__.py index bed31f8f76..eaeb923c3f 100644 --- a/src/parcels/interpolators/__init__.py +++ b/src/parcels/interpolators/__init__.py @@ -1,9 +1,9 @@ from ._uxinterpolators import ( Ux_Velocity, - UxConstantFaceConstantZC, - UxConstantFaceLinearZF, UxConstantEdgeConstantZC, UxConstantEdgeLinearZF, + UxConstantFaceConstantZC, + UxConstantFaceLinearZF, UxLinearNodeConstantZC, UxLinearNodeLinearZF, ) From c246edac1ec1f48efa057c4e4e7a8c351e023661 Mon Sep 17 00:00:00 2001 From: Wyatt Sieminski Date: Fri, 23 Oct 2026 20:32:43 +0000 Subject: [PATCH 4/4] Added new interpolators to __all__ section of interpolator __init__ --- src/parcels/interpolators/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/parcels/interpolators/__init__.py b/src/parcels/interpolators/__init__.py index bed31f8f76..69200c017b 100644 --- a/src/parcels/interpolators/__init__.py +++ b/src/parcels/interpolators/__init__.py @@ -34,6 +34,8 @@ "Ux_Velocity", "UxConstantFaceConstantZC", "UxConstantFaceLinearZF", + "UxConstantEdgeConstantZC", + "UxConstantEdgeLinearZF", "UxLinearNodeConstantZC", "UxLinearNodeLinearZF", ]