Update XGrid axis dim to use cell count

Author: VeckoTheGecko

parcels/basegrid.py

@@ -159,7 +159,7 @@ def axes(self) -> list[str]:
     @abstractmethod
     def get_axis_dim(self, axis: str) -> int:
         """
-        Return the dimensionality (number of cells/edges) along a specific axis.
+        Return the dimensionality (number of cells/faces) along a specific axis.
 
         Parameters
         ----------

parcels/xgrid.py

@@ -18,13 +18,11 @@
 _XGCM_AXES = Mapping[_XGCM_AXIS_DIRECTION, xgcm.Axis]
 
 
-def get_cell_edge_count_along_dim(axis: xgcm.Axis | None) -> int:
-    if axis is None:
-        return 1
+def get_cell_count_along_dim(axis: xgcm.Axis) -> int:
     first_coord = list(axis.coords.items())[0]
     _, coord_var = first_coord
 
-    return axis._ds[coord_var].size
+    return axis._ds[coord_var].size - 1
 
 
 def get_time(axis: xgcm.Axis) -> npt.NDArray:
@@ -129,7 +127,7 @@ 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))
+        return get_cell_count_along_dim(self.xgcm_grid.axes[axis])
 
     @property
     def _z4d(self) -> Literal[0, 1]:

tests/v4/test_xgrid.py

@@ -16,9 +16,9 @@
     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"], "xdim", X - 1),
+    GridTestCase(datasets["ds_2d_left"], "ydim", Y - 1),
+    GridTestCase(datasets["ds_2d_left"], "zdim", Z - 1),
 ]
 
 
@@ -53,9 +53,9 @@ def test_xgrid_axes(ds):
 @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
+    assert grid.get_axis_dim("Z") == Z - 1
+    assert grid.get_axis_dim("Y") == Y - 1
+    assert grid.get_axis_dim("X") == X - 1
 
 
 def test_invalid_xgrid_field_array():