Implementing new Interpolation API

Author: erikvansebille

src/parcels/_core/field.py

@@ -216,12 +216,13 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
         else:
             _ei = particles.ei[:, self.igrid]
 
-        tau, ti = _search_time_index(self, time)
-        position = self.grid.search(z, y, x, ei=_ei)
+        position = {"time": time, "z": z, "lat": y, "lon": x}
+        position["T"] = _search_time_index(self, time)
+        position.update(self.grid.search(z, y, x, ei=_ei))
         _update_particles_ei(particles, position, self)
         _update_particle_states_position(particles, position)
 
-        value = self._interp_method(self, ti, position, tau, time, z, y, x)
+        value = self._interp_method(position, self)
 
         _update_particle_states_interp_value(particles, value)
 
@@ -300,20 +301,21 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
         else:
             _ei = particles.ei[:, self.igrid]
 
-        tau, ti = _search_time_index(self.U, time)
-        position = self.grid.search(z, y, x, ei=_ei)
+        position = {"time": time, "z": z, "lat": y, "lon": x}
+        position["T"] = _search_time_index(self.U, time)
+        position.update(self.grid.search(z, y, x, ei=_ei))
         _update_particles_ei(particles, position, self)
         _update_particle_states_position(particles, position)
 
         if self._vector_interp_method is None:
-            u = self.U._interp_method(self.U, ti, position, tau, time, z, y, x)
-            v = self.V._interp_method(self.V, ti, position, tau, time, z, y, x)
+            u = self.U._interp_method(position, self.U)
+            v = self.V._interp_method(position, self.V)
             if "3D" in self.vector_type:
-                w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x)
+                w = self.W._interp_method(position, self.W)
             else:
                 w = 0.0
         else:
-            (u, v, w) = self._vector_interp_method(self, ti, position, tau, time, z, y, x)
+            (u, v, w) = self._vector_interp_method(position, self)
 
         if applyConversion:
             u = self.U.units.to_target(u, z, y, x)

src/parcels/_core/index_search.py

@@ -75,14 +75,14 @@ def _search_time_index(field: Field, time: datetime):
     if the sampled value is outside the time value range.
     """
     if field.time_interval is None:
-        return np.zeros(shape=time.shape, dtype=np.float32), np.zeros(shape=time.shape, dtype=np.int32)
+        return np.zeros(shape=time.shape, dtype=np.int32), np.zeros(shape=time.shape, dtype=np.float32)
 
     if not field.time_interval.is_all_time_in_interval(time):
         _raise_time_extrapolation_error(time, field=None)
 
     ti = np.searchsorted(field.data.time.data, time, side="right") - 1
     tau = (time - field.data.time.data[ti]) / (field.data.time.data[ti + 1] - field.data.time.data[ti])
-    return np.atleast_1d(tau), np.atleast_1d(ti)
+    return np.atleast_1d(ti), np.atleast_1d(tau)
 
 
 def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: np.ndarray):

src/parcels/interpolators.py

@@ -30,28 +30,16 @@
 
 
 def ZeroInterpolator(
-    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,
+    field: Field,
 ) -> np.float32 | np.float64:
     """Template function used for the signature check of the lateral interpolation methods."""
     return 0.0
 
 
 def ZeroInterpolator_Vector(
-    vectorfield: VectorField,
-    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,
+    vectorfield: VectorField,
 ) -> np.float32 | np.float64:
     """Template function used for the signature check of the interpolation methods for velocity fields."""
     return 0.0
@@ -105,19 +93,14 @@ def _get_corner_data_Agrid(
 
 
 def XLinear(
-    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,
+    field: Field,
 ):
     """Trilinear interpolation on a regular grid."""
     xi, xsi = position["X"]
     yi, eta = position["Y"]
     zi, zeta = position["Z"]
+    ti, tau = position["T"]
 
     axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
     data = field.data
@@ -149,14 +132,8 @@ def XLinear(
 
 
 def CGrid_Velocity(
-    vectorfield: VectorField,
-    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,
+    vectorfield: VectorField,
 ):
     """
     Interpolation kernel for velocity fields on a C-Grid.
