Add axes and get_axis_dim to BaseGrid and other cleanup

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 @@
         """
         ...
 
-    @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 @@
         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,4 +136,106 @@
         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]
+            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.
+        """
         ...
+
+
+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/field.py

@@ -10,7 +10,6 @@
 import xarray as xr
 
 from parcels._core.utils.time import TimeInterval
-from parcels._core.utils.unstructured import get_vertical_location_from_dims
 from parcels._reprs import default_repr
 from parcels._typing import (
     Mesh,
@@ -158,7 +157,7 @@ def __init__(
         self.grid = grid
 
         try:
-            self.time_interval = get_time_interval(data)
+            self.time_interval = _get_time_interval(data)
         except ValueError as e:
             e.add_note(
                 f"Error getting time interval for field {name!r}. Are you sure that the time dimension on the xarray dataset is stored as timedelta, datetime or cftime datetime objects?"
@@ -206,41 +205,6 @@ def units(self, value):
             raise ValueError(f"Units must be a UnitConverter object, got {type(value)}")
         self._units = value
 
-    @property
-    def lat(self):
-        if type(self.data) is ux.UxDataArray:
-            if self.data.attrs["location"] == "node":
-                return self.grid.node_lat
-            elif self.data.attrs["location"] == "face":
-                return self.grid.face_lat
-            elif self.data.attrs["location"] == "edge":
-                return self.grid.edge_lat
-        else:
-            return self.grid.lat
-
-    @property
-    def lon(self):
-        if type(self.data) is ux.UxDataArray:
-            if self.data.attrs["location"] == "node":
-                return self.grid.node_lon
-            elif self.data.attrs["location"] == "face":
-                return self.grid.face_lon
-            elif self.data.attrs["location"] == "edge":
-                return self.grid.edge_lon
-        else:
-            return self.grid.lon
-
-    @property
-    def depth(self):
-        if type(self.data) is ux.UxDataArray:
-            vertical_location = get_vertical_location_from_dims(self.data.dims)
-            if vertical_location == "center":
-                return self.grid.nz1
-            elif vertical_location == "face":
-                return self.grid.nz
-        else:
-            return self.grid.depth
-
     @property
     def xdim(self):
         if type(self.data) is xr.DataArray:
@@ -316,13 +280,6 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         else:
             return value
 
-    def _rescale_and_set_minmax(self, data):
-        data[np.isnan(data)] = 0
-        return data
-
-    def __getattr__(self, key: str):
-        return getattr(self.data, key)
-
     def __getitem__(self, key):
         self._check_velocitysampling()
         try:
@@ -333,9 +290,6 @@ def __getitem__(self, key):
         except tuple(AllParcelsErrorCodes.keys()) as error:
             return _deal_with_errors(error, key, vector_type=None)
 
-    def __contains__(self, key: str):
-        return key in self.data
-
 
 class VectorField:
     """VectorField class that holds vector field data needed to execute particles."""
@@ -350,9 +304,9 @@ def __init__(
         self.grid = U.grid
 
         if W is None:
-            assert_same_time_interval((U, V))
+            _assert_same_time_interval((U, V))
         else:
-            assert_same_time_interval((U, V, W))
+            _assert_same_time_interval((U, V, W))
 
         self.time_interval = U.time_interval
 
@@ -384,14 +338,6 @@ def vector_interp_method(self, method: Callable):
         _assert_same_function_signature(method, ref=ZeroInterpolator)
         self._vector_interp_method = method
 
-    # @staticmethod
-    # TODO : def _check_grid_dimensions(grid1, grid2):
-    #     return (
-    #         np.allclose(grid1.lon, grid2.lon)
-    #         and np.allclose(grid1.lat, grid2.lat)
-    #         and np.allclose(grid1.depth, grid2.depth)
-    #         and np.allclose(grid1.time, grid2.time)
-    #     )
     def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         """Interpolate field values in space and time.
 
@@ -498,14 +444,14 @@ def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux
             )
 
 
-def get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None:
+def _get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None:
     if len(data.time) == 1:
         return None
 
     return TimeInterval(data.time.values[0], data.time.values[-1])
 
 
-def assert_same_time_interval(fields: list[Field]) -> None:
+def _assert_same_time_interval(fields: list[Field]) -> None:
     if len(fields) == 0:
         return