Adding XNearest for nearest neighbour interpolation

Author: erikvansebille

parcels/application_kernels/interpolation.py

@@ -17,6 +17,7 @@
     "UXPiecewiseConstantFace",
     "UXPiecewiseLinearNode",
     "XLinear",
+    "XNearest",
     "ZeroInterpolator",
 ]
 
@@ -111,6 +112,69 @@ def XLinear(
     return value.compute() if isinstance(value, dask.Array) else value
 
 
+def XNearest(
+    field: Field,
+    ti: int,
+    position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]],
+    tau: np.float32 | np.float64,
+    t: np.float32 | np.float64,
+    z: np.float32 | np.float64,
+    y: np.float32 | np.float64,
+    x: np.float32 | np.float64,
+):
+    """
+    Nearest-Neighbour spatial interpolation on a regular grid.
+    Note that this still uses linear interpolation in time.
+    """
+    xi, xsi = position["X"]
+    yi, eta = position["Y"]
+    zi, zeta = position["Z"]
+
+    axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
+    data = field.data
+
+    lenT = 2 if np.any(tau > 0) else 1
+
+    # Spatial coordinates: left if barycentric < 0.5, otherwise right
+    zi_1 = np.clip(zi + 1, 0, data.shape[1] - 1)
+    zi_full = np.where(zeta < 0.5, zi, zi_1)
+
+    yi_1 = np.clip(yi + 1, 0, data.shape[2] - 1)
+    yi_full = np.where(eta < 0.5, yi, yi_1)
+
+    xi_1 = np.clip(xi + 1, 0, data.shape[3] - 1)
+    xi_full = np.where(xsi < 0.5, xi, xi_1)
+
+    # Time coordinates: 1 point at ti, then 1 point at ti+1
+    if lenT == 1:
+        ti_full = ti
+    else:
+        ti_1 = np.clip(ti + 1, 0, data.shape[0] - 1)
+        ti_full = np.concatenate([ti, ti_1])
+        xi_full = np.repeat(xi_full, 2)
+        yi_full = np.repeat(yi_full, 2)
+        zi_full = np.repeat(zi_full, 2)
+
+    # Create DataArrays for indexing
+    selection_dict = {
+        axis_dim["X"]: xr.DataArray(xi_full, dims=("points")),
+        axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")),
+    }
+    if "Z" in axis_dim:
+        selection_dict[axis_dim["Z"]] = xr.DataArray(zi_full, dims=("points"))
+    if "time" in data.dims:
+        selection_dict["time"] = xr.DataArray(ti_full, dims=("points"))
+
+    corner_data = data.isel(selection_dict).data.reshape(lenT, len(xsi))
+
+    if lenT == 2:
+        value = corner_data[0, :] * (1 - tau) + corner_data[1, :] * tau
+    else:
+        value = corner_data[0, :]
+
+    return value.compute() if isinstance(value, dask.Array) else value
+
+
 def UXPiecewiseConstantFace(
     field: Field,
     ti: int,

tests/v4/test_interpolation.py

@@ -6,7 +6,7 @@
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels._index_search import _search_time_index
 from parcels.application_kernels.advection import AdvectionRK4_3D
-from parcels.application_kernels.interpolation import UXPiecewiseLinearNode, XLinear, ZeroInterpolator
+from parcels.application_kernels.interpolation import UXPiecewiseLinearNode, XLinear, XNearest, ZeroInterpolator
 from parcels.field import Field, VectorField
 from parcels.fieldset import FieldSet
 from parcels.particle import Particle, Variable
@@ -49,9 +49,25 @@ def field():
     "func, t, z, y, x, expected",
     [
         pytest.param(ZeroInterpolator, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 0, id="Zero"),
-        pytest.param(XLinear, np.timedelta64(0, "s"), 0, 0.49, 0.51, 1.49, id="Linear-1"),
-        pytest.param(XLinear, np.timedelta64(1, "s"), 0, 0.49, 0.51, 6.49, id="Linear-2"),
+        pytest.param(
+            XLinear,
+            [np.timedelta64(0, "s"), np.timedelta64(1, "s")],
+            [0, 0],
+            [0.49, 0.49],
+            [0.51, 0.51],
+            [1.49, 6.49],
+            id="Linear",
+        ),
         pytest.param(XLinear, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 13.99, id="Linear-2"),
+        pytest.param(
+            XNearest,
+            [np.timedelta64(0, "s"), np.timedelta64(3, "s")],
+            [0.2, 0.2],
+            [0.2, 0.2],
+            [0.51, 0.51],
+            [1.0, 16.0],
+            id="Nearest",
+        ),
     ],
 )
 def test_raw_2d_interpolation(field, func, t, z, y, x, expected):