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

Author: VeckoTheGecko

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])

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(

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)
 

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,15 +113,23 @@ 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")
     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)
+    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,15 +139,23 @@ 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")
 
     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]")
+    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")