Fixing last failing tests in vectorized kernels

Author: erikvansebille

parcels/application_kernels/advectiondiffusion.py

@@ -28,18 +28,18 @@ def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
     dWx = np.random.normal(0, np.sqrt(np.fabs(dt)))
     dWy = np.random.normal(0, np.sqrt(np.fabs(dt)))
 
-    Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres]
-    Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres]
+    Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres, particle]
+    Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres, particle]
     dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres)
 
-    u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon]
-    bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon])
+    u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon, particle]
+    bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon, particle])
 
-    Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon]
-    Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon]
+    Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon, particle]
+    Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon, particle]
     dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres)
 
-    by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon])
+    by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon, particle])
 
     # Particle positions are updated only after evaluating all terms.
     particle.dlon += u * dt + 0.5 * dKdx * (dWx**2 + dt) + bx * dWx
@@ -64,19 +64,19 @@ def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
     dWx = np.random.normal(0, np.sqrt(np.fabs(dt)))
     dWy = np.random.normal(0, np.sqrt(np.fabs(dt)))
 
-    u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon]
+    u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon, particle]
 
-    Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres]
-    Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres]
+    Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres, particle]
+    Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres, particle]
     dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres)
     ax = u + dKdx
-    bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon])
+    bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon, particle])
 
-    Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon]
-    Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon]
+    Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon, particle]
+    Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon, particle]
     dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres)
     ay = v + dKdy
-    by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon])
+    by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon, particle])
 
     # Particle positions are updated only after evaluating all terms.
     particle.dlon += ax * dt + bx * dWx

parcels/application_kernels/interpolation.py

@@ -17,7 +17,6 @@
     "UXPiecewiseConstantFace",
     "UXPiecewiseLinearNode",
     "XBiLinear",
-    "XBiLinearPeriodic",
     "XTriLinear",
 ]
 
@@ -57,47 +56,6 @@ def XBiLinear(
     return val
 
 
-def XBiLinearPeriodic(
-    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,
-):
-    """Bilinear interpolation on a regular grid with periodic boundary conditions in horizontal directions."""
-    xi, xsi = position["X"]
-    yi, eta = position["Y"]
-    zi, _ = position["Z"]
-
-    xi = np.where(xi > len(field.grid.lon) - 2, 0, xi)
-    xsi = (x - field.grid.lon[xi]) / (field.grid.lon[xi + 1] - field.grid.lon[xi])
-    yi = np.where(yi > len(field.grid.lat) - 2, 0, yi)
-    eta = (y - field.grid.lat[yi]) / (field.grid.lat[yi + 1] - field.grid.lat[yi])
-
-    axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
-
-    data = field.data
-    val = np.zeros_like(tau)
-
-    timeslices = [ti, ti + 1] if tau.any() > 0 else [ti]
-    for tii, tau_factor in zip(timeslices, [1 - tau, tau], strict=False):
-        xi = xr.DataArray(xi, dims="points")
-        yi = xr.DataArray(yi, dims="points")
-        zi = xr.DataArray(zi, dims="points")
-        ti = xr.DataArray(tii, dims="points")
-        F00 = data.isel({axis_dim["X"]: xi, axis_dim["Y"]: yi, axis_dim["Z"]: zi, "time": ti}).values.flatten()
-        F10 = data.isel({axis_dim["X"]: xi + 1, axis_dim["Y"]: yi, axis_dim["Z"]: zi, "time": ti}).values.flatten()
-        F01 = data.isel({axis_dim["X"]: xi, axis_dim["Y"]: yi + 1, axis_dim["Z"]: zi, "time": ti}).values.flatten()
-        F11 = data.isel({axis_dim["X"]: xi + 1, axis_dim["Y"]: yi + 1, axis_dim["Z"]: zi, "time": ti}).values.flatten()
-        val += (
-            (1 - xsi) * (1 - eta) * F00 + xsi * (1 - eta) * F10 + (1 - xsi) * eta * F01 + xsi * eta * F11
-        ) * tau_factor
-    return val
-
-
 def XTriLinear(
     field: Field,
     ti: int,

parcels/field.py

@@ -23,12 +23,10 @@
 )
 from parcels.tools.statuscodes import (
     AllParcelsErrorCodes,
-    FieldOutOfBoundError,
-    FieldOutOfBoundSurfaceError,
-    FieldSamplingError,
+    StatusCode,
 )
 from parcels.uxgrid import UxGrid
-from parcels.xgrid import XGrid, _transpose_xfield_data_to_tzyx
+from parcels.xgrid import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, XGrid, _transpose_xfield_data_to_tzyx
 
 from ._index_search import _search_time_index
 
@@ -257,19 +255,11 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         # else:
         #    _ei = particle.ei[self.igrid]
 
-        try:
-            tau, ti = _search_time_index(self, time)
-            position = self.grid.search(z, y, x, ei=_ei)
-            value = self._interp_method(self, ti, position, tau, time, z, y, x)
-
-            # TODO fix outof bounds sampling
-            # if np.isnan(value):
-            #     # Detect Out-of-bounds sampling and raise exception
-            #     _raise_field_out_of_bound_error(z, y, x)
+        tau, ti = _search_time_index(self, time)
+        position = self.grid.search(z, y, x, ei=_ei)
+        _update_particle_states(particle, position)
 
-        except (FieldSamplingError, FieldOutOfBoundError, FieldOutOfBoundSurfaceError) as e:
-            e.add_note(f"Error interpolating field '{self.name}'.")
-            raise e
+        value = self._interp_method(self, ti, position, tau, time, z, y, x)
 
         if applyConversion:
             value = self.units.to_target(value, z, y, x)
@@ -345,31 +335,28 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         # else:
         #    _ei = particle.ei[self.igrid]
 
-        try:
-            tau, ti = _search_time_index(self.U, time)
-            position = self.grid.search(z, y, x, ei=_ei)
-            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)
-                if "3D" in self.vector_type:
-                    w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x)
-            else:
-                (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x)
+        tau, ti = _search_time_index(self.U, time)
+        position = self.grid.search(z, y, x, ei=_ei)
+        _update_particle_states(particle, position)
 
