Add get_axis_dim and axes to BaseGrid, XGrid, and UxGrid

Also added some simple tests in test_uxgrid.py

Author: VeckoTheGecko

parcels/basegrid.py

@@ -135,3 +135,38 @@ def unravel_index(self, ei: int) -> dict[str, int]:
             Raised if the method is not implemented for the current grid type.
         """
         ...
+
+    @property
+    @abstractmethod
+    def axes(self) -> list[str]:
+        """
+        Return a list of axis names that are part of this grid.
+
+        Returns
+        -------
+        list[str]
+            List of axis names, e.g. ["Z", "Y", "X"] for a 3D structured grid or ["Z", "FACE"] for an unstructured grid.
+        """
+        ...
+
+    @abstractmethod
+    def get_axis_dim(self, axis: str) -> int:
+        """
+        Return the dimensionality (number of cells/edges) along a specific axis.
+
+        Parameters
+        ----------
+        axis : str
+            The name of the axis to get the dimensionality for. Must be one of the values returned by self.axes.
+
+        Returns
+        -------
+        int
+            The number of cells/edges along the specified axis.
+
+        Raises
+        ------
+        ValueError
+            If the specified axis is not part of this grid.
+        """
+        ...

parcels/uxgrid.py

@@ -54,6 +54,19 @@ def depth(self):
             return np.zeros(1)
         return self.z.values
 
+    @property
+    def axes(self) -> list[_UXGRID_AXES]:
+        return ["Z", "FACE"]
+
+    def get_axis_dim(self, axis: _UXGRID_AXES) -> int:
+        if axis not in self.axes:
+            raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}")
+
+        if axis == "Z":
+            return len(self.z.values)
+        elif axis == "FACE":
+            return self.uxgrid.n_face
+
     def search(self, z, y, x, ei=None):
         tol = 1e-10
 

parcels/xgrid.py

@@ -1,4 +1,5 @@
 from collections.abc import Hashable, Mapping
+from functools import cached_property
 from typing import Literal, cast
 
 import numpy as np
@@ -112,21 +113,23 @@ def _datetimes(self):
     def time(self):
         return self._datetimes.astype(np.float64) / 1e9
 
-    @property
-    def xdim(self):
-        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("X"))
+    @cached_property
+    def xdim(self) -> int:
+        return self.get_axis_dim("X")
 
-    @property
-    def ydim(self):
-        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("Y"))
+    @cached_property
+    def ydim(self) -> int:
+        return self.get_axis_dim("Y")
 
-    @property
-    def zdim(self):
-        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("Z"))
+    @cached_property
+    def zdim(self) -> int:
+        return self.get_axis_dim("Z")
 
-    @property
-    def tdim(self):
-        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("T"))
+    def get_axis_dim(self, axis: _XGRID_AXES) -> int:
+        if axis not in self.axes:
+            raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}")
+
+        return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get(axis))
 
     @property
     def _z4d(self) -> Literal[0, 1]:
@@ -188,7 +191,9 @@ def ravel_index(self, axis_indices: dict[_XGRID_AXES, int]) -> int:
         xi = axis_indices.get("X", 0)
         yi = axis_indices.get("Y", 0)
         zi = axis_indices.get("Z", 0)
-        return xi + self.xdim * yi + self.xdim * self.ydim * zi
+        xdim = self.get_axis_dim("X")
+        ydim = self.get_axis_dim("Y")
+        return xi + xdim * yi + xdim * ydim * zi
 
     def unravel_index(self, ei) -> dict[_XGRID_AXES, int]:
         zi = ei // (self.xdim * self.ydim)

tests/v4/test_uxgrid.py

@@ -0,0 +1,23 @@
+import pytest
+
+from parcels._datasets.unstructured.generic import datasets as uxdatasets
+from parcels.uxgrid import UxGrid
+
+
+@pytest.mark.parametrize("uxds", [pytest.param(uxds, id=key) for key, uxds in uxdatasets.items()])
+def test_uxgrid_init_on_generic_datasets(uxds):
+    UxGrid(uxds.uxgrid, z=uxds.coords["nz"])
+
+
+@pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]])
+def test_uxgrid_axes(uxds):
+    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"])
+    assert grid.axes == ["Z", "FACE"]
+
+
+@pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]])
+def test_xgrid_get_axis_dim(uxds):
+    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"])
+
+    assert grid.get_axis_dim("FACE") == 721
+    assert grid.get_axis_dim("Z") == 2

tests/v4/test_xgrid.py

@@ -6,7 +6,7 @@
 from numpy.testing import assert_allclose
 
 from parcels import xgcm
-from parcels._datasets.structured.generic import T, X, Y, Z, datasets
+from parcels._datasets.structured.generic import X, Y, Z, datasets
 from parcels.xgrid import XGrid, _search_1d_array
 
 GridTestCase = namedtuple("GridTestCase", ["Grid", "attr", "expected"])
@@ -16,10 +16,6 @@
     GridTestCase(datasets["ds_2d_left"], "lat", datasets["ds_2d_left"].YG.values),
     GridTestCase(datasets["ds_2d_left"], "depth", datasets["ds_2d_left"].ZG.values),
     GridTestCase(datasets["ds_2d_left"], "time", datasets["ds_2d_left"].time.values.astype(np.float64) / 1e9),
-    GridTestCase(datasets["ds_2d_left"], "xdim", X),
-    GridTestCase(datasets["ds_2d_left"], "ydim", Y),
-    GridTestCase(datasets["ds_2d_left"], "zdim", Z),
-    GridTestCase(datasets["ds_2d_left"], "tdim", T),
 ]
 
 
@@ -45,9 +41,18 @@ def test_xgrid_init_on_generic_datasets(ds):
     XGrid(xgcm.Grid(ds, periodic=False))
 
 
-def test_xgrid_axes():
-    # Tests that the xgrid.axes property correctly identifies the axes and ordering
-    ...
+@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
+def test_xgrid_axes(ds):
+    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    assert grid.axes == ["Z", "Y", "X"]
+
+
+@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
+def test_xgrid_get_axis_dim(ds):
+    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    assert grid.get_axis_dim("Z") == Z
+    assert grid.get_axis_dim("Y") == Y
+    assert grid.get_axis_dim("X") == X
 
 
 def test_invalid_xgrid_field_array():