Add unit test test_transpose_xfield_data_to_tzyx

Author: VeckoTheGecko

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}"

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
 

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)