-            if applyConversion:
-                u = self.U.units.to_target(u, z, y, x)
-                v = self.V.units.to_target(v, z, y, x)
-                if "3D" in self.vector_type:
-                    w = self.W.units.to_target(w, z, y, x) if self.W else 0.0
+        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)
+            if "3D" in self.vector_type:
+                w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x)
+        else:
+            (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x)
 
+        if applyConversion:
+            u = self.U.units.to_target(u, z, y, x)
+            v = self.V.units.to_target(v, z, y, x)
             if "3D" in self.vector_type:
-                return (u, v, w)
-            else:
-                return (u, v)
+                w = self.W.units.to_target(w, z, y, x) if self.W else 0.0
 
-        except (FieldSamplingError, FieldOutOfBoundError, FieldOutOfBoundSurfaceError) as e:
-            e.add_note(f"Error interpolating field '{self.name}'.")
-            raise e
+        if "3D" in self.vector_type:
+            return (u, v, w)
+        else:
+            return (u, v)
 
     def __getitem__(self, key):
         try:
@@ -381,6 +368,17 @@ def __getitem__(self, key):
             return _deal_with_errors(error, key, vector_type=self.vector_type)
 
 
+def _update_particle_states(particle, position):
+    """Update the particle states based on the position dictionary."""
+    if particle and "X" in position:  # TODO also support uxgrid search
+        particle.state = np.where(position["X"][0] < 0, StatusCode.ErrorOutOfBounds, particle.state)
+        particle.state = np.where(position["Y"][0] < 0, StatusCode.ErrorOutOfBounds, particle.state)
+        particle.state = np.where(position["Z"][0] == RIGHT_OUT_OF_BOUNDS, StatusCode.ErrorOutOfBounds, particle.state)
+        particle.state = np.where(
+            position["Z"][0] == LEFT_OUT_OF_BOUNDS, StatusCode.ErrorThroughSurface, particle.state
+        )
+
+
 def _assert_valid_uxdataarray(data: ux.UxDataArray):
     """Verifies that all the required attributes are present in the xarray.DataArray or
     uxarray.UxDataArray object.

parcels/xgrid.py

@@ -21,6 +21,9 @@
 
 _DEFAULT_XGCM_KWARGS = {"periodic": False}
 
+LEFT_OUT_OF_BOUNDS = -2
+RIGHT_OUT_OF_BOUNDS = -1
+
 
 def get_cell_count_along_dim(axis: xgcm.Axis) -> int:
     first_coord = list(axis.coords.items())[0]
@@ -271,15 +274,6 @@ def search(self, z, y, x, ei=None):
         ds = self.xgcm_grid._ds
 
         zi, zeta = _search_1d_array(ds.depth.values, z)
-        # if any(zi == -1):  # TODO throw error only for those particles where zi == -1
-        #     if any(zeta < 0):
-        #         raise FieldOutOfBoundError(
-        #             f"Depth {z} is out of bounds for the grid with depth values {ds.depth.values}."
-        #         )
-        #     elif any(zeta > 1):
-        #         raise FieldOutOfBoundSurfaceError(
-        #             f"Depth {z} is out of the surface for the grid with depth values {ds.depth.values}."
-        #         )
 
         if ds.lon.ndim == 1:
             yi, eta = _search_1d_array(ds.lat.values, y)
@@ -477,7 +471,6 @@ def _search_1d_array(
     Searches for the particle location in a 1D array and returns barycentric coordinate along dimension.
 
     Assumptions:
-    - particle position x is within the bounds of the array
     - array is strictly monotonically increasing.
 
     Parameters
@@ -489,9 +482,9 @@ def _search_1d_array(
 
     Returns
     -------
-    int
-        Index of the element just before the position x in the array.
-    float
+    array of int
+        Index of the element just before the position x in the array. Note that this index is -2 if the index is left out of bounds and -1 if the index is right out of bounds.
+    array of float
         Barycentric coordinate.
     """
     # TODO v4: We probably rework this to deal with 0D arrays before this point (as we already know field dimensionality)
