Repository: OceanParcels/Parcels (python)

Issue:
StructuredGrid support for dimensions with length-one
Following the discussion at https://github.com/OceanParcels/Parcels/pull/2037#issuecomment-3003488793, I'd suggest that we consider to actively coerce all structured-grid fields to 4D `(time, depth, lat, lon)`, and then allow for any of the dimensions to have length-1. These length-1 dimensions don't have to coordinates. 

Then, in the `index_search`, any length-1 dimension will return a `(0, 0)` for grid-index and barycentric coordinate. Note that a custom interpolator may then have to be provided

Note that this is effectively already how fields without depth or without time are treated. We would then extend it also to longitude and/or latitude



Pull request description:
Update structured Field data ingestion transposing to 4D (TZYX) and add `XGrid.from_dataset()` 
<!-- Feel free to remove list items that are not relevant for your changes. -->

Changes:
- Add `XGrid.from_dataset` (rather than having to construct a xgcm grid first). This execution path drops all the Field data from the dataset before constructing the xgcm grid (and also makes it easier to call for devs as its 1 less import). Seeb73813c769aeaa92be181bf3fee1ba5275d317ce commit message for more info
- Update Field init to transpose data to 4D array, inserting additional dummy axes where necessary

- [x] Chose the correct base branch (`v4-dev` for v4 changes)
- [x] Fixes #2047 and works towards #2054
- [x] Added tests (for the transposing to 4D I didn't add a unit test, but added an integration test over all the generic datasets)
- [x] Added documentation

Will do a self review tomorrow then mark ready for review.


Patch:
diff --git a/parcels/field.py b/parcels/field.py
index dba61dbd..1a1fbb71 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -29,7 +29,7 @@ from parcels.tools.statuscodes import (
     _raise_field_out_of_bound_error,
 )
 from parcels.uxgrid import UxGrid
-from parcels.xgrid import XGrid
+from parcels.xgrid import XGrid, _transpose_xfield_data_to_tzyx
 
 from ._index_search import _search_time_index
 
@@ -146,6 +146,9 @@ class Field:
 
         _assert_compatible_combination(data, grid)
 
+        if isinstance(grid, XGrid):
+            data = _transpose_xfield_data_to_tzyx(data, grid.xgcm_grid)
+
         self.name = name
         self.data = data
         self.grid = grid
@@ -186,8 +189,9 @@ class Field:
         else:
             raise ValueError("Unsupported mesh type in data array attributes. Choose either: 'spherical' or 'flat'")
 
-        if "time" not in self.data.dims:
-            raise ValueError("Field is missing a 'time' dimension. ")
+        if self.data.shape[0] > 1:
+            if "time" not in self.data.coords:
+                raise ValueError("Field data is missing a 'time' coordinate.")
 
     @property
     def units(self):
@@ -439,7 +443,7 @@ def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux
 
 
 def _get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None:
-    if len(data.time) == 1:
+    if data.shape[0] == 1:
         return None
 
     return TimeInterval(data.time.values[0], data.time.values[-1])
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index c423000f..5937c7f5 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -132,7 +132,6 @@ class FieldSet:
         """
         da = xr.DataArray(
             data=np.full((1, 1, 1, 1), value),
-            dims=["time", "ZG", "YG", "XG"],
         )
         grid = XGrid(xgcm.Grid(da))
         self.add_field(
diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 2cfda507..81391ee8 100644
--- a/parcels/xgrid.py
+++ b/parcels/xgrid.py
@@ -1,4 +1,4 @@
-from collections.abc import Hashable, Mapping
+from collections.abc import Hashable, Mapping, Sequence
 from functools import cached_property
 from typing import Literal, cast
 
@@ -10,13 +10,17 @@ from parcels import xgcm
 from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.basegrid import BaseGrid
 
-_XGRID_AXES_ORDERING = "ZYX"
 _XGRID_AXES = Literal["X", "Y", "Z"]
+_XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX"
 
 _XGCM_AXIS_DIRECTION = Literal["X", "Y", "Z", "T"]
 _XGCM_AXIS_POSITION = Literal["center", "left", "right", "inner", "outer"]
 _XGCM_AXES = Mapping[_XGCM_AXIS_DIRECTION, xgcm.Axis]
 
+_FIELD_DATA_ORDERING: Sequence[_XGCM_AXIS_DIRECTION] = "TZYX"
+
+_DEFAULT_XGCM_KWARGS = {"periodic": False}
+
 
 def get_cell_count_along_dim(axis: xgcm.Axis) -> int:
     first_coord = list(axis.coords.items())[0]
@@ -34,6 +38,48 @@ def _get_xgrid_axes(grid: xgcm.Grid) -> list[_XGRID_AXES]:
     return sorted(spatial_axes, key=_XGRID_AXES_ORDERING.index)
 
 
+def _drop_field_data(ds: xr.Dataset) -> xr.Dataset:
+    """
+    Removes DataArrays from the dataset that are associated with field data so that
+    when passed to the XGCM grid, the object only functions as an in memory representation
+    of the grid.
+    """
+    return ds.drop_vars(ds.data_vars)
+
+
+def _transpose_xfield_data_to_tzyx(da: xr.DataArray, xgcm_grid: xgcm.Grid) -> xr.DataArray:
+    """
+    Transpose a DataArray of any shape into a 4D array of order TZYX. Uses xgcm to determine
+    the axes, and inserts mock dimensions of size 1 for any axes not present in the DataArray.
+    """
+    ax_dims = [(get_axis_from_dim_name(xgcm_grid.axes, dim), dim) for dim in da.dims]
+
+    if all(ax_dim[0] is None for ax_dim in ax_dims):
+        # Assuming its a 1D constant field (hence has no axes)
+        assert da.shape == (1, 1, 1, 1)
+        return da.rename({old_dim: f"mock{axis}" for old_dim, axis in zip(da.dims, _FIELD_DATA_ORDERING, strict=True)})
+
+    # All dimensions must be associated with an axis in the grid
+    if any(ax_dim[0] is None for ax_dim in ax_dims):
+        raise ValueError(
+            f"DataArray {da.name!r} with dims {da.dims} has dimensions that are not associated with a direction on the provided grid."
+        )
+
+    axes_not_in_field = set(_FIELD_DATA_ORDERING) - set(ax_dim[0] for ax_dim in ax_dims)
+
+    mock_dims_to_create = {}
+    for ax in axes_not_in_field:
+        mock_dims_to_create[f"mock{ax}"] = 1
+        ax_dims.append((ax, f"mock{ax}"))
+
+    if mock_dims_to_create:
+        da = da.expand_dims(mock_dims_to_create, create_index_for_new_dim=False)
+
+    ax_dims = sorted(ax_dims, key=lambda x: _FIELD_DATA_ORDERING.index(x[0]))
+
+    return da.transpose(*[ax_dim[1] for ax_dim in ax_dims])
+
+
 class XGrid(BaseGrid):
     """
     Class to represent a structured grid in Parcels. Wraps a xgcm-like Grid object (we use a trimmed down version of the xgcm.Grid class that is vendored with Parcels).
@@ -53,6 +99,18 @@ class XGrid(BaseGrid):
         if len(set(grid.axes) & {"X", "Y", "Z"}) > 0:  # Only if spatial grid is >0D (see #2054 for further development)
             assert_valid_lat_lon(ds["lat"], ds["lon"], grid.axes)
 
+    @classmethod
+    def from_dataset(cls, ds: xr.Dataset, mesh="flat", xgcm_kwargs=None):
+        """WARNING: unstable API, subject to change in future versions."""  # TODO v4: make private or remove warning on v4 release
+        if xgcm_kwargs is None:
+            xgcm_kwargs = {}
+
+        xgcm_kwargs = {**_DEFAULT_XGCM_KWARGS, **xgcm_kwargs}
+
+        ds = _drop_field_data(ds)
+        grid = xgcm.Grid(ds, **xgcm_kwargs)
+        return cls(grid, mesh=mesh)
+
     @property
     def axes(self) -> list[_XGRID_AXES]:
         return _get_xgrid_axes(self.xgcm_grid)


Test patch:
diff --git a/tests/utils.py b/tests/utils.py
index e0fd5f34..a76b7e8a 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,12 +1,19 @@
 """General helper functions and utilies for test suite."""
 
+from __future__ import annotations
+
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 import xarray as xr
 
 import parcels
 from parcels import FieldSet
+from parcels.xgrid import _FIELD_DATA_ORDERING, get_axis_from_dim_name
+
+if TYPE_CHECKING:
+    from parcels.xgrid import XGrid
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
 TEST_ROOT = PROJECT_ROOT / "tests"
@@ -116,3 +123,13 @@ def create_fieldset_zeros_simple(xdim=40, ydim=100, withtime=False):
 
 def assert_empty_folder(path: Path):
     assert [p.name for p in path.iterdir()] == []
+
+
+def assert_valid_field_data(data: xr.DataArray, grid: XGrid):
+    assert len(data.shape) == 4, f"Field data should have 4 dimensions (time, depth, lat, lon), got dims {data.dims}"
+
+    for ax_expected, dim in zip(_FIELD_DATA_ORDERING, data.dims, strict=True):
+        ax_actual = get_axis_from_dim_name(grid.xgcm_grid.axes, dim)
+        if ax_actual is None:
+            continue  # None is ok
+        assert ax_actual == ax_expected, f"Expected axis {ax_expected} for dimension '{dim}', got {ax_actual}"
diff --git a/tests/v4/test_datasets.py b/tests/v4/test_datasets.py
index e5201e58..d2d21040 100644
--- a/tests/v4/test_datasets.py
+++ b/tests/v4/test_datasets.py
@@ -1,11 +1,11 @@
+from parcels import xgcm
 from parcels._datasets.structured.generic import datasets
-from parcels.xgcm import Grid
 
 
 def test_left_indexed_dataset():
     """Checks that 'ds_2d_left' is right indexed on all variables."""
     ds = datasets["ds_2d_left"]
-    grid = Grid(ds)
+    grid = xgcm.Grid(ds)
 
     for _axis_name, axis in grid.axes.items():
         for pos, _dim_name in axis.coords.items():
@@ -15,7 +15,7 @@ def test_left_indexed_dataset():
 def test_right_indexed_dataset():
     """Checks that 'ds_2d_right' is right indexed on all variables."""
     ds = datasets["ds_2d_right"]
-    grid = Grid(ds)
+    grid = xgcm.Grid(ds)
     for _axis_name, axis in grid.axes.items():
         for pos, _dim_name in axis.coords.items():
             assert pos in ["center", "right"]
diff --git a/tests/v4/test_field.py b/tests/v4/test_field.py
index a555a672..6f106889 100644
--- a/tests/v4/test_field.py
+++ b/tests/v4/test_field.py
@@ -5,7 +5,7 @@ import pytest
 import uxarray as ux
 import xarray as xr
 
-from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, VectorField, xgcm
+from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, VectorField
 from parcels._datasets.structured.generic import T as T_structured
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
@@ -15,7 +15,7 @@ from parcels.xgrid import XGrid
 
 def test_field_init_param_types():
     data = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(data))
