From 9f33ac9dd32195f82e5596f653a8f7b3d13fb643 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 17 Jul 2025 16:23:55 +0200
Subject: [PATCH 01/10] Update grid imports

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

diff --git a/tests/v4/test_datasets.py b/tests/v4/test_datasets.py
index e5201e58b7..d2d21040bb 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"]

From fcb2c1ec56d3e38b88b1c56ea8aa73c07d4d2c65 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 17 Jul 2025 16:44:10 +0200
Subject: [PATCH 02/10] Add XGrid from_dataset helper with default to drop
 Field data

It's not necessary for us to have the Field data on the XGrid/xgcm.Grid object itself. It doesn't make a big difference in the end, but dropping the field data on grid ingestion would allow us to keep things logically separate.
---
 parcels/xgrid.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 2cfda507e7..5c31f789bf 100644
--- a/parcels/xgrid.py
+++ b/parcels/xgrid.py
@@ -16,6 +16,7 @@
 _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]
+_DEFAULT_XGCM_KWARGS = {"periodic": False}
 
 
 def get_cell_count_along_dim(axis: xgcm.Axis) -> int:
@@ -34,6 +35,15 @@ 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)
+
+
 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 +63,18 @@ def __init__(self, grid: xgcm.Grid, mesh="flat"):
         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."""
+        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)

From 7908ab9d9554f309b347a3fc865239b677edf4f7 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 17 Jul 2025 17:15:46 +0200
Subject: [PATCH 03/10] Update test suite to use XGrid.from_dataset

---
 tests/v4/test_field.py               | 18 ++++++-------
 tests/v4/test_fieldset.py            | 19 +++++++-------
 tests/v4/test_index_search.py        |  3 +--
 tests/v4/test_kernel.py              |  3 +--
 tests/v4/test_particleset.py         |  3 +--
 tests/v4/test_particleset_execute.py |  3 +--
 tests/v4/test_xgrid.py               | 38 ++++++++++++++--------------
 7 files changed, 41 insertions(+), 46 deletions(-)

diff --git a/tests/v4/test_field.py b/tests/v4/test_field.py
index a555a672f2..6f1068896a 100644
--- a/tests/v4/test_field.py
+++ b/tests/v4/test_field.py
@@ -5,7 +5,7 @@
 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 @@
 
 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 invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, inval
 
 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 f623b7f1c4..c88e52577d 100644
--- a/tests/v4/test_fieldset.py
+++ b/tests/v4/test_fieldset.py
@@ -5,7 +5,6 @@
 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
@@ -21,7 +20,7 @@
 @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 +54,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 +67,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'"):
@@ -89,12 +88,12 @@ 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))
+    grid2 = XGrid.from_dataset(ds2)
     field2 = Field("field2", ds2["U (A grid)"], grid2, mesh_type="flat")
 
     fieldset = FieldSet([field1, field2])
@@ -116,14 +115,14 @@ def test_fieldset_init_incompatible_calendars():
     ds1 = ds.copy()
     ds1["time"] = xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True)
 
-    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))
+    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"):
@@ -133,7 +132,7 @@ 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))
+    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"):
@@ -141,7 +140,7 @@ def test_fieldset_add_field_incompatible_calendars(fieldset):
 
     ds_test = ds.copy()
     ds_test["time"] = np.linspace(0, 100, T_structured, dtype="timedelta64[s]")
-    grid = XGrid(xgcm.Grid(ds_test))
+    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 7f9290f12c..5d5c602b94 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 @@
 @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 1125d925cc..cb521bc787 100644
--- a/tests/v4/test_kernel.py
+++ b/tests/v4/test_kernel.py
@@ -6,7 +6,6 @@
     Field,
     FieldSet,
     ParticleSet,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.xgrid import XGrid
@@ -16,7 +15,7 @@
 @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 7a9dd7be64..ace5e8dba6 100644
--- a/tests/v4/test_particleset.py
+++ b/tests/v4/test_particleset.py
@@ -13,7 +13,6 @@
     ParticleSet,
     ParticleSetWarning,
     Variable,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.xgrid import XGrid
@@ -23,7 +22,7 @@
 @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 f63210add0..b1994d4d4e 100644
--- a/tests/v4/test_particleset_execute.py
+++ b/tests/v4/test_particleset_execute.py
@@ -10,7 +10,6 @@
     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 @@
 @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 210ca6b742..936f50c923 100644
--- a/tests/v4/test_xgrid.py
+++ b/tests/v4/test_xgrid.py
@@ -5,11 +5,10 @@
 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, _drop_field_data, _search_1d_array
 
-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 +33,25 @@ 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_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 +71,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 +79,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 +89,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 +100,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 +147,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 +160,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 +170,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 +180,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 +191,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}"

From dd3c07d8897458a0db1e2769808d59139fe1670f Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 17 Jul 2025 17:56:12 +0200
Subject: [PATCH 04/10] Transpose Field data to 4D arrays

Needed to update the tests because the time axis information was getting lost (due to xarray behaviour with keep_attrs).

Fixes #2047
---
 parcels/field.py          | 10 ++++-----
 parcels/fieldset.py       |  1 -
 parcels/xgrid.py          | 44 +++++++++++++++++++++++++++++++++++----
 tests/v4/test_fieldset.py | 26 ++++++++++++++++++-----
 4 files changed, 66 insertions(+), 15 deletions(-)

diff --git a/parcels/field.py b/parcels/field.py
index dba61dbd6a..f1549d0338 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -29,7 +29,7 @@
     _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 @@ def __init__(
 
         _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,9 +189,6 @@ def __init__(
         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. ")
-
     @property
     def units(self):
         return self._units
@@ -439,7 +439,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 c423000fc4..5937c7f573 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -132,7 +132,6 @@ def add_constant_field(self, name: str, value, mesh: Mesh = "flat"):
         """
         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 5c31f789bf..4cd205b48d 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,12 +10,15 @@
 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}
 
 
