Generalise ravel/unravel to any grid

Author: VeckoTheGecko

parcels/basegrid.py

@@ -4,6 +4,8 @@
 from enum import IntEnum
 from typing import TYPE_CHECKING
 
+import numpy as np
+
 if TYPE_CHECKING:
     import numpy as np
 
@@ -67,7 +69,6 @@ def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int,
         """
         ...
 
-    @abstractmethod
     def ravel_index(self, axis_indices: dict[str, int]) -> int:
         """
         Convert a dictionary of axis indices to a single encoded index (ei).
@@ -101,9 +102,10 @@ def ravel_index(self, axis_indices: dict[str, int]) -> int:
         NotImplementedError
             Raised if the method is not implemented for the current grid type.
         """
-        ...
+        dims = np.array([self.get_axis_dim(axis) for axis in self.axes], dtype=int)
+        indices = np.array([axis_indices[axis] for axis in self.axes], dtype=int)
+        return _ravel(dims, indices)
 
-    @abstractmethod
     def unravel_index(self, ei: int) -> dict[str, int]:
         """
         Convert a single encoded index (ei) back to a dictionary of axis indices.
@@ -134,14 +136,19 @@ def unravel_index(self, ei: int) -> dict[str, int]:
         NotImplementedError
             Raised if the method is not implemented for the current grid type.
         """
-        ...
+        dims = np.array([self.get_axis_dim(axis) for axis in self.axes], dtype=int)
+        indices = _unravel(dims, ei)
+        return dict(zip(self.axes, indices, strict=True))
 
     @property
     @abstractmethod
     def axes(self) -> list[str]:
         """
         Return a list of axis names that are part of this grid.
 
+        This list must at least be of length 1, and `get_axis_dim` should
+        return a valid integer for each axis name in the list.
+
         Returns
         -------
         list[str]
