From 85f025be11cd122f8ac4bba463db2a97860c2549 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 19 Aug 2025 09:17:31 +0200 Subject: [PATCH 1/2] Adding unit test for spatial hash on 1/4 degree NEMO curvilinear grid The unit test currently fails, because the spatial hash index returned is the cell on the anti-meridian (between 179.875 and 179.875) --- tests/v4/test_index_search.py | 36 ++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/tests/v4/test_index_search.py b/tests/v4/test_index_search.py index 5c1a47307..5606913fd 100644 --- a/tests/v4/test_index_search.py +++ b/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) From 3eb280342bb80666cca40da419eef2ee16922c0c Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 1 Sep 2025 12:25:28 +0200 Subject: [PATCH 2/2] Vectorizing _reconnect_bnd_indices --- parcels/_index_search.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/parcels/_index_search.py b/parcels/_index_search.py index 5fd2d0620..f560d8167 100644 --- a/parcels/_index_search.py +++ b/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