@@ -35,7 +38,7 @@ 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:
+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
@@ -44,6 +47,39 @@ def drop_field_data(ds: xr.Dataset) -> xr.Dataset:
     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 dummy 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"dummy{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)
+
+    dummy_dims_to_create = {}
+    for ax in axes_not_in_field:
+        dummy_dims_to_create[f"dummy{ax}"] = 1
+        ax_dims.append((ax, f"dummy{ax}"))
+
+    if dummy_dims_to_create:
+        da = da.expand_dims(dummy_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).
@@ -71,7 +107,7 @@ def from_dataset(cls, ds: xr.Dataset, mesh="flat", xgcm_kwargs=None):
 
         xgcm_kwargs = {**_DEFAULT_XGCM_KWARGS, **xgcm_kwargs}
 
-        ds = drop_field_data(ds)
+        ds = _drop_field_data(ds)
         grid = xgcm.Grid(ds, **xgcm_kwargs)
         return cls(grid, mesh=mesh)
 
diff --git a/tests/v4/test_fieldset.py b/tests/v4/test_fieldset.py
index c88e52577d..fb9a10a6a6 100644
--- a/tests/v4/test_fieldset.py
+++ b/tests/v4/test_fieldset.py
@@ -92,7 +92,7 @@ def test_fieldset_time_interval():
     field1 = Field("field1", ds["U (A grid)"], grid1, mesh_type="flat")
 
     ds2 = ds.copy()
-    ds2["time"] = ds2["time"] + np.timedelta64(timedelta(days=1))
+    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")
 
@@ -113,7 +113,11 @@ 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.from_dataset(ds1)
     U = Field("U", ds1["U (A grid)"], grid, mesh_type="flat")
@@ -121,7 +125,11 @@ def test_fieldset_init_incompatible_calendars():
     UV = VectorField("UV", U, V)
 
     ds2 = ds.copy()