@@ -170,3 +177,65 @@ def get_axis_dim(self, axis: str) -> int:
             If the specified axis is not part of this grid.
         """
         ...
+
+
+def _unravel(dims, ei):
+    """
+    Converts a flattened (raveled) index back to multi-dimensional indices.
+
+    Args:
+        dims (1d-array-like): The dimensions along each axis
+        ei (int): The flattened index to convert
+
+    Returns
+    -------
+        array-like: Indices along each axis corresponding to the given flattened index
+
+    Example:
+        >>> dims = [2, 3, 4]
+        >>> ei = 9
+        >>> unravel(dims, ei)
+        array([0, 2, 1])
+        # Calculation:
+        # i0 = 9 // (3*4) = 9 // 12 = 0
+        # remainder = 9 % 12 = 9
+        # i1 = 9 // 4 = 2
+        # i2 = 9 % 4 = 1
+    """
+    strides = np.cumprod(dims[::-1])[::-1]
+
+    indices = np.empty(len(dims), dtype=int)
+
+    for i in range(len(dims) - 1):
+        indices[i] = ei // strides[i + 1]
+        ei = ei % strides[i + 1]
+
+    indices[-1] = ei
+    return indices
+
+
+def _ravel(dims, indices):
+    """
+    Converts indices to a flattened (raveled) index.
+
+    Args:
+        dims (1d-array-like): The dimensions along each axis
+        indices (array-like): Indices along each axis to convert
+
+    Returns
+    -------
+        int: The flattened index corresponding to the given indices
+
+    Example:
+        >>> dims = [2, 3, 4]
+        >>> indices = [0, 2, 1]
+        >>> ravel(dims, indices)
+        9
+        # Calculation: 0 * (3 * 4) + 2 * (4) + 1 = 0 + 8 + 1 = 9
+    """
+    strides = np.cumprod(dims[::-1])[::-1]
+    ei = 0
+    for i in range(len(dims) - 1):
+        ei += indices[i] * strides[i + 1]
+
+    return ei + indices[-1]

parcels/uxgrid.py

@@ -114,11 +114,3 @@ def _get_barycentric_coordinates(self, y, x, fi):
         bcoord = np.asarray(_barycentric_coordinates(nodes, coord))
         err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs(np.dot(bcoord, nodes[:, 1]) - coord[1])
         return bcoord, err
-
-    def ravel_index(self, axis_indices: dict[_UXGRID_AXES, int]):
-        return axis_indices["FACE"] + self.uxgrid.n_face * axis_indices["Z"]
-
-    def unravel_index(self, ei) -> dict[_UXGRID_AXES, int]:
-        zi = ei // self.uxgrid.n_face
-        fi = ei % self.uxgrid.n_face
-        return {"Z": zi, "FACE": fi}

parcels/xgrid.py

@@ -187,22 +187,6 @@ def search(self, z, y, x, ei=None):
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
-    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)
-        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)
-        ei = ei % (self.xdim * self.ydim)
-
-        yi = ei // self.xdim
-        xi = ei % self.xdim
-        return {"Z": zi, "Y": yi, "X": xi}
-
 
 def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None:
     """For a given dimension name in a grid, returns the direction axis it is on."""

tests/v4/test_basegrid.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import itertools
+
+import numpy as np
+import pytest
+
+from parcels.basegrid import BaseGrid
+
+
+class TestGrid(BaseGrid):
+    def __init__(self, axis_dim: dict[str, int]):
+        self.axis_dim = axis_dim
+
+    def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int, float | np.ndarray]]:
+        pass
+
+    @property
+    def axes(self) -> list[str]:
+        return list(self.axis_dim.keys())
+
+    def get_axis_dim(self, axis: str) -> int:
+        return self.axis_dim[axis]
+
+
+@pytest.mark.parametrize(
+    "grid",
+    [
+        TestGrid({"Z": 10, "Y": 20, "X": 30}),
+        TestGrid({"Z": 5, "Y": 15}),
+        TestGrid({"Z": 8}),
+        TestGrid({"Z": 12, "FACE": 25}),
+    ],
+)
+def test_basegrid_ravel_unravel_index(grid):
+    axes = grid.axes
+    dimensionalities = (grid.get_axis_dim(axis) for axis in axes)
+    all_possible_axis_indices = itertools.product(*[range(dim) for dim in dimensionalities])
+
+    encountered_eis = []
+
+    for axis_indices_numeric in all_possible_axis_indices:
+        axis_indices = dict(zip(axes, axis_indices_numeric, strict=True))
+
+        ei = grid.ravel_index(axis_indices)
+        axis_indices_test = grid.unravel_index(ei)
+        assert axis_indices_test == axis_indices
+        encountered_eis.append(ei)
+
+    encountered_eis = sorted(encountered_eis)
+    assert len(set(encountered_eis)) == len(encountered_eis), "Raveled indices are not unique."
+    assert np.allclose(np.diff(np.array(encountered_eis)), 1), "Raveled indices are not consecutive integers."
+    assert encountered_eis[0] == 0, "Raveled indices do not start at 0."

tests/v4/test_xgrid.py

@@ -90,30 +90,6 @@ def test_invalid_lon_lat():
         XGrid(xgcm.Grid(ds, periodic=False))
 
 
-def test_xgrid_ravel_unravel_index():
-    ds = datasets["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
-
-    xdim = grid.xdim
-    ydim = grid.ydim
-    zdim = grid.zdim
-
-    encountered_eis = []
-    for xi in range(xdim):
-        for yi in range(ydim):
-            for zi in range(zdim):
-                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
-                encountered_eis.append(ei)
-
-    encountered_eis = sorted(encountered_eis)
-    assert len(set(encountered_eis)) == len(encountered_eis), "Raveled indices are not unique."
-    assert np.allclose(np.diff(np.array(encountered_eis)), 1), "Raveled indices are not consecutive integers."
-    assert encountered_eis[0] == 0, "Raveled indices do not start at 0."
-
-
 @pytest.mark.parametrize(
     "ds",
     [