Porting interpolation unit test from v3 to v4

Author: erikvansebille

tests/test_interpolation.py

@@ -1,6 +1,4 @@
-import numpy as np
 import pytest
-import xarray as xr
 
 import parcels._interpolation as interpolation
 from tests.utils import create_fieldset_zeros_3d
@@ -31,50 +29,6 @@ def some_function():
     assert f() == g() == "test"
 
 
-def create_interpolation_data():
-    """Reference data used for testing interpolation.
-
-    Most interpolation will be focussed around index
-    (depth, lat, lon) = (zi, yi, xi) = (1, 1, 1) with ti=0.
-    """
-    z0 = np.array(  # each x is +1 from the previous, each y is +2 from the previous
-        [
-            [0.0, 1.0, 2.0, 3.0],
-            [2.0, 3.0, 4.0, 5.0],
-            [4.0, 5.0, 6.0, 7.0],
-            [6.0, 7.0, 8.0, 9.0],
-        ]
-    )
-    spatial_data = [z0, z0 + 3, z0 + 6, z0 + 9]  # each z is +3 from the previous
-    return xr.DataArray([spatial_data, spatial_data, spatial_data], dims=("time", "depth", "lat", "lon"))
-
-
-@pytest.fixture
-def data_2d():
-    """2D slice of the reference data at depth=0."""
-    return create_interpolation_data().isel(depth=0).values
-
-
-@pytest.mark.v4remove
-@pytest.mark.xfail(reason="GH1946")
-@pytest.mark.parametrize(
-    "func, eta, xsi, expected",
-    [
-        pytest.param(interpolation._nearest_2d, 0.49, 0.49, 3.0, id="nearest_2d-1"),
-        pytest.param(interpolation._nearest_2d, 0.49, 0.51, 4.0, id="nearest_2d-2"),
-        pytest.param(interpolation._nearest_2d, 0.51, 0.49, 5.0, id="nearest_2d-3"),
-        pytest.param(interpolation._nearest_2d, 0.51, 0.51, 6.0, id="nearest_2d-4"),
-        pytest.param(interpolation._tracer_2d, None, None, 6.0, id="tracer_2d"),
-    ],
-)
-def test_raw_2d_interpolation(data_2d, func, eta, xsi, expected):
-    """Test the 2D interpolation functions on the raw arrays."""
-    tau, ti = 0, 0
-    yi, xi = 1, 1
-    ctx = interpolation.InterpolationContext2D(data_2d, tau, eta, xsi, ti, yi, xi)
-    assert func(ctx) == expected
-
-
 @pytest.mark.v4remove
 @pytest.mark.xfail(reason="GH1946")
 @pytest.mark.usefixtures("tmp_interpolator_registry")

tests/v4/test_interpolation.py

@@ -4,8 +4,9 @@
 
 from parcels._datasets.structured.generated import simple_UV_dataset
 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
+from parcels.application_kernels.interpolation import UXPiecewiseLinearNode, XLinear, ZeroInterpolator
 from parcels.field import Field, VectorField
 from parcels.fieldset import FieldSet
 from parcels.particle import Particle, Variable
@@ -16,6 +17,52 @@
 from tests.utils import TEST_DATA
 
 
+@pytest.fixture
+def field():
+    """Reference data used for testing interpolation."""
+    z0 = np.array(  # each x is +1 from the previous, each y is +2 from the previous
+        [
+            [0.0, 1.0, 2.0, 3.0],
+            [2.0, 3.0, 4.0, 5.0],
+            [4.0, 5.0, 6.0, 7.0],
+            [6.0, 7.0, 8.0, 9.0],
+        ]
+    )
+    spatial_data = np.array([z0, z0 + 3, z0 + 6, z0 + 9])  # each z is +3 from the previous
+    temporal_data = np.array([spatial_data, spatial_data + 10, spatial_data + 20])  # each t is +10 from the previous
+
+    ds = xr.Dataset(
+        {"U": (["time", "depth", "lat", "lon"], temporal_data)},
+        coords={
+            "time": (["time"], [np.timedelta64(t, "s") for t in [0, 2, 4]], {"axis": "T"}),
+            "depth": (["depth"], [0, 1, 2, 3], {"axis": "Z"}),
+            "lat": (["lat"], [0, 1, 2, 3], {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "lon": (["lon"], [0, 1, 2, 3], {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "x": (["x"], [0.5, 1.5, 2.5, 3.5], {"axis": "X"}),
+            "y": (["y"], [0.5, 1.5, 2.5, 3.5], {"axis": "Y"}),
+        },
+    )
+    return Field("U", ds["U"], XGrid.from_dataset(ds))
+
+
+@pytest.mark.parametrize(
+    "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(1, "s"), 2.5, 0.49, 0.51, 13.99, id="Linear-2"),
+    ],
+)
+def test_raw_2d_interpolation(field, func, t, z, y, x, expected):
+    """Test the interpolation functions on the Field."""
+    tau, ti = _search_time_index(field, t)
+    position = field.grid.search(z, y, x)
+
+    value = func(field, ti, position, tau, 0, 0, y, x)
+    np.testing.assert_equal(value, expected)
+
+
 @pytest.mark.parametrize("mesh_type", ["spherical", "flat"])
 def test_interpolation_mesh_type(mesh_type, npart=10):
     ds = simple_UV_dataset(mesh_type=mesh_type)