From d9322f007f394052d1d181fb85c93378cdea83b1 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 11:10:28 +0200
Subject: [PATCH 01/12] Remove FieldSet.dimrange

---
 parcels/fieldset.py | 28 ----------------------------
 1 file changed, 28 deletions(-)

diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index fbb3055668..6dea50cc1c 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -82,34 +82,6 @@ def time_interval(self):
             return None
         return functools.reduce(lambda x, y: x.intersection(y), time_intervals)
 
-    def dimrange(self, dim):
-        """Returns maximum value of a dimension (lon, lat, depth or time)
-        on 'left' side and minimum value on 'right' side for all grids
-        in a gridset. Useful for finding e.g. longitude range that
-        overlaps on all grids in a gridset.
-        """
-        maxleft, minright = (-np.inf, np.inf)
-        dim2ds = {
-            "depth": ["nz1", "nz"],
-            "lat": ["node_lat", "face_lat", "edge_lat"],
-            "lon": ["node_lon", "face_lon", "edge_lon"],
-            "time": ["time"],
-        }
-        for ds in self.datasets:
-            for field in ds.data_vars:
-                for d in dim2ds[dim]:  # check all possible dimensions
-                    if d in ds[field].dims:
-                        if dim == "depth":
-                            maxleft = max(maxleft, ds[field][d].min().data)
-                            minright = min(minright, ds[field][d].max().data)
-                        else:
-                            maxleft = max(maxleft, ds[field][d].data[0])
-                            minright = min(minright, ds[field][d].data[-1])
-        maxleft = 0 if maxleft == -np.inf else maxleft  # if all len(dim) == 1
-        minright = 0 if minright == np.inf else minright  # if all len(dim) == 1
-
-        return maxleft, minright
-
     def add_field(self, field: Field, name: str | None = None):
         """Add a :class:`parcels.field.Field` object to the FieldSet.
 

From 578d77ddc90b4b9c7d147b3cc681c2ae422d74f2 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 11:11:08 +0200
Subject: [PATCH 02/12] Remove commented code

---
 parcels/fieldset.py | 27 ---------------------------
 1 file changed, 27 deletions(-)

diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index 6dea50cc1c..c423000fc4 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -176,33 +176,6 @@ def gridset(self) -> list[BaseGrid]:
                 grids.append(field.grid)
         return grids
 
-    # def computeTimeChunk(self, time=0.0, dt=1):
-    #     """Load a chunk of three data time steps into the FieldSet.
-    #     This is used when FieldSet uses data imported from netcdf,
-    #     with default option deferred_load. The loaded time steps are at or immediatly before time
-    #     and the two time steps immediately following time if dt is positive (and inversely for negative dt)
-
-    #     Parameters
-    #     ----------
-    #     time :
-    #         Time around which the FieldSet data are to be loaded.
-    #         Time is provided as a double, relatively to Fieldset.time_origin.
-    #         Default is 0.
-    #     dt :
-    #         time step of the integration scheme, needed to set the direction of time chunk loading.
-    #         Default is 1.
-    #     """
-    #     nextTime = np.inf if dt > 0 else -np.inf
-
-    #     if abs(nextTime) == np.inf or np.isnan(nextTime):  # Second happens when dt=0
-    #         return nextTime
-    #     else:
-    #         nSteps = int((nextTime - time) / dt)
-    #         if nSteps == 0:
-    #             return nextTime
-    #         else:
-    #             return time + nSteps * dt
-
 
 class CalendarError(Exception):  # TODO: Move to a parcels errors module
     """Exception raised when the calendar of a field is not compatible with the rest of the Fields. The user should ensure that they only add fields to a FieldSet that have compatible CFtime calendars."""

From dc1856012797a74007b5dddae9efbda302c49bc6 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 11:16:47 +0200
Subject: [PATCH 03/12] Remove fieldfilebuffer.py

Wasn't used anymore in the codebase
---
 parcels/field.py           |   8 --
 parcels/fieldfilebuffer.py | 184 -------------------------------------
 2 files changed, 192 deletions(-)
 delete mode 100644 parcels/fieldfilebuffer.py

