Implement XGrid.search with tests for 2D lon lat

Author: VeckoTheGecko

parcels/_index_search.py

@@ -17,11 +17,12 @@
     _raise_field_sampling_error,
     _raise_time_extrapolation_error,
 )
-from parcels.xgrid import XGrid
 
 from .grid import GridType
 
 if TYPE_CHECKING:
+    from parcels.xgrid import XGrid
+
     from .field import Field
     # from .grid import Grid
 

parcels/xgrid.py

@@ -6,6 +6,7 @@
 import xarray as xr
 
 from parcels import xgcm
+from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.basegrid import BaseGrid
 from parcels.tools.converters import TimeConverter
 
@@ -182,18 +183,26 @@ def _gtype(self):
     def search(self, z, y, x, ei=None, search2D=False):
         ds = self.xgcm_grid._ds
 
+        if search2D:
+            zi = 0
+        else:
+            zi, _ = _search_1d_array(ds.depth.values, z)
+
         if ds.lon.ndim == 1:
-            yi, bcoord_y = _search_1d_array(ds.lat.values, y)
-            xi, bcoord_x = _search_1d_array(ds.lon.values, x)
+            yi, eta = _search_1d_array(ds.lat.values, y)
+            xi, xsi = _search_1d_array(ds.lon.values, x)
+            return (zi, yi, xi), np.array([eta, xsi, 1 - eta, 1 - xsi])
 
-            if search2D:
-                zi = 0
-            else:
-                zi, _ = _search_1d_array(ds.depth.values, z)
+        yi, xi = None, None
+        if ei is not None:
+            _, yi, xi = self.unravel_index(ei)
+
+        if ds.lon.ndim == 2:
+            eta, xsi, yi, xi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
 
-            return (zi, yi, xi), np.array([bcoord_y, bcoord_x, 1 - bcoord_y, 1 - bcoord_x])
+            return (zi, yi, xi), np.array([eta, xsi, 1 - eta, 1 - xsi])
 
-        raise NotImplementedError("Searching in 2D arrays is not implemented yet.")
+        raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
     def ravel_index(self, zi, yi, xi):
         """
@@ -415,6 +424,6 @@ def _search_1d_array(
     float
         Barycentric coordinate.
     """
-    i = np.argmin(arr < x)
+    i = np.argmin(arr <= x) - 1
     barry = (x - arr[i]) / (arr[i + 1] - arr[i])
     return i, barry

tests/v4/test_xgrid.py

@@ -10,10 +10,7 @@
 from parcels._datasets.structured.generic import T, X, Y, Z, datasets
 from parcels.grid import Grid as OldGrid
 from parcels.tools.converters import TimeConverter
-from parcels.xgrid import (
-    XGrid,
-    _generate_cells,
-)
+from parcels.xgrid import XGrid, _generate_cells, _search_1d_array
 
 GridTestCase = namedtuple("GridTestCase", ["Grid", "attr", "expected"])
 
@@ -187,3 +184,56 @@ def _assert_point_is(
             raise ValueError(f"Invalid method: {direction}")
 
     np.testing.assert_allclose(reference_cell + delta, test_cell)
+
+
+@pytest.mark.parametrize(
+    "ds",
+    [
+        pytest.param(datasets["ds_2d_left"], id="1D lon/lat"),
+        pytest.param(datasets["2d_left_rotated"], id="2D lon/lat"),
+    ],
+)  # for key, ds in datasets.items()])
+def test_xgrid_search_cpoints(ds):
+    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    lat_array, lon_array = get_2d_fpoint_mesh(grid)
+    lat_array, lon_array = corner_to_cell_center_points(lat_array, lon_array)
+
+    for xi in range(grid.xdim - 1):
+        for yi in range(grid.ydim - 1):
+            lat, lon = lat_array[yi, xi], lon_array[yi, xi]
+            (zi_test, yi_test, xi_test), bcoords = grid.search(0, lat, lon, ei=None, search2D=True)
+            assert xi == xi_test
+            assert yi == yi_test
+            assert zi_test == 0
+
+            # assert np.isclose(bcoords[0], 0.5) #? Should this not be the case with the cell center points?
+            # assert np.isclose(bcoords[1], 0.5)
+
+
+def get_2d_fpoint_mesh(grid: XGrid):
+    lat, lon = grid.lat, grid.lon
+    if lon.ndim == 1:
+        lat, lon = np.meshgrid(lat, lon, indexing="ij")
+    return lat, lon
+
+
+def corner_to_cell_center_points(lat, lon):
+    """Convert F points to C points."""
+    lon_c = (lon[:-1, :-1] + lon[:-1, 1:]) / 2
+    lat_c = (lat[:-1, :-1] + lat[1:, :-1]) / 2
+    return lat_c, lon_c
+
+
+@pytest.mark.parametrize(
+    "array, x, expected_xi, expected_xsi",
+    [
+        (np.array([1, 2, 3, 4, 5]), 1.1, 0, 0.1),
+        (np.array([1, 2, 3, 4, 5]), 2.1, 1, 0.1),
+        (np.array([1, 2, 3, 4, 5]), 3.1, 2, 0.1),
+        (np.array([1, 2, 3, 4, 5]), 4.5, 3, 0.5),
+    ],
+)
+def test_search_1d_array(array, x, expected_xi, expected_xsi):
+    xi, xsi = _search_1d_array(array, x)
+    assert xi == expected_xi
+    assert np.isclose(xsi, expected_xsi)