Testing and fixing NEMO spatialhash index_search

parcels/_index_search.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 
+from parcels._typing import Mesh
 from parcels.tools.statuscodes import (
     _raise_field_out_of_bound_error,
     _raise_field_sampling_error,
@@ -108,21 +109,13 @@ def _search_indices_curvilinear_2d(
     return (yi, eta, xi, xsi)
 
 
-def _reconnect_bnd_indices(yi: int, xi: int, ydim: int, xdim: int, sphere_mesh: bool):
-    if xi < 0:
-        if sphere_mesh:
-            xi = xdim - 2
-        else:
-            xi = 0
-    if xi > xdim - 2:
-        if sphere_mesh:
-            xi = 0
-        else:
-            xi = xdim - 2
-    if yi < 0:
-        yi = 0
-    if yi > ydim - 2:
-        yi = ydim - 2
-        if sphere_mesh:
-            xi = xdim - xi
+def _reconnect_bnd_indices(yi: int, xi: int, ydim: int, xdim: int, mesh: Mesh):
+    xi = np.where(xi < 0, (xdim - 2) if mesh == "spherical" else 0, xi)
+    xi = np.where(xi > xdim - 2, 0 if mesh == "spherical" else (xdim - 2), xi)
+
+    xi = np.where(yi > ydim - 2, xdim - xi if mesh == "spherical" else xi, xi)
+
+    yi = np.where(yi < 0, 0, yi)
+    yi = np.where(yi > ydim - 2, ydim - 2, yi)
+
     return yi, xi

tests/v4/test_index_search.py

@@ -1,12 +1,13 @@
 import numpy as np
 import pytest
+import xarray as xr
 
 from parcels._datasets.structured.generic import datasets
 from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.field import Field
-from parcels.xgrid import (
-    XGrid,
-)
+from parcels.tools.exampledata_utils import download_example_dataset
+from parcels.xgcm import Grid
+from parcels.xgrid import XGrid
 
 
 @pytest.fixture
@@ -51,3 +52,32 @@ def test_grid_indexing_fpoints(field_cone):
             ]
             assert x > np.min(cell_lon) and x < np.max(cell_lon)
             assert y > np.min(cell_lat) and y < np.max(cell_lat)
+
+
+def test_indexing_nemo_curvilinear():
+    data_folder = download_example_dataset("NemoCurvilinear_data")
+    ds = xr.open_mfdataset(
+        data_folder.glob("*.nc4"), combine="nested", data_vars="minimal", coords="minimal", compat="override"
+    )
+    ds = ds.isel({"time_counter": 0, "time": 0, "z_a": 0}, drop=True).rename(
+        {"glamf": "lon", "gphif": "lat", "z": "depth"}
+    )
+    xgcm_grid = Grid(ds, coords={"X": {"left": "x"}, "Y": {"left": "y"}}, periodic=False)
+    grid = XGrid(xgcm_grid)
+
+    # Test points on the NEMO 1/4 degree curvilinear grid
+    lats = np.array([-30, 0, 88])
+    lons = np.array([30, 60, -150])
+
+    yi, eta, xi, xsi = _search_indices_curvilinear_2d(grid, lats, lons)
+
+    # Construct cornerpoints px
+    px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]])
+
+    # Maximum 5 degree difference between px values
+    for i in range(lons.shape[0]):
+        np.testing.assert_allclose(px[1, i], px[:, i], atol=5)
+
+    # Reconstruct lons values from cornerpoints
+    xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3]
+    np.testing.assert_allclose(xx, lons, atol=1e-6)