diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index cb8594cff..bc50017a6 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -28,6 +28,8 @@ from parcels.interpolators import ( CGrid_Velocity, Ux_Velocity, + UxConstantEdgeConstantZC, + UxConstantEdgeLinearZF, UxConstantFaceConstantZC, UxConstantFaceLinearZF, UxLinearNodeConstantZC, @@ -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 4aa06526e..2f131c546 100644 --- a/src/parcels/interpolators/__init__.py +++ b/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", ] diff --git a/src/parcels/interpolators/_uxinterpolators.py b/src/parcels/interpolators/_uxinterpolators.py index 2666b57e9..6db1000eb 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.""" diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py index 01d765fd7..cebee79aa 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],