diff --git a/parcels/field.py b/parcels/field.py
index 3268f11c4e..e328d19b23 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -384,14 +384,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.
 
diff --git a/parcels/fieldfilebuffer.py b/parcels/fieldfilebuffer.py
deleted file mode 100644
index c1a0e7dd12..0000000000
--- a/parcels/fieldfilebuffer.py
+++ /dev/null
@@ -1,184 +0,0 @@
-import datetime
-import warnings
-from pathlib import Path
-
-import numpy as np
-import xarray as xr
-
-from parcels._typing import InterpMethodOption, PathLike
-from parcels.tools.warnings import FileWarning
-
-
-class NetcdfFileBuffer:
-    def __init__(
-        self,
-        filename,
-        dimensions,
-        indices,
-        *,
-        interp_method: InterpMethodOption = "linear",
-        gridindexingtype="nemo",
-    ):
-        self.filename: PathLike | list[PathLike] = filename
-        self.dimensions = dimensions  # Dict with dimension keys for file data
-        self.indices = indices
-        self.dataset = None
-        self.interp_method = interp_method
-        self.gridindexingtype = gridindexingtype
-
-    def __enter__(self):
-        self.dataset = open_xarray_dataset(self.filename)
-        return self
-
-    def __exit__(self, type, value, traceback):
-        self.close()
-
-    def close(self):
-        if self.dataset is not None:
-            self.dataset.close()
-            self.dataset = None
-
-    @property
-    def xdim(self):
-        lon = self.dataset[self.dimensions["lon"]]
-        xdim = lon.size if len(lon.shape) == 1 else lon.shape[-1]
-        if self.gridindexingtype in ["croco"]:
-            xdim -= 1
-        return xdim
-
-    @property
-    def ydim(self):
-        lat = self.dataset[self.dimensions["lat"]]
-        ydim = lat.size if len(lat.shape) == 1 else lat.shape[-2]
-        if self.gridindexingtype in ["croco"]:
-            ydim -= 1
-        return ydim
-
-    @property
-    def zdim(self):
-        if "depth" not in self.dimensions:
-            return 1
-        depth = self.dataset[self.dimensions["depth"]]
-        zdim = depth.size if len(depth.shape) == 1 else depth.shape[-3]
-        if self.gridindexingtype in ["croco"]:
-            zdim -= 1
-        return zdim
-
-    @property
-    def latlon(self):
-        lon = self.dataset[self.dimensions["lon"]]
-        lat = self.dataset[self.dimensions["lat"]]
-        if self.gridindexingtype not in ["croco"]:
-            if len(lon.shape) == 3:  # some lon, lat have a time dimension 1
-                lon = lon[0, :, :]
-                lat = lat[0, :, :]
-            elif len(lon.shape) == 4:  # some lon, lat have a time and depth dimension 1
-                lon = lon[0, 0, :, :]
-                lat = lat[0, 0, :, :]
-
-        if len(lon.shape) > 1:
-            if is_rectilinear(lon, lat):
-                lon = lon[0, :]
-                lat = lat[:, 0]
-        return lat, lon
-
-    @property
-    def depth(self):
-        self.indices["depth"] = range(self.zdim)
-        if "depth" in self.dimensions:
-            return self.dataset[self.dimensions["depth"]]
-        return np.zeros(1)
-
-    @property
-    def data_full_zdim(self):
-        return self.zdim
-
-    def _check_extend_depth(self, data, dim):
-        return (
-            self.indices["depth"][-1] == self.data_full_zdim - 1
-            and data.shape[dim] == self.data_full_zdim - 1
-            and self.interp_method in ["bgrid_velocity", "bgrid_w_velocity", "bgrid_tracer"]
-        )
-
-    def _apply_indices(self, data, ti):
-        if len(data.shape) == 1:
-            if self.indices["depth"] is not None:
-                data = data[self.indices["depth"]]
-        elif len(data.shape) == 3:
-            if self._check_extend_depth(data, 0):
-                data = data[self.indices["depth"][:-1], :, :]
-            elif len(self.indices["depth"]) > 1:
-                data = data[self.indices["depth"], :, :]
-            else:
-                data = data[ti, :, :]
-        else:
-            if self._check_extend_depth(data, 1):
-                data = data[ti, self.indices["depth"][:-1], :, :]
-            else:
-                data = data[ti, self.indices["depth"], :, :]
-        return data
-
-    @property
-    def data(self):
-        return self.data_access()
-
-    def data_access(self):
-        data = self.dataset[self.name]
-        ti = range(data.shape[0])
-        return np.array(self._apply_indices(data, ti))
-
-    @property
-    def time(self):
-        return self.time_access()
-
-    def time_access(self):
-        if "time" not in self.dimensions:
-            return np.array([None])
-
-        time_da = self.dataset[self.dimensions["time"]]
-        time = (
-            np.array([time_da[self.dimensions["time"]].data])
-            if len(time_da.shape) == 0
-            else np.array(time_da[self.dimensions["time"]])
-        )
-        if isinstance(time[0], datetime.datetime):
-            raise NotImplementedError(
-                "Parcels currently only parses dates ranging from 1678 AD to 2262 AD, which are stored by xarray as np.datetime64. If you need a wider date range, please open an Issue on the parcels github page."
-            )
-        return time
-
-
-def open_xarray_dataset(filename: Path | str) -> xr.Dataset:
-    try:
-        # Unfortunately we need to do if-else here, cause the lock-parameter is either False or a Lock-object
-        # (which we would rather want to have being auto-managed).
-        # If 'lock' is not specified, the Lock-object is auto-created and managed by xarray internally.
-        ds = xr.open_mfdataset(filename, decode_cf=True)
-        ds["decoded"] = True
-    except:
-        warnings.warn(  # TODO: Is this warning necessary? What cases does this except block get triggered - is it to do with the bare except???
-            f"File {filename} could not be decoded properly by xarray (version {xr.__version__}). "
-            "It will be opened with no decoding. Filling values might be wrongly parsed.",
-            FileWarning,
-            stacklevel=2,
-        )
-
-        ds = xr.open_mfdataset(filename, decode_cf=False)
-        ds["decoded"] = False
-    return ds
-
-
-def is_rectilinear(lon_subset, lat_subset) -> bool:
-    """Test if all columns and rows are the same for lon and lat (in which case grid is rectilinear).
-
-    lon_subset and lat_subset are 2D numpy arrays
-    """
-    for xi in range(1, lon_subset.shape[0]):
-        if not np.allclose(lon_subset[0, :], lon_subset[xi, :]):
-            return False
-
-    for yi in range(1, lat_subset.shape[1]):
-        if not np.allclose(lat_subset[:, 0], lat_subset[:, yi]):
-            return False
-
-    return True

