Review feedback

Author: VeckoTheGecko

parcels/_index_search.py

@@ -352,7 +352,7 @@ def _search_indices_curvilinear_2d(
     if not ((0 <= xsi <= 1) and (0 <= eta <= 1)):
         _raise_field_sampling_error(y, x)
 
-    return (eta, xsi, yi, xi)
+    return (yi, eta, xi, xsi)
 
 
 ## TODO :  Still need to implement the search_indices_curvilinear

parcels/basegrid.py

@@ -38,8 +38,8 @@ def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int,
             A dictionary mapping spatial axis names to tuples of (index, barycentric_coordinates).
             The returned axes depend on the grid dimensionality and type:
 
-            - 3D structured grid: {"X": (xi, xsi), "Y": (yi, eta), "Z": (zi, zeta)}
-            - 2D structured grid: {"X": (xi, xsi), "Y": (yi, eta)}
+            - 3D structured grid: {"Z": (zi, zeta), "Y": (yi, eta), "X": (xi, xsi)}
+            - 2D structured grid: {"Y": (yi, eta), "X": (xi, xsi)}
             - 1D structured grid (depth): {"Z": (zi, zeta)}
             - Unstructured grid: {"Z": (zi, zeta), "FACE": (fi, bcoords)}
 
@@ -74,8 +74,8 @@ def ravel_index(self, axis_indices: dict[str, int]) -> int:
             A dictionary mapping axis names to their corresponding indices.
             The expected keys depend on the grid dimensionality and type:
 
-            - 3D structured grid: {"X": xi, "Y": yi, "Z": zi}
-            - 2D structured grid: {"X": xi, "Y": yi}
+            - 3D structured grid: {"Z": zi, "Y": yi, "X": xi}
+            - 2D structured grid: {"Y": yi, "X": xi}
             - 1D structured grid: {"Z": zi}
             - Unstructured grid: {"Z": zi, "FACE": fi}
 
@@ -114,8 +114,8 @@ def unravel_index(self, ei: int) -> dict[str, int]:
             A dictionary mapping axis names to their corresponding indices.
             The returned keys depend on the grid dimensionality and type:
 
-            - 3D structured grid: {"X": xi, "Y": yi, "Z": zi}
-            - 2D structured grid: {"X": xi, "Y": yi}
+            - 3D structured grid: {"Z": zi, "Y": yi, "X": xi}
+            - 2D structured grid: {"Y": yi, "X": xi}
             - 1D structured grid: {"Z": zi}
             - Unstructured grid: {"Z": zi, "FACE": fi}
 

parcels/xgrid.py

@@ -16,7 +16,7 @@
 _XGCM_AXES = Mapping[_XGCM_AXIS_DIRECTION, xgcm.Axis]
 
 
-def get_n_cell_edges_along_dim(axis: xgcm.Axis | None) -> int:
+def get_cell_edge_count_along_dim(axis: xgcm.Axis | None) -> int:
     if axis is None:
         return 1
     first_coord = list(axis.coords.items())[0]
@@ -36,7 +36,7 @@ class XGrid(BaseGrid):
     This class provides methods and properties required for indexing and interpolating on the grid.
 
     Assumptions:
-    - If using Parcels in the context of a periodic simulation, the provided grid already has a halo
+    - If using Parcels in the context of a spatially periodic simulation, the provided grid already has a halo
 
     """
 
@@ -112,19 +112,19 @@ def time(self):
 
     @property
     def xdim(self):
-        return get_n_cell_edges_along_dim(self.xgcm_grid.axes.get("X"))
+        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("X"))
 
     @property
     def ydim(self):
-        return get_n_cell_edges_along_dim(self.xgcm_grid.axes.get("Y"))
+        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("Y"))
 
     @property
     def zdim(self):
-        return get_n_cell_edges_along_dim(self.xgcm_grid.axes.get("Z"))
+        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("Z"))
 
     @property
     def tdim(self):
-        return get_n_cell_edges_along_dim(self.xgcm_grid.axes.get("T"))
+        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("T"))
 
     @property
     def time_origin(self):
@@ -177,7 +177,7 @@ def search(self, z, y, x, ei=None):
         if ds.lon.ndim == 1:
             yi, eta = _search_1d_array(ds.lat.values, y)
             xi, xsi = _search_1d_array(ds.lon.values, x)
-            return {"X": (xi, xsi), "Y": (yi, eta), "Z": (zi, zeta)}
+            return {"Z": (zi, zeta), "Y": (yi, eta), "X": (xi, xsi)}
 
         yi, xi = None, None
         if ei is not None:
@@ -186,9 +186,9 @@ def search(self, z, y, x, ei=None):
             yi = axis_indices.get("Y")
 
         if ds.lon.ndim == 2:
-            eta, xsi, yi, xi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
+            yi, eta, xi, xsi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
 
-            return {"X": (xi, xsi), "Y": (yi, eta), "Z": (zi, zeta)}
+            return {"Z": (zi, zeta), "Y": (yi, eta), "X": (xi, xsi)}
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
@@ -204,11 +204,7 @@ def unravel_index(self, ei) -> dict[_AXIS_DIRECTION, int]:
 
         yi = ei // self.xdim
         xi = ei % self.xdim
-        return {
-            "X": xi,
-            "Y": yi,
-            "Z": zi,
-        }
+        return {"Z": zi, "Y": yi, "X": xi}
 
 
 def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None:
@@ -219,7 +215,7 @@ def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION |
     return None
 
 
-def get_position_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_POSITION | None:
+def get_xgcm_position_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_POSITION | None:
     """For a given dimension, returns the position of the variable in the grid."""
     for axis in axes.values():
         var_to_position = {var: position for position, var in axis.coords.items()}
@@ -281,16 +277,16 @@ def assert_valid_lat_lon(da_lat, da_lon, axes: _XGCM_AXES):
     assert_all_dimensions_correspond_with_axis(da_lon, axes)
     assert_all_dimensions_correspond_with_axis(da_lat, axes)
 
-    dim_to_position = {dim: get_position_from_dim_name(axes, dim) for dim in da_lon.dims}
-    dim_to_position.update({dim: get_position_from_dim_name(axes, dim) for dim in da_lat.dims})
+    dim_to_position = {dim: get_xgcm_position_from_dim_name(axes, dim) for dim in da_lon.dims}
+    dim_to_position.update({dim: get_xgcm_position_from_dim_name(axes, dim) for dim in da_lat.dims})
 
     for dim in da_lon.dims:
-        if get_position_from_dim_name(axes, dim) == "center":
+        if get_xgcm_position_from_dim_name(axes, dim) == "center":
             raise ValueError(
                 f"Longitude DataArray {da_lon.name!r} with dims {da_lon.dims} is defined on the center of the grid, but must be defined on the F points."
             )
     for dim in da_lat.dims:
-        if get_position_from_dim_name(axes, dim) == "center":
+        if get_xgcm_position_from_dim_name(axes, dim) == "center":
             raise ValueError(
                 f"Latitude DataArray {da_lat.name!r} with dims {da_lat.dims} is defined on the center of the grid, but must be defined on the F points."
             )

tests/v4/test_index_search.py

@@ -30,7 +30,7 @@ def test_grid_indexing_fpoints(field_cone):
             x = grid.lon[yi_expected, xi_expected] + 0.00001
             y = grid.lat[yi_expected, xi_expected] + 0.00001
 
-            eta, xsi, yi, xi = _search_indices_curvilinear_2d(grid, y, x)
+            yi, eta, xi, xsi = _search_indices_curvilinear_2d(grid, y, x)
             if eta > 0.9:
                 yi_expected -= 1
             if xsi > 0.9:

tests/v4/test_xgrid.py

@@ -129,7 +129,7 @@ def test_xgrid_ravel_unravel_index():
     for xi in range(xdim):
         for yi in range(ydim):
             for zi in range(zdim):
-                axis_indices = {"X": xi, "Y": yi, "Z": zi}
+                axis_indices = {"Z": zi, "Y": yi, "X": xi}
                 ei = grid.ravel_index(axis_indices)
                 axis_indices_test = grid.unravel_index(ei)
                 assert axis_indices_test == axis_indices
@@ -155,7 +155,7 @@ def test_xgrid_search_cpoints(ds):
 
     for xi in range(grid.xdim - 1):
         for yi in range(grid.ydim - 1):
-            axis_indices = {"X": xi, "Y": yi, "Z": 0}
+            axis_indices = {"Z": 0, "Y": yi, "X": xi}
 
             lat, lon = lat_array[yi, xi], lon_array[yi, xi]
             axis_indices_bcoords = grid.search(0, lat, lon, ei=None)