+    grid = XGrid.from_dataset(data)
     with pytest.raises(ValueError, match="Expected `name` to be a string"):
         Field(name=123, data=data["data_g"], grid=grid)
 
@@ -32,7 +32,7 @@ def test_field_init_param_types():
 @pytest.mark.parametrize(
     "data,grid",
     [
-        pytest.param(ux.UxDataArray(), XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])), id="uxdata-grid"),
+        pytest.param(ux.UxDataArray(), XGrid.from_dataset(datasets_structured["ds_2d_left"]), id="uxdata-grid"),
         pytest.param(
             xr.DataArray(),
             UxGrid(
@@ -57,7 +57,7 @@ def test_field_incompatible_combination(data, grid):
     [
         pytest.param(
             datasets_structured["ds_2d_left"]["data_g"],
-            XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])),
+            XGrid.from_dataset(datasets_structured["ds_2d_left"]),
             id="ds_2d_left",
         ),  # TODO: Perhaps this test should be expanded to cover more datasets?
     ],
@@ -80,10 +80,10 @@ def test_field_init_fail_on_float_time_dim():
     (users are expected to use timedelta64 or datetime).
     """
     ds = datasets_structured["ds_2d_left"].copy()
-    ds["time"] = np.arange(0, T_structured, dtype="float64")
+    ds["time"] = (ds["time"].dims, np.arange(0, T_structured, dtype="float64"), ds["time"].attrs)
 
     data = ds["data_g"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     with pytest.raises(
         ValueError,
         match="Error getting time interval.*. Are you sure that the time dimension on the xarray dataset is stored as timedelta, datetime or cftime datetime objects\?",
@@ -100,7 +100,7 @@ def test_field_init_fail_on_float_time_dim():
     [
         pytest.param(
             datasets_structured["ds_2d_left"]["data_g"],
-            XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])),
+            XGrid.from_dataset(datasets_structured["ds_2d_left"]),
             id="ds_2d_left",
         ),
     ],
@@ -119,7 +119,7 @@ def test_vectorfield_init_different_time_intervals():
 
 def test_field_invalid_interpolator():
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
 
     def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid):
         return 0.0
@@ -131,7 +131,7 @@ def test_field_invalid_interpolator():
 
 def test_vectorfield_invalid_interpolator():
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
 
     def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid):
         return 0.0
diff --git a/tests/v4/test_fieldset.py b/tests/v4/test_fieldset.py
index f623b7f1..c3871018 100644
--- a/tests/v4/test_fieldset.py
+++ b/tests/v4/test_fieldset.py
@@ -5,7 +5,6 @@ import numpy as np
 import pytest
 import xarray as xr
 
-from parcels import xgcm
 from parcels._datasets.structured.circulation_models import (
     datasets as datasets_circulation_models,  # noqa: F401
 )  # just making sure the import works. Will eventually be used in tests
@@ -14,6 +13,7 @@ from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.field import Field, VectorField
 from parcels.fieldset import CalendarError, FieldSet, _datetime_to_msg
 from parcels.xgrid import XGrid
+from tests import utils
 
 ds = datasets_structured["ds_2d_left"]
 
@@ -21,7 +21,7 @@ ds = datasets_structured["ds_2d_left"]
 @pytest.fixture
 def fieldset() -> FieldSet:
     """Fixture to create a FieldSet object for testing."""
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     UV = VectorField("UV", U, V)
@@ -55,7 +55,7 @@ def test_fieldset_add_constant_field(fieldset):
 
 
 def test_fieldset_add_field(fieldset):
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     field = Field("test_field", ds["U (A grid)"], grid, mesh_type="flat")
     fieldset.add_field(field)
     assert fieldset.test_field == field
@@ -68,7 +68,7 @@ def test_fieldset_add_field_wrong_type(fieldset):
 
 
 def test_fieldset_add_field_already_exists(fieldset):
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     field = Field("test_field", ds["U (A grid)"], grid, mesh_type="flat")
     fieldset.add_field(field, "test_field")
     with pytest.raises(ValueError, match="FieldSet already has a Field with name 'test_field'"):
@@ -85,16 +85,32 @@ def test_fieldset_gridset(fieldset):
     assert len(fieldset.gridset) == 2
 
 
+@pytest.mark.parametrize("ds", [pytest.param(ds, id=k) for k, ds in datasets_structured.items()])
+def test_fieldset_from_structured_generic_datasets(ds):
+    grid = XGrid.from_dataset(ds)
+    fields = []
+    for var in ds.data_vars:
+        fields.append(Field(var, ds[var], grid, mesh_type="flat"))
+
+    fieldset = FieldSet(fields)
+
+    assert len(fieldset.fields) == len(ds.data_vars)
+    for field in fieldset.fields.values():
+        utils.assert_valid_field_data(field.data, field.grid)
+
+    assert len(fieldset.gridset) == 1
+
+
 def test_fieldset_gridset_multiple_grids(): ...
 
 
 def test_fieldset_time_interval():
-    grid1 = XGrid(xgcm.Grid(ds))
+    grid1 = XGrid.from_dataset(ds)
     field1 = Field("field1", ds["U (A grid)"], grid1, mesh_type="flat")
 
     ds2 = ds.copy()
-    ds2["time"] = ds2["time"] + np.timedelta64(timedelta(days=1))
-    grid2 = XGrid(xgcm.Grid(ds2))
+    ds2["time"] = (ds2["time"].dims, ds2["time"].data + np.timedelta64(timedelta(days=1)), ds2["time"].attrs)
+    grid2 = XGrid.from_dataset(ds2)
     field2 = Field("field2", ds2["U (A grid)"], grid2, mesh_type="flat")
 
     fieldset = FieldSet([field1, field2])
@@ -114,16 +130,24 @@ def test_fieldset_time_interval_constant_fields():
 
 def test_fieldset_init_incompatible_calendars():
     ds1 = ds.copy()
-    ds1["time"] = xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True)
+    ds1["time"] = (
+        ds1["time"].dims,
+        xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True),
+        ds1["time"].attrs,
+    )
 
-    grid = XGrid(xgcm.Grid(ds1))
+    grid = XGrid.from_dataset(ds1)
     U = Field("U", ds1["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds1["V (A grid)"], grid, mesh_type="flat")
     UV = VectorField("UV", U, V)
 
     ds2 = ds.copy()
-    ds2["time"] = xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True)
-    grid2 = XGrid(xgcm.Grid(ds2))
+    ds2["time"] = (
+        ds2["time"].dims,
+        xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True),
+        ds2["time"].attrs,
+    )
+    grid2 = XGrid.from_dataset(ds2)
     incompatible_calendar = Field("test", ds2["data_g"], grid2, mesh_type="flat")
 
     with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"):
@@ -132,16 +156,24 @@ def test_fieldset_init_incompatible_calendars():
 
 def test_fieldset_add_field_incompatible_calendars(fieldset):
     ds_test = ds.copy()
-    ds_test["time"] = xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True)
-    grid = XGrid(xgcm.Grid(ds_test))
+    ds_test["time"] = (
+        ds_test["time"].dims,
+        xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True),
+        ds_test["time"].attrs,
+    )
+    grid = XGrid.from_dataset(ds_test)
     field = Field("test_field", ds_test["data_g"], grid, mesh_type="flat")
 
     with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"):
         fieldset.add_field(field, "test_field")
 
     ds_test = ds.copy()
-    ds_test["time"] = np.linspace(0, 100, T_structured, dtype="timedelta64[s]")
-    grid = XGrid(xgcm.Grid(ds_test))
+    ds_test["time"] = (
+        ds_test["time"].dims,
+        np.linspace(0, 100, T_structured, dtype="timedelta64[s]"),
+        ds_test["time"].attrs,
+    )
+    grid = XGrid.from_dataset(ds_test)
     field = Field("test_field", ds_test["data_g"], grid, mesh_type="flat")
 
     with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"):
diff --git a/tests/v4/test_index_search.py b/tests/v4/test_index_search.py
index 7f9290f1..5d5c602b 100644
--- a/tests/v4/test_index_search.py
+++ b/tests/v4/test_index_search.py
@@ -1,7 +1,6 @@
 import numpy as np
 import pytest
 
-from parcels import xgcm
 from parcels._datasets.structured.generic import datasets
 from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.field import Field
@@ -13,7 +12,7 @@ from parcels.xgrid import (
 @pytest.fixture
 def field_cone():
     ds = datasets["2d_left_unrolled_cone"]
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     field = Field(
         name="test_field",
         data=ds["data_g"],
diff --git a/tests/v4/test_kernel.py b/tests/v4/test_kernel.py
index 1125d925..cb521bc7 100644
--- a/tests/v4/test_kernel.py
+++ b/tests/v4/test_kernel.py
@@ -6,7 +6,6 @@ from parcels import (
     Field,
     FieldSet,
     ParticleSet,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.xgrid import XGrid
@@ -16,7 +15,7 @@ from tests.common_kernels import MoveEast, MoveNorth
 @pytest.fixture
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     return FieldSet([U, V])
diff --git a/tests/v4/test_particleset.py b/tests/v4/test_particleset.py
index 7a9dd7be..ace5e8db 100644
--- a/tests/v4/test_particleset.py
+++ b/tests/v4/test_particleset.py
@@ -13,7 +13,6 @@ from parcels import (
     ParticleSet,
     ParticleSetWarning,
     Variable,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.xgrid import XGrid
@@ -23,7 +22,7 @@ from tests.common_kernels import DoNothing
 @pytest.fixture
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     return FieldSet([U, V])
diff --git a/tests/v4/test_particleset_execute.py b/tests/v4/test_particleset_execute.py
index f63210ad..b1994d4d 100644
--- a/tests/v4/test_particleset_execute.py
+++ b/tests/v4/test_particleset_execute.py
@@ -10,7 +10,6 @@ from parcels import (
     StatusCode,
     UXPiecewiseConstantFace,
     VectorField,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
@@ -22,7 +21,7 @@ from tests.common_kernels import DoNothing
 @pytest.fixture
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     return FieldSet([U, V])
diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 210ca6b7..30bcc7c1 100644
--- a/tests/v4/test_xgrid.py
+++ b/tests/v4/test_xgrid.py
@@ -1,3 +1,4 @@
+import itertools
 from collections import namedtuple
 
 import numpy as np
@@ -5,11 +6,11 @@ import pytest
 import xarray as xr
 from numpy.testing import assert_allclose
 
-from parcels import xgcm
 from parcels._datasets.structured.generic import X, Y, Z, datasets
-from parcels.xgrid import XGrid, _search_1d_array
+from parcels.xgrid import XGrid, _search_1d_array, _transpose_xfield_data_to_tzyx
+from tests import utils
 
-GridTestCase = namedtuple("GridTestCase", ["Grid", "attr", "expected"])
+GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"])
 
 test_cases = [
     GridTestCase(datasets["ds_2d_left"], "lon", datasets["ds_2d_left"].XG.values),
@@ -34,25 +35,39 @@ def assert_equal(actual, expected):
 
 @pytest.mark.parametrize("ds, attr, expected", test_cases)
 def test_xgrid_properties_ground_truth(ds, attr, expected):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     actual = getattr(grid, attr)
     assert_equal(actual, expected)
 
 
 @pytest.mark.parametrize("ds", [pytest.param(ds, id=key) for key, ds in datasets.items()])
-def test_xgrid_init_on_generic_datasets(ds):
-    XGrid(xgcm.Grid(ds, periodic=False))
+def test_xgrid_from_dataset_on_generic_datasets(ds):
+    XGrid.from_dataset(ds)
 
 
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_xgrid_axes(ds):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     assert grid.axes == ["Z", "Y", "X"]
 
 
+@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
+def test_transpose_xfield_data_to_tzyx(ds):
+    da = ds["data_g"]
+    grid = XGrid.from_dataset(ds)
+
+    all_combinations = (itertools.combinations(da.dims, n) for n in range(len(da.dims)))
+    all_combinations = itertools.chain(*all_combinations)
+    for subset_dims in all_combinations:
+        isel = {dim: 0 for dim in subset_dims}
+        da_subset = da.isel(isel, drop=True)
+        da_test = _transpose_xfield_data_to_tzyx(da_subset, grid.xgcm_grid)
+        utils.assert_valid_field_data(da_test, grid)
+
+
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_xgrid_get_axis_dim(ds):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     assert grid.get_axis_dim("Z") == Z - 1
     assert grid.get_axis_dim("Y") == Y - 1
     assert grid.get_axis_dim("X") == X - 1
@@ -72,7 +87,7 @@ def test_invalid_lon_lat():
         ValueError,
         match=".*is defined on the center of the grid, but must be defined on the F points\.",
     ):
-        XGrid(xgcm.Grid(ds, periodic=False))
+        XGrid.from_dataset(ds)
 
     ds = datasets["ds_2d_left"].copy()
     ds["lon"], _ = xr.broadcast(ds["YG"], ds["XG"])
@@ -80,7 +95,7 @@ def test_invalid_lon_lat():
         ValueError,
         match=".*have different dimensionalities\.",
     ):
-        XGrid(xgcm.Grid(ds, periodic=False))
+        XGrid.from_dataset(ds)
 
     ds = datasets["ds_2d_left"].copy()
     ds["lon"], ds["lat"] = xr.broadcast(ds["YG"], ds["XG"])
@@ -90,7 +105,7 @@ def test_invalid_lon_lat():
         ValueError,
         match=".*must be defined on the X and Y axes and transposed to have dimensions in order of Y, X\.",
     ):
-        XGrid(xgcm.Grid(ds, periodic=False))
+        XGrid.from_dataset(ds)
 
 
 @pytest.mark.parametrize(
@@ -101,7 +116,7 @@ def test_invalid_lon_lat():
     ],
 )  # for key, ds in datasets.items()])
 def test_xgrid_search_cpoints(ds):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     lat_array, lon_array = get_2d_fpoint_mesh(grid)
     lat_array, lon_array = corner_to_cell_center_points(lat_array, lon_array)
 
@@ -148,10 +163,10 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
 
 
 @pytest.mark.parametrize(
-    "grid, da_name, expected",
+    "ds, da_name, expected",
     [
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)),
+            datasets["ds_2d_left"],
             "U (C grid)",
             {
                 "XG": (np.int64(0), np.float64(0.0)),
@@ -161,7 +176,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
             id="MITgcm indexing style U (C grid)",
         ),
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)),
+            datasets["ds_2d_left"],
             "V (C grid)",
             {
                 "XC": (np.int64(-1), np.float64(0.5)),
@@ -171,7 +186,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
             id="MITgcm indexing style V (C grid)",
         ),
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)),
+            datasets["ds_2d_right"],
             "U (C grid)",
             {
                 "XG": (np.int64(0), np.float64(0.0)),
@@ -181,7 +196,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
             id="NEMO indexing style U (C grid)",
         ),
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)),
+            datasets["ds_2d_right"],
             "V (C grid)",
             {
                 "XC": (np.int64(0), np.float64(0.5)),
@@ -192,10 +207,11 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
         ),
     ],
 )
-def test_xgrid_localize_zero_position(grid, da_name, expected):
+def test_xgrid_localize_zero_position(ds, da_name, expected):
     """Test localize function using left and right datasets."""
+    grid = XGrid.from_dataset(ds)
+    da = ds[da_name]
     position = grid.search(0, 0, 0)
-    da = grid.xgcm_grid._ds[da_name]
 
     local_position = grid.localize(position, da.dims)
     assert local_position == expected, f"Expected {expected}, got {local_position}"