Move axis assertion to dedicated function

Author: VeckoTheGecko

src/parcels/_core/field.py

@@ -21,7 +21,7 @@
 from parcels._core.utils.string import _assert_str_and_python_varname
 from parcels._core.utils.time import TimeInterval
 from parcels._core.uxgrid import UxGrid
-from parcels._core.xgrid import XGrid, _transpose_xfield_data_to_tzyx
+from parcels._core.xgrid import XGrid, _transpose_xfield_data_to_tzyx, assert_all_field_dims_have_axis
 from parcels._python import assert_same_function_signature
 from parcels._reprs import default_repr
 from parcels._typing import VectorType
@@ -111,6 +111,7 @@ def __init__(
         _assert_compatible_combination(data, grid)
 
         if isinstance(grid, XGrid):
+            assert_all_field_dims_have_axis(data, grid.xgcm_grid)
             data = _transpose_xfield_data_to_tzyx(data, grid.xgcm_grid)
 
         self.name = name

src/parcels/_core/xgrid.py

@@ -48,20 +48,23 @@ 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 mock dimensions of size 1 for any axes not present in the DataArray.
-    """
+def assert_all_field_dims_have_axis(da: xr.DataArray, xgcm_grid: xgcm.Grid) -> None:
     ax_dims = [(get_axis_from_dim_name(xgcm_grid.axes, dim), dim) for dim in da.dims]
-
     for dim in ax_dims:
         if dim[0] is None:
             raise ValueError(
                 f'Dimension "{dim[1]}" has no axis attribute. '
                 f'HINT: You may want to add an {{"axis": A}} to your DataSet["{dim[1]}"], where A is one of "X", "Y", "Z" or "T"'
             )
 
+
+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)