review feedback

Author: VeckoTheGecko

parcels/_core/constants.py

@@ -1,3 +0,0 @@
-# Used in the searching of 1D arrays for time and space
-LEFT_OUT_OF_BOUNDS = -2
-RIGHT_OUT_OF_BOUNDS = -1

parcels/_core/field.py

@@ -8,12 +8,11 @@
 import uxarray as ux
 import xarray as xr
 
-from parcels._core.constants import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS
 from parcels._core.converters import (
     UnitConverter,
     _unitconverters_map,
 )
-from parcels._core.index_search import GRID_SEARCH_ERROR, _search_time_index
+from parcels._core.index_search import GRID_SEARCH_ERROR, LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, _search_time_index
 from parcels._core.particle import KernelParticle
 from parcels._core.statuscodes import (
     AllParcelsErrorCodes,

parcels/_core/index_search.py

@@ -13,6 +13,53 @@
 
 
 GRID_SEARCH_ERROR = -3
+LEFT_OUT_OF_BOUNDS = -2
+RIGHT_OUT_OF_BOUNDS = -1
+
+
+def _search_1d_array(
+    arr: np.array,
+    x: float,
+) -> tuple[int, int]:
+    """
+    Searches for particle locations in a 1D array and returns barycentric coordinate along dimension.
+
+    Assumptions:
+    - array is strictly monotonically increasing.
+
+    Parameters
+    ----------
+    arr : np.array
+        1D array to search in.
+    x : float
+        Position in the 1D array to search for.
+
+    Returns
+    -------
+    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)
+    if len(arr) < 2:
+        return np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x)
+    index = np.searchsorted(arr, x, side="right") - 1
+    # Use broadcasting to avoid repeated array access
+    arr_index = arr[index]
+    arr_next = arr[np.clip(index + 1, 1, len(arr) - 1)]  # Ensure we don't go out of bounds
+    bcoord = (x - arr_index) / (arr_next - arr_index)
+
+    # TODO check how we can avoid searchsorted when grid spacing is uniform
+    # dx = arr[1] - arr[0]
+    # index = ((x - arr[0]) / dx).astype(int)
+    # index = np.clip(index, 0, len(arr) - 2)
+    # bcoord = (x - arr[index]) / dx
+
+    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)
 
 
 def _search_time_index(field: Field, time: datetime):

parcels/_core/utils/array.py

@@ -6,8 +6,7 @@
 import uxarray as ux
 
 from parcels._core.basegrid import BaseGrid
-from parcels._core.index_search import GRID_SEARCH_ERROR, uxgrid_point_in_cell
-from parcels._core.utils.array import _search_1d_array
+from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, uxgrid_point_in_cell
 from parcels._typing import assert_valid_mesh
 
 _UXGRID_AXES = Literal["Z", "FACE"]

parcels/_core/uxgrid.py

@@ -8,8 +8,7 @@
 import xgcm
 
 from parcels._core.basegrid import BaseGrid
-from parcels._core.index_search import _search_indices_curvilinear_2d
-from parcels._core.utils.array import _search_1d_array
+from parcels._core.index_search import _search_1d_array, _search_indices_curvilinear_2d
 from parcels._typing import assert_valid_mesh
 
 _XGRID_AXES = Literal["X", "Y", "Z"]

parcels/_core/xgrid.py

@@ -6,8 +6,7 @@
 import xarray as xr
 from numpy.testing import assert_allclose
 
-from parcels._core.constants import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS
-from parcels._core.utils.array import _search_1d_array
+from parcels._core.index_search import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, _search_1d_array
 from parcels._core.xgrid import (
     XGrid,
     _transpose_xfield_data_to_tzyx,