From 92cd143858abe11f9a80dddcf91f0b5f955cf747 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 29 Oct 2026 09:44:58 +0100 Subject: [PATCH] Relax hardcoding of lon and lat in sgrid node_coordinates --- src/parcels/_core/model.py | 1 + src/parcels/_core/xgrid.py | 45 +++++++++++++++++++++++--------------- tests/test_xgrid.py | 26 ++++++++++++++++++++++ 3 files changed, 54 insertions(+), 18 deletions(-) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index b7ab76679..2554f6b81 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -351,6 +351,7 @@ def create_empty_constant_field_model(mesh: SphericalMesh | FlatMesh) -> Structu sgrid.FaceNodePadding("XC", "lon", sgrid.Padding.LOW), sgrid.FaceNodePadding("YC", "lat", sgrid.Padding.LOW), ), + node_coordinates=("lon", "lat"), ), ), mesh=mesh, diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py index 20b6cd788..c96fee26b 100644 --- a/src/parcels/_core/xgrid.py +++ b/src/parcels/_core/xgrid.py @@ -131,15 +131,18 @@ def __init__(self, model_data: xr.Dataset, mesh: ptyping.TMesh): self._spatialhash = None ds = model_data - # Set the coordinates for the dataset (needed to be done explicitly for curvilinear grids) - if "lon" in ds: - ds.set_coords("lon") - if "lat" in ds: - ds.set_coords("lat") - axes = self.axes - if len(set(axes) & {"X", "Y"}) > 0: # Only if spatial grid is >0D (see #2054 for further development) - assert_valid_lat_lon(ds["lat"], ds["lon"], self.sgrid_metadata) + if len(set(axes) & {"X", "Y"}) > 0: # TODO also support grids with only lon or lat (see #2054) + if self.sgrid_metadata.node_coordinates is None: + raise ValueError("SGRID metadata must define 'node_coordinates' for a grid with horizontal axes.") + x_coordinate_name, y_coordinate_name = self.sgrid_metadata.node_coordinates + missing_coordinates = [name for name in (x_coordinate_name, y_coordinate_name) if name not in ds] + if missing_coordinates: + raise ValueError( + f"SGRID metadata references node coordinates {missing_coordinates}, which are not present in the dataset." + ) + ds = ds.set_coords([x_coordinate_name, y_coordinate_name]) + assert_valid_lat_lon(ds[y_coordinate_name], ds[x_coordinate_name], self.sgrid_metadata) if "Z" in axes: depth = ds["depth"] @@ -184,10 +187,11 @@ def lon(self): """ if "X" not in self.axes: return np.zeros(1) + x_coordinate_name = self.sgrid_metadata.node_coordinates[0] # ensure lon is loaded into memory for dask-backed datasets, as it is used in the search method - if is_dask_collection(self._ds["lon"].data): - self._ds["lon"].load() - return self._ds["lon"].values + if is_dask_collection(self._ds[x_coordinate_name].data): + self._ds[x_coordinate_name].load() + return self._ds[x_coordinate_name].values @property def lat(self): @@ -199,10 +203,11 @@ def lat(self): """ if "Y" not in self.axes: return np.zeros(1) + y_coordinate_name = self.sgrid_metadata.node_coordinates[1] # ensure lat is loaded into memory for dask-backed datasets, as it is used in the search method - if is_dask_collection(self._ds["lat"].data): - self._ds["lat"].load() - return self._ds["lat"].values + if is_dask_collection(self._ds[y_coordinate_name].data): + self._ds[y_coordinate_name].load() + return self._ds[y_coordinate_name].values @property def depth(self): @@ -347,18 +352,22 @@ def _gtype(self): def search(self, z, y, x, ei=None, ti=None, tau=None): ds = self._ds + if self.sgrid_metadata.node_coordinates is not None: + x_coordinate_name, y_coordinate_name = self.sgrid_metadata.node_coordinates + x_coordinates = ds[x_coordinate_name] + y_coordinates = ds[y_coordinate_name] - if "X" not in self.axes or ds.lon.ndim == 1: + if "X" not in self.axes or x_coordinates is None or x_coordinates.ndim == 1: if "Y" in self.axes: - yi, eta = _search_1d_array(ds.lat.values, y) + yi, eta = _search_1d_array(y_coordinates.values, y) else: yi, eta = np.zeros(y.shape, dtype=int), np.zeros(y.shape, dtype=float) if "X" in self.axes: - xi, xsi = _search_1d_array(ds.lon.values, x) + xi, xsi = _search_1d_array(x_coordinates.values, x) else: xi, xsi = np.zeros(x.shape, dtype=int), np.zeros(x.shape, dtype=float) - elif ds.lon.ndim == 2: + elif x_coordinates is not None and x_coordinates.ndim == 2: yi, xi = None, None if ei is not None: axis_indices = self.unravel_index(ei) diff --git a/tests/test_xgrid.py b/tests/test_xgrid.py index 38151cd1c..bbc2bef21 100644 --- a/tests/test_xgrid.py +++ b/tests/test_xgrid.py @@ -71,6 +71,32 @@ def test_xgrid_axes(fieldset): assert fieldset.U.grid.axes == ["Z", "Y", "X"] +@pytest.mark.parametrize( + "ds", + [datasets["ds_2d_left"], datasets["2d_left_rotated"]], + ids=["rectilinear", "curvilinear"], +) +def test_xgrid_uses_sgrid_node_coordinate_names(ds): + reference_grid = FieldSet.from_sgrid_conventions(ds).data_g.grid + ds = ds.sgrid.rename({"lon": "dim1", "lat": "dim2"}) + grid = FieldSet.from_sgrid_conventions(ds).data_g.grid + + np.testing.assert_array_equal(grid.lon, ds["dim1"].values) + np.testing.assert_array_equal(grid.lat, ds["dim2"].values) + + if grid.lon.ndim == 1: + y = np.array([(grid.lat[1] + grid.lat[2]) / 2]) + x = np.array([(grid.lon[1] + grid.lon[2]) / 2]) + else: + y = np.array([grid.lat[1, 1]]) + x = np.array([grid.lon[1, 1]]) + result = grid.search(0, y, x) + reference_result = reference_grid.search(0, y, x) + for axis in ("X", "Y"): + np.testing.assert_array_equal(result[axis]["index"], reference_result[axis]["index"]) + np.testing.assert_allclose(result[axis]["bcoord"], reference_result[axis]["bcoord"]) + + @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]]) @pytest.mark.parametrize("mesh", [FlatMesh, SphericalMesh]) def test_xgrid_mesh_inferred_from_units(ds, mesh):