Merge branch 'v4-dev' into particledata_as_dict

Author: erikvansebille

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,8 +189,9 @@ 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. ")
+        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])

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,13 +10,17 @@
 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 @@ 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."""  # 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)

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

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

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