From d00880fceaa57f86fd5f83b6e85a30b92dc9dec9 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 11:34:39 +0200
Subject: [PATCH 04/12] Remove Field._rescale_and_set_minmax

---
 parcels/field.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/parcels/field.py b/parcels/field.py
index e328d19b23..27c3dfb9cf 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -316,10 +316,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)
 

From 389373002b6cdfad86d40b6f6e24b864d8ddf1d6 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 11:41:40 +0200
Subject: [PATCH 05/12] Remove __contains__ and __getattr__ from Field

Overloaded behaviour that wasn't justified/used in the rest of the codebase.
---
 parcels/field.py                  |  6 ------
 tests/v4/test_uxarray_fieldset.py | 12 ++++++------
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/parcels/field.py b/parcels/field.py
index 27c3dfb9cf..bbbc267d50 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -316,9 +316,6 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         else:
             return value
 
-    def __getattr__(self, key: str):
-        return getattr(self.data, key)
-
     def __getitem__(self, key):
         self._check_velocitysampling()
         try:
@@ -329,9 +326,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."""
diff --git a/tests/v4/test_uxarray_fieldset.py b/tests/v4/test_uxarray_fieldset.py
index f1ed7262fe..3399740116 100644
--- a/tests/v4/test_uxarray_fieldset.py
+++ b/tests/v4/test_uxarray_fieldset.py
@@ -77,16 +77,16 @@ def uvw_fesom_channel(ds_fesom_channel) -> VectorField:
 def test_fesom_fieldset(ds_fesom_channel, uv_fesom_channel):
     fieldset = FieldSet([uv_fesom_channel, uv_fesom_channel.U, uv_fesom_channel.V])
     # Check that the fieldset has the expected properties
-    assert (fieldset.U == ds_fesom_channel.U).all()
-    assert (fieldset.V == ds_fesom_channel.V).all()
+    assert (fieldset.U.data == ds_fesom_channel.U).all()
+    assert (fieldset.V.data == ds_fesom_channel.V).all()
 
 
 def test_fesom_in_particleset(ds_fesom_channel, uv_fesom_channel):
     fieldset = FieldSet([uv_fesom_channel, uv_fesom_channel.U, uv_fesom_channel.V])
 
     # Check that the fieldset has the expected properties
-    assert (fieldset.U == ds_fesom_channel.U).all()
-    assert (fieldset.V == ds_fesom_channel.V).all()
+    assert (fieldset.U.data == ds_fesom_channel.U).all()
+    assert (fieldset.V.data == ds_fesom_channel.V).all()
     pset = ParticleSet(fieldset, pclass=Particle)
     assert pset.fieldset == fieldset
 
@@ -94,8 +94,8 @@ def test_fesom_in_particleset(ds_fesom_channel, uv_fesom_channel):
 def test_set_interp_methods(ds_fesom_channel, uv_fesom_channel):
     fieldset = FieldSet([uv_fesom_channel, uv_fesom_channel.U, uv_fesom_channel.V])
     # Check that the fieldset has the expected properties
-    assert (fieldset.U == ds_fesom_channel.U).all()
-    assert (fieldset.V == ds_fesom_channel.V).all()
+    assert (fieldset.U.data == ds_fesom_channel.U).all()
+    assert (fieldset.V.data == ds_fesom_channel.V).all()
 
     # Set the interpolation method for each field
     fieldset.U.interp_method = UXPiecewiseConstantFace

From 1b9fcbff34327699d82c03c66898612ca94c2adc Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 11:44:33 +0200
Subject: [PATCH 06/12] Make internal helper functions private

---
 parcels/field.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/parcels/field.py b/parcels/field.py
index bbbc267d50..e790fddcb2 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -158,7 +158,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?"
@@ -340,9 +340,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
 
@@ -480,14 +480,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
 

From 25c589bcb4355cf72c1f07151233c4f8391c83c0 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Mon, 7 Jul 2025 16:34:29 +0200
Subject: [PATCH 07/12] Add uxarray to pyproject dep list

Was in environment.yaml file - wasn't in the pyproject.toml file though
---
 pyproject.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index daee2826ce..8fde067a20 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = [
   "pytest",
   "scipy",
   "xarray",
+  "uxarray",
   "pooch",
 ]
 
@@ -68,6 +69,7 @@ pymbolic = "*"
 scipy = ">=0.16.0"
 tqdm = "*"
 xarray = ">=0.10.8"
+uxarray = ">=2025.3.0"
 cftime = ">=1.3.1"
 dask = ">=2.0"
 scikit-learn = "*"

From ce49751ed52f2f4c568fc9c0af1fb08543c95632 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 14:26:40 +0200
Subject: [PATCH 08/12] Remove lon, lat, and depth from Field class

These variables mainly apply to structured grids. With the new API, they aren't needed in the codebase (we can add this on the repective grid classes if needed).
---
 parcels/field.py | 36 ------------------------------------
 1 file changed, 36 deletions(-)

diff --git a/parcels/field.py b/parcels/field.py
index e790fddcb2..a93eb2cc00 100644
--- a/parcels/field.py
+++ b/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,
@@ -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:

From f6228ff1e5d028d1121742de4c2887c8ef697225 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 14:45:41 +0200
Subject: [PATCH 09/12] Add get_axis_dim and axes to BaseGrid, XGrid, and
 UxGrid

Also added some simple tests in test_uxgrid.py
---
 parcels/basegrid.py     | 35 +++++++++++++++++++++++++++++++++++
 parcels/uxgrid.py       | 13 +++++++++++++
 parcels/xgrid.py        | 31 ++++++++++++++++++-------------
 tests/v4/test_uxgrid.py | 23 +++++++++++++++++++++++
 tests/v4/test_xgrid.py  | 21 +++++++++++++--------
 5 files changed, 102 insertions(+), 21 deletions(-)
 create mode 100644 tests/v4/test_uxgrid.py

diff --git a/parcels/basegrid.py b/parcels/basegrid.py
index e1a866e225..bac0374301 100644
--- a/parcels/basegrid.py
+++ b/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.
+        """
+        ...
diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py
index 641213b7a0..2acab9c1ce 100644
--- a/parcels/uxgrid.py
+++ b/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
 
diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 21d3101851..50e0e134d5 100644
--- a/parcels/xgrid.py
+++ b/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)
diff --git a/tests/v4/test_uxgrid.py b/tests/v4/test_uxgrid.py
new file mode 100644
index 0000000000..9c675b278b
--- /dev/null
+++ b/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
diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 301a16c131..9de07dcbf3 100644
--- a/tests/v4/test_xgrid.py
+++ b/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():

From 40503c7bd41e6abbc1544cc0a9d30d6dca093f1f Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 15:54:23 +0200
Subject: [PATCH 10/12] Generalise ravel/unravel to any grid

---
 parcels/basegrid.py       | 77 +++++++++++++++++++++++++++++++++++++--
 parcels/uxgrid.py         |  8 ----
 parcels/xgrid.py          | 16 --------
 tests/v4/test_basegrid.py | 53 +++++++++++++++++++++++++++
 tests/v4/test_xgrid.py    | 24 ------------
 5 files changed, 126 insertions(+), 52 deletions(-)
 create mode 100644 tests/v4/test_basegrid.py

diff --git a/parcels/basegrid.py b/parcels/basegrid.py
index bac0374301..b0ef53d7f7 100644
--- a/parcels/basegrid.py
+++ b/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,7 +136,9 @@ 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
@@ -142,6 +146,9 @@ 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]
diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py
index 2acab9c1ce..e95eb84b53 100644
--- a/parcels/uxgrid.py
+++ b/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}
diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 50e0e134d5..00e9fcd7ac 100644
--- a/parcels/xgrid.py
+++ b/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."""
diff --git a/tests/v4/test_basegrid.py b/tests/v4/test_basegrid.py
new file mode 100644
index 0000000000..86a4f12ce2
--- /dev/null
+++ b/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."
diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 9de07dcbf3..29b0ed6059 100644
--- a/tests/v4/test_xgrid.py
+++ b/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",
     [

From 38bd14fabda6174e1436258eb491db85f9eae641 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 10 Jul 2025 16:18:52 +0200
Subject: [PATCH 11/12] Add back test cases for XGrid xdim,ydim,zdim

---
 tests/v4/test_xgrid.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 29b0ed6059..a177ce8904 100644
--- a/tests/v4/test_xgrid.py
+++ b/tests/v4/test_xgrid.py
@@ -16,6 +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),
 ]
 
 

From 0de90ebe3e897630b5aaeb89af635e0a45ad1b6f Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Fri, 11 Jul 2025 12:52:28 +0200
Subject: [PATCH 12/12] Update XGrid axis dim to use cell count

---
 parcels/basegrid.py    |  2 +-
 parcels/xgrid.py       |  8 +++-----
 tests/v4/test_xgrid.py | 12 ++++++------
 3 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/parcels/basegrid.py b/parcels/basegrid.py
index b0ef53d7f7..2765dd59ec 100644
--- a/parcels/basegrid.py
+++ b/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
         ----------
diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 00e9fcd7ac..c3c110485c 100644
--- a/parcels/xgrid.py
+++ b/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]:
diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index a177ce8904..2af1e1f019 100644
--- a/tests/v4/test_xgrid.py
+++ b/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():