Patch return of XGrid.search() to match BaseGrid

Author: VeckoTheGecko

parcels/xgrid.py

@@ -191,7 +191,7 @@ def search(self, z, y, x, ei=None, search2D=False):
         if ds.lon.ndim == 1:
             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])
+            return self.ravel_index(zi, yi, xi), np.array([eta, xsi, 1 - eta, 1 - xsi])
 
         yi, xi = None, None
         if ei is not None:
@@ -200,7 +200,7 @@ def search(self, z, y, x, ei=None, search2D=False):
         if ds.lon.ndim == 2:
             eta, xsi, yi, xi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
 
-            return (zi, yi, xi), np.array([eta, xsi, 1 - eta, 1 - xsi])
+            return self.ravel_index(zi, yi, xi), np.array([eta, xsi, 1 - eta, 1 - xsi])
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 

tests/v4/test_xgrid.py

@@ -157,7 +157,8 @@ def test_xgrid_search_cpoints(ds):
     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)
+            ei, bcoords = grid.search(0, lat, lon, ei=None, search2D=True)
+            zi_test, yi_test, xi_test = grid.unravel_index(ei)
             assert xi == xi_test
             assert yi == yi_test
             assert zi_test == 0