@@ -500,11 +493,11 @@ def _search_1d_array(
     index = np.searchsorted(arr, x, side="right") - 1
     index_next = np.clip(index + 1, 1, len(arr) - 1)  # Ensure we don't go out of bounds
 
-    bcoord = np.where(
-        index <= len(arr) - 2,
-        (x - arr[index]) / (arr[index_next] - arr[index]),
-        np.nan,  # If at the end of the array, we return np.nan
-    )
+    bcoord = (x - arr[index]) / (arr[index_next] - arr[index])
+
+    index = np.where(x < arr[0], LEFT_OUT_OF_BOUNDS, index)
+    index = np.where(x >= arr[-1], RIGHT_OUT_OF_BOUNDS, index)
+
     return np.atleast_1d(index), np.atleast_1d(bcoord)
 
 

tests/v4/test_advection.py

@@ -5,7 +5,7 @@
 from parcels._datasets.structured.generated import simple_UV_dataset
 from parcels.application_kernels.advection import AdvectionEE, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45
 from parcels.application_kernels.advectiondiffusion import AdvectionDiffusionEM, AdvectionDiffusionM1
-from parcels.application_kernels.interpolation import XBiLinear, XBiLinearPeriodic, XTriLinear
+from parcels.application_kernels.interpolation import XBiLinear, XTriLinear
 from parcels.field import Field, VectorField
 from parcels.fieldset import FieldSet
 from parcels.particle import Particle, Variable
@@ -46,8 +46,7 @@ def test_advection_zonal(mesh_type, npart=10):
 
 def periodicBC(particle, fieldset, time):
     particle.total_dlon += particle.dlon
-    particle.lon = np.fmod(particle.lon, fieldset.U.grid.lon[-1])
-    particle.lat = np.fmod(particle.lat, fieldset.U.grid.lat[-1])
+    particle.lon = np.fmod(particle.lon, 2)
 
 
 def test_advection_zonal_periodic():
@@ -56,9 +55,15 @@ def test_advection_zonal_periodic():
     ds["lon"].data = np.array([0, 2])
     ds["lat"].data = np.array([0, 2])
 
+    # add a halo
+    halo = ds.isel(XG=0)
+    halo.lon.values = ds.lon.values[1] + 1
+    halo.XG.values = ds.XG.values[1] + 2
+    ds = xr.concat([ds, halo], dim="XG")
+
     grid = XGrid.from_dataset(ds)
-    U = Field("U", ds["U"], grid, interp_method=XBiLinearPeriodic)
-    V = Field("V", ds["V"], grid, interp_method=XBiLinearPeriodic)
+    U = Field("U", ds["U"], grid, interp_method=XBiLinear)
+    V = Field("V", ds["V"], grid, interp_method=XBiLinear)
     UV = VectorField("UV", U, V)
     fieldset = FieldSet([U, V, UV])
 
@@ -95,7 +100,7 @@ def test_advection_3D_outofbounds(direction, wErrorThroughSurface):
     ds = simple_UV_dataset(mesh_type="flat")
     grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U"], grid, interp_method=XTriLinear)
-    U.data[:] = 0.01  # Set U to 0 at the surface
+    U.data[:] = 0.01  # Set U to small value (to avoid horizontal out of bounds)
     V = Field("V", ds["V"], grid, interp_method=XTriLinear)
     W = Field("W", ds["V"], grid, interp_method=XTriLinear)  # Use V as W for testing
     W.data[:] = -1.0 if direction == "up" else 1.0
@@ -104,18 +109,22 @@ def test_advection_3D_outofbounds(direction, wErrorThroughSurface):
     fieldset = FieldSet([U, V, W, UVW, UV])
 
     def DeleteParticle(particle, fieldset, time):  # pragma: no cover
-        if particle.state == StatusCode.ErrorOutOfBounds or particle.state == StatusCode.ErrorThroughSurface:
-            particle.state = StatusCode.Delete
+        particle.state = np.where(particle.state == StatusCode.ErrorOutOfBounds, StatusCode.Delete, particle.state)
+        particle.state = np.where(particle.state == StatusCode.ErrorThroughSurface, StatusCode.Delete, particle.state)
 
     def SubmergeParticle(particle, fieldset, time):  # pragma: no cover
-        if particle.state == StatusCode.ErrorThroughSurface:
-            dt = particle.dt / np.timedelta64(1, "s")
-            (u, v) = fieldset.UV[particle]
-            particle.dlon = u * dt
-            particle.dlat = v * dt
-            particle.ddepth = 0.0
-            particle.depth = 0
-            particle.state = StatusCode.Evaluate
+        if len(particle.state) == 0:
+            return
+        inds = np.argwhere(particle.state == StatusCode.ErrorThroughSurface).flatten()
+        if len(inds) == 0:
+            return
+        dt = particle.dt / np.timedelta64(1, "s")
+        (u, v) = fieldset.UV[particle[inds]]
+        particle.dlon[inds] = u * dt
+        particle.dlat[inds] = v * dt
+        particle.ddepth[inds] = 0.0
+        particle.depth[inds] = 0
+        particle.state[inds] = StatusCode.Evaluate
 
     kernels = [AdvectionRK4_3D]
     if wErrorThroughSurface:

tests/v4/test_particleset_execute.py

@@ -103,16 +103,14 @@ def test_execution_fail_python_exception(fieldset, npart=10):
     pset = ParticleSet(fieldset, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart))
 
     def PythonFail(particle, fieldset, time):  # pragma: no cover
-        if particle.time >= fieldset.time_interval.left + np.timedelta64(10, "s"):
+        inds = np.argwhere(particle.time >= fieldset.time_interval.left + np.timedelta64(10, "s"))
+        if inds.size > 0:
             raise RuntimeError("Enough is enough!")
-        else:
-            pass
 
     with pytest.raises(RuntimeError):
         pset.execute(PythonFail, runtime=np.timedelta64(20, "s"), dt=np.timedelta64(2, "s"))
     assert len(pset) == npart
-    assert pset.time[0] == fieldset.time_interval.left + np.timedelta64(10, "s")
-    assert all([time == fieldset.time_interval.left + np.timedelta64(0, "s") for time in pset.time[1:]])
+    assert all(pset.time == fieldset.time_interval.left + np.timedelta64(10, "s"))
 
 
 def test_uxstommelgyre_pset_execute():