@@ -166,6 +143,8 @@ def CGrid_Velocity(
     xi, xsi = position["X"]
     yi, eta = position["Y"]
     zi, zeta = position["Z"]
+    ti, tau = position["T"]
+    lon = position["lon"]
 
     U = vectorfield.U.data
     V = vectorfield.V.data
@@ -180,8 +159,8 @@ def CGrid_Velocity(
         py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]])
 
     if grid._mesh == "spherical":
-        px[0] = np.where(px[0] < x - 225, px[0] + 360, px[0])
-        px[0] = np.where(px[0] > x + 225, px[0] - 360, px[0])
+        px[0] = np.where(px[0] < lon - 225, px[0] + 360, px[0])
+        px[0] = np.where(px[0] > lon + 225, px[0] - 360, px[0])
         px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:])
         px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:])
     c1 = i_u._geodetic_distance(
@@ -296,7 +275,7 @@ def CGrid_Velocity(
 
     # check whether the grid conversion has been applied correctly
     xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3]
-    u = np.where(np.abs((xx - x) / x) > 1e-4, np.nan, u)
+    u = np.where(np.abs((xx - lon) / lon) > 1e-4, np.nan, u)
 
     if vectorfield.W:
         data = vectorfield.W.data
@@ -348,14 +327,8 @@ def CGrid_Velocity(
 
 
 def CGrid_Tracer(
-    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,
+    field: Field,
 ):
     """Interpolation kernel for tracer fields on a C-Grid.
 
@@ -365,6 +338,7 @@ def CGrid_Tracer(
     xi, _ = position["X"]
     yi, _ = position["Y"]
     zi, _ = position["Z"]
+    ti, tau = position["T"]
 
     axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
     data = field.data
@@ -403,31 +377,26 @@ def CGrid_Tracer(
 
 
 def _Spatialslip(
-    vectorfield: VectorField,
-    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,
+    vectorfield: VectorField,
     a: np.float32,
     b: np.float32,
 ):
     """Helper function for spatial boundary condition interpolation for velocity fields."""
     xi, xsi = position["X"]
     yi, eta = position["Y"]
     zi, zeta = position["Z"]
+    ti, tau = position["T"]
 
     axis_dim = vectorfield.U.grid.get_axis_dim_mapping(vectorfield.U.data.dims)
     lenT = 2 if np.any(tau > 0) else 1
     lenZ = 2 if np.any(zeta > 0) else 1
     npart = len(xsi)
 
-    u = XLinear(vectorfield.U, ti, position, tau, t, z, y, x)
-    v = XLinear(vectorfield.V, ti, position, tau, t, z, y, x)
+    u = XLinear(position, vectorfield.U)
+    v = XLinear(position, vectorfield.V)
     if vectorfield.W:
-        w = XLinear(vectorfield.W, ti, position, tau, t, z, y, x)
+        w = XLinear(position, vectorfield.W)
 
     corner_dataU = _get_corner_data_Agrid(vectorfield.U.data, ti, zi, yi, xi, lenT, lenZ, npart, axis_dim)
     corner_dataV = _get_corner_data_Agrid(vectorfield.V.data, ti, zi, yi, xi, lenT, lenZ, npart, axis_dim)
@@ -519,42 +488,24 @@ def is_land(ti: int, zi: int, yi: int, xi: int):
 
 
 def XFreeslip(
-    vectorfield: VectorField,
-    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,
+    vectorfield: VectorField,
 ):
     """Free-slip boundary condition interpolation for velocity fields."""
-    return _Spatialslip(vectorfield, ti, position, tau, t, z, y, x, a=1.0, b=0.0)
+    return _Spatialslip(position, vectorfield, a=1.0, b=0.0)
 
 
 def XPartialslip(
-    vectorfield: VectorField,
-    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,
+    vectorfield: VectorField,
 ):
     """Partial-slip boundary condition interpolation for velocity fields."""
-    return _Spatialslip(vectorfield, ti, position, tau, t, z, y, x, a=0.5, b=0.5)
+    return _Spatialslip(position, vectorfield, a=0.5, b=0.5)
 
 
 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,
+    field: Field,
 ):
     """
     Nearest-Neighbour spatial interpolation on a regular grid.
@@ -563,6 +514,7 @@ def XNearest(
     xi, xsi = position["X"]
     yi, eta = position["Y"]
     zi, zeta = position["Z"]
+    ti, tau = position["T"]
 
     axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
     data = field.data
@@ -610,49 +562,39 @@ def XNearest(
 
 
 def UXPiecewiseConstantFace(
-    field: Field,
-    ti: int,
     position: dict[_UXGRID_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,
+    field: Field,
 ):
     """
     Piecewise constant interpolation kernel for face registered data.
     This interpolation method is appropriate for fields that are
     face registered, such as u,v in FESOM.
     """
-    return field.data.values[ti, position["Z"][0], position["FACE"][0]]
+    return field.data.values[position["T"][0], position["Z"][0], position["FACE"][0]]
 
 
 def UXPiecewiseLinearNode(
-    field: Field,
-    ti: int,
     position: dict[_UXGRID_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,
+    field: Field,
 ):
     """
     Piecewise linear interpolation kernel for node registered data located at vertical interface levels.
     This interpolation method is appropriate for fields that are node registered such as the vertical
     velocity W in FESOM2. Effectively, it applies barycentric interpolation in the lateral direction
     and piecewise linear interpolation in the vertical direction.
     """
-    k, fi = position["Z"][0], position["FACE"][0]
+    ti, _ = position["T"]
+    zi, fi = position["Z"][0], position["FACE"][0]
+    z = position["z"]
     bcoords = position["FACE"][1]
     node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values
     # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels.
     # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1.
     # First, do barycentric interpolation in the lateral direction for each interface level
-    fzk = np.sum(field.data.values[ti[:, None], k[:, None], node_ids] * bcoords, axis=-1)
-    fzkp1 = np.sum(field.data.values[ti[:, None], k[:, None] + 1, node_ids] * bcoords, axis=-1)
+    fzk = np.sum(field.data.values[ti[:, None], zi[:, None], node_ids] * bcoords, axis=-1)
+    fzkp1 = np.sum(field.data.values[ti[:, None], zi[:, None] + 1, node_ids] * bcoords, axis=-1)
 
     # Then, do piecewise linear interpolation in the vertical direction
-    zk = field.grid.z.values[k]
-    zkp1 = field.grid.z.values[k + 1]
+    zk = field.grid.z.values[zi]
+    zkp1 = field.grid.z.values[zi + 1]
     return (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk)  # Linear interpolation in the vertical direction

tests/test_field.py

@@ -117,7 +117,7 @@ def test_field_invalid_interpolator():
     ds = datasets_structured["ds_2d_left"]
     grid = XGrid.from_dataset(ds)
 
-    def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid):
+    def invalid_interpolator_wrong_signature(position, invalid):
         return 0.0
 
     # Test invalid interpolator with wrong signature
@@ -129,7 +129,7 @@ def test_vectorfield_invalid_interpolator():
     ds = datasets_structured["ds_2d_left"]
     grid = XGrid.from_dataset(ds)
 
-    def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid):
+    def invalid_interpolator_wrong_signature(position, invalid):
         return 0.0
 
     # Create component fields

tests/test_interpolation.py

@@ -84,10 +84,10 @@ def field():
 )
 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)
+    position["T"] = _search_time_index(field, t)
 
-    value = func(field, ti, position, tau, 0, 0, y, x)
+    value = func(position, field)
     np.testing.assert_equal(value, expected)
 
 

tests/test_particleset_execute.py

@@ -298,8 +298,8 @@ def test_raise_general_error(): ...
 
 
 def test_errorinterpolation(fieldset):
-    def NaNInterpolator(field, ti, position, tau, t, z, y, x):  # pragma: no cover
-        return np.nan * np.zeros_like(x)
+    def NaNInterpolator(position, field):  # pragma: no cover
+        return np.nan * np.zeros_like(position["lon"])
 
     def SampleU(particles, fieldset):  # pragma: no cover
         fieldset.U[particles.time, particles.z, particles.lat, particles.lon, particles]