-    ds2["time"] = xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True)
+    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")
 
@@ -131,7 +139,11 @@ 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)
+    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")
 
@@ -139,7 +151,11 @@ def test_fieldset_add_field_incompatible_calendars(fieldset):
         fieldset.add_field(field, "test_field")
 
     ds_test = ds.copy()
-    ds_test["time"] = np.linspace(0, 100, T_structured, dtype="timedelta64[s]")
+    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")
 

From 1089dd2eec1516acd780837fb2eb5f2eb694846f Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 17 Jul 2025 18:45:51 +0200
Subject: [PATCH 05/10] Add test test_fieldset_from_structured_generic_datasets

---
 tests/v4/test_fieldset.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/tests/v4/test_fieldset.py b/tests/v4/test_fieldset.py
index fb9a10a6a6..e67e41b423 100644
--- a/tests/v4/test_fieldset.py
+++ b/tests/v4/test_fieldset.py
@@ -12,7 +12,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 parcels.xgrid import _FIELD_DATA_ORDERING, XGrid, get_axis_from_dim_name
 
 ds = datasets_structured["ds_2d_left"]
 
@@ -84,6 +84,30 @@ 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():
+        assert (
+            len(field.data.shape) == 4
+        ), f"Field data should have 4 dimensions (time, depth, lat, lon), got {len(field.data.shape)}"
+
+        for ax_expected, dim in zip(_FIELD_DATA_ORDERING, field.data.dims, strict=True):
+            ax_actual = get_axis_from_dim_name(field.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}"
+
+    assert len(fieldset.gridset) == 1
+
+
 def test_fieldset_gridset_multiple_grids(): ...
 
 

From ff95ea6c7bdd9c4a14ae42625f626855f2d6294e Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Thu, 17 Jul 2025 19:24:39 +0200
Subject: [PATCH 06/10] Add unit test test_transpose_xfield_data_to_tzyx

---
 tests/utils.py            | 17 +++++++++++++++++
 tests/v4/test_fieldset.py | 13 +++----------
 tests/v4/test_xgrid.py    | 18 +++++++++++++++++-
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/tests/utils.py b/tests/utils.py
index e0fd5f34fe..ca8157f6ee 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 {len(data.shape)}"
+
+    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_fieldset.py b/tests/v4/test_fieldset.py
index e67e41b423..c3871018be 100644
--- a/tests/v4/test_fieldset.py
+++ b/tests/v4/test_fieldset.py
@@ -12,7 +12,8 @@
 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 _FIELD_DATA_ORDERING, XGrid, get_axis_from_dim_name
+from parcels.xgrid import XGrid
+from tests import utils
 
 ds = datasets_structured["ds_2d_left"]
 
@@ -95,15 +96,7 @@ def test_fieldset_from_structured_generic_datasets(ds):
 
     assert len(fieldset.fields) == len(ds.data_vars)
     for field in fieldset.fields.values():
-        assert (
-            len(field.data.shape) == 4
-        ), f"Field data should have 4 dimensions (time, depth, lat, lon), got {len(field.data.shape)}"
-
-        for ax_expected, dim in zip(_FIELD_DATA_ORDERING, field.data.dims, strict=True):
-            ax_actual = get_axis_from_dim_name(field.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}"
+        utils.assert_valid_field_data(field.data, field.grid)
 
     assert len(fieldset.gridset) == 1
 
diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 936f50c923..7e0b7faad8 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
@@ -6,7 +7,8 @@
 from numpy.testing import assert_allclose
 
 from parcels._datasets.structured.generic import X, Y, Z, datasets
-from parcels.xgrid import XGrid, _drop_field_data, _search_1d_array
+from parcels.xgrid import XGrid, _drop_field_data, _search_1d_array, _transpose_xfield_data_to_tzyx
+from tests import utils
 
 GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"])
 
@@ -49,6 +51,20 @@ def test_xgrid_axes(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.from_dataset(ds)

From 6ba1c94d851f7f7e5d9f95f25ebc317448f8ced6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 17 Jul 2025 17:34:20 +0000
Subject: [PATCH 07/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/v4/test_xgrid.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 7e0b7faad8..30bcc7c1aa 100644
--- a/tests/v4/test_xgrid.py
+++ b/tests/v4/test_xgrid.py
@@ -7,7 +7,7 @@
 from numpy.testing import assert_allclose
 
 from parcels._datasets.structured.generic import X, Y, Z, datasets
-from parcels.xgrid import XGrid, _drop_field_data, _search_1d_array, _transpose_xfield_data_to_tzyx
+from parcels.xgrid import XGrid, _search_1d_array, _transpose_xfield_data_to_tzyx
 from tests import utils
 
 GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"])

From 422ef97b02bbf3ba6616599c3741ce3fc5f00927 Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Fri, 18 Jul 2025 10:43:22 +0200
Subject: [PATCH 08/10] Update naming to 'mock' axes

---
 parcels/xgrid.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 4cd205b48d..4135cd0af5 100644
--- a/parcels/xgrid.py
+++ b/parcels/xgrid.py
@@ -50,14 +50,14 @@ def _drop_field_data(ds: xr.Dataset) -> xr.Dataset:
 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 dummy dimensions of size 1 for any axes not present in the DataArray.
+    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"dummy{axis}" for old_dim, axis in zip(da.dims, _FIELD_DATA_ORDERING, strict=True)})
+        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):
@@ -67,13 +67,13 @@ def _transpose_xfield_data_to_tzyx(da: xr.DataArray, xgcm_grid: xgcm.Grid) -> xr
 
     axes_not_in_field = set(_FIELD_DATA_ORDERING) - set(ax_dim[0] for ax_dim in ax_dims)
 
-    dummy_dims_to_create = {}
+    mock_dims_to_create = {}
     for ax in axes_not_in_field:
-        dummy_dims_to_create[f"dummy{ax}"] = 1
-        ax_dims.append((ax, f"dummy{ax}"))
+        mock_dims_to_create[f"mock{ax}"] = 1
+        ax_dims.append((ax, f"mock{ax}"))
 
-    if dummy_dims_to_create:
-        da = da.expand_dims(dummy_dims_to_create, create_index_for_new_dim=False)
+    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]))
 

From 2820e51753a806757362b50347ad407f3ea8ab9d Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Fri, 18 Jul 2025 10:55:12 +0200
Subject: [PATCH 09/10] Review feedback

---
 parcels/xgrid.py | 2 +-
 tests/utils.py   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index 4135cd0af5..81391ee818 100644
--- a/parcels/xgrid.py
+++ b/parcels/xgrid.py
@@ -101,7 +101,7 @@ def __init__(self, grid: xgcm.Grid, mesh="flat"):
 
     @classmethod
     def from_dataset(cls, ds: xr.Dataset, mesh="flat", xgcm_kwargs=None):
-        """WARNING: unstable API, subject to change in future versions."""
+        """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 = {}
 
diff --git a/tests/utils.py b/tests/utils.py
index ca8157f6ee..a76b7e8a6e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -126,7 +126,7 @@ def assert_empty_folder(path: Path):
 
 
 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 {len(data.shape)}"
+    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)

From ddb4b4a3ee012d7a098f9c0777fd882c6240f7bd Mon Sep 17 00:00:00 2001
From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com>
Date: Fri, 18 Jul 2025 14:04:04 +0200
Subject: [PATCH 10/10] Add time coordinate validation in Field init

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

diff --git a/parcels/field.py b/parcels/field.py
index f1549d0338..1a1fbb7175 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -189,6 +189,10 @@ def __init__(
         else:
             raise ValueError("Unsupported mesh type in data array attributes. Choose either: 'spherical' or 'flat'")
 
+        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):
         return self._units