Update Field init to take XGrid

Author: VeckoTheGecko

parcels/field.py

@@ -30,7 +30,6 @@
     _raise_field_out_of_bound_error,
 )
 from parcels.uxgrid import UxGrid
-from parcels.xgcm.grid import Grid
 from parcels.xgrid import XGrid
 
 from ._index_search import _search_time_index
@@ -145,7 +144,7 @@ def __init__(
         self,
         name: str,
         data: xr.DataArray | ux.UxDataArray,
-        grid: UxGrid | Grid,
+        grid: UxGrid | XGrid,
         mesh_type: Mesh = "flat",
         interp_method: Callable | None = None,
     ):
@@ -155,7 +154,7 @@ def __init__(
             )
         if not isinstance(name, str):
             raise ValueError(f"Expected `name` to be a string, got {type(name)}.")
-        if not isinstance(grid, (UxGrid, Grid)):
+        if not isinstance(grid, (UxGrid, XGrid)):
             raise ValueError(f"Expected `grid` to be a parcels UxGrid, or parcels Grid object, got {type(grid)}.")
 
         assert_valid_mesh(mesh_type)
@@ -174,13 +173,6 @@ def __init__(
             )
             raise e
 
-        # For compatibility with parts of the codebase that rely on v3 definition of Grid.
-        # Should be worked to be removed in v4
-        if isinstance(grid, Grid):
-            self.gridadapter = XGrid(grid)
-        else:
-            self.gridadapter = None
-
         try:
             if isinstance(data, ux.UxDataArray):
                 _assert_valid_uxdataarray(data)
@@ -232,7 +224,7 @@ def lat(self):
             elif self.data.attrs["location"] == "edge":
                 return self.grid.edge_lat
         else:
-            return self.gridadapter.lat
+            return self.grid.lat
 
     @property
     def lon(self):
@@ -244,7 +236,7 @@ def lon(self):
             elif self.data.attrs["location"] == "edge":
                 return self.grid.edge_lon
         else:
-            return self.gridadapter.lon
+            return self.grid.lon
 
     @property
     def depth(self):
@@ -255,26 +247,26 @@ def depth(self):
             elif vertical_location == "face":
                 return self.grid.nz
         else:
-            return self.gridadapter.depth
+            return self.grid.depth
 
     @property
     def xdim(self):
         if type(self.data) is xr.DataArray:
-            return self.gridadapter.xdim
+            return self.grid.xdim
         else:
             raise NotImplementedError("xdim not implemented for unstructured grids")
 
     @property
     def ydim(self):
         if type(self.data) is xr.DataArray:
-            return self.gridadapter.ydim
+            return self.grid.ydim
         else:
             raise NotImplementedError("ydim not implemented for unstructured grids")
 
     @property
     def zdim(self):
         if type(self.data) is xr.DataArray:
-            return self.gridadapter.zdim
+            return self.grid.zdim
         else:
             if "nz1" in self.data.dims:
                 return self.data.sizes["nz1"]
@@ -523,14 +515,14 @@ def _assert_valid_uxgrid(grid):
         )
 
 
-def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux.Grid | Grid):
+def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux.Grid | XGrid):
     if isinstance(data, ux.UxDataArray):
         if not isinstance(grid, UxGrid):
             raise ValueError(
                 f"Incompatible data-grid combination. Data is a uxarray.UxDataArray, expected `grid` to be a UxGrid object, got {type(grid)}."
             )
     elif isinstance(data, xr.DataArray):
-        if not isinstance(grid, Grid):
+        if not isinstance(grid, XGrid):
             raise ValueError(
                 f"Incompatible data-grid combination. Data is a xarray.DataArray, expected `grid` to be a parcels Grid object, got {type(grid)}."
             )

parcels/fieldset.py

@@ -7,11 +7,12 @@
 import numpy as np
 import xarray as xr
 
+from parcels import xgcm
 from parcels._core.utils.time import get_datetime_type_calendar
 from parcels._core.utils.time import is_compatible as datetime_is_compatible
 from parcels._typing import Mesh
 from parcels.field import Field, VectorField
-from parcels.xgcm.grid import Grid
+from parcels.xgrid import XGrid
 
 if TYPE_CHECKING:
     from parcels._typing import DatetimeLike
@@ -171,7 +172,7 @@ def add_constant_field(self, name: str, value, mesh: Mesh = "flat"):
                 "time": (["time"], np.arange(1), {"axis": "T"}),
             },
         )
-        grid = Grid(da)
+        grid = XGrid(xgcm.Grid(da))
         self.add_field(
             Field(
                 name,

tests/v4/test_field.py

@@ -3,39 +3,34 @@
 import uxarray as ux
 import xarray as xr
 
-from parcels import Field
+from parcels import Field, xgcm
 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
 from parcels.uxgrid import UxGrid
-from parcels.xgcm import Grid
+from parcels.xgrid import XGrid
 
 
 def test_field_init_param_types():
-    data = xr.DataArray(
-        attrs={
-            "location": "node",
-            "mesh": "flat",
-        }
-    )
-    grid = Grid(data)
+    data = datasets_structured["ds_2d_left"]
+    grid = XGrid(xgcm.Grid(data))
     with pytest.raises(ValueError, match="Expected `name` to be a string"):
-        Field(name=123, data=data, grid=grid)
+        Field(name=123, data=data["data_g"], grid=grid)
 
     with pytest.raises(ValueError, match="Expected `data` to be a uxarray.UxDataArray or xarray.DataArray"):
         Field(name="test", data=123, grid=grid)
 
-    with pytest.raises(ValueError, match="Expected `grid` to be a parcels UxGrid, or parcels Grid"):
-        Field(name="test", data=data, grid=123)
+    with pytest.raises(ValueError, match="Expected `grid` to be a parcels UxGrid, or parcels XGrid"):
+        Field(name="test", data=data["data_g"], grid=123)
 
     with pytest.raises(ValueError, match="Invalid value 'invalid'. Valid options are.*"):
-        Field(name="test", data=data, grid=grid, mesh_type="invalid")
+        Field(name="test", data=data["data_g"], grid=grid, mesh_type="invalid")
 
 
 @pytest.mark.parametrize(
     "data,grid",
     [
-        pytest.param(ux.UxDataArray(), Grid(xr.Dataset()), id="uxdata-grid"),
+        pytest.param(ux.UxDataArray(), XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])), id="uxdata-grid"),
         pytest.param(
             xr.DataArray(),
             UxGrid(datasets_unstructured["stommel_gyre_delaunay"].uxgrid),
@@ -56,7 +51,9 @@ def test_field_incompatible_combination(data, grid):
     "data,grid",
     [
         pytest.param(
-            datasets_structured["ds_2d_left"]["data_g"], Grid(datasets_structured["ds_2d_left"]), id="ds_2d_left"
+            datasets_structured["ds_2d_left"]["data_g"],
+            XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])),
+            id="ds_2d_left",
         ),  # TODO: Perhaps this test should be expanded to cover more datasets?
     ],
 )
@@ -79,7 +76,7 @@ def test_field_init_fail_on_bad_time_type(numpy_dtype):
     ds["time"] = np.arange(0, T_structured, dtype=numpy_dtype)
 
     data = ds["data_g"]
-    grid = Grid(ds)
+    grid = XGrid(xgcm.Grid(ds))
     with pytest.raises(
         ValueError,
         match="Error getting time interval.*. Are you sure that the time dimension on the xarray dataset is stored as datetime or cftime datetime objects\?",
@@ -95,7 +92,9 @@ def test_field_init_fail_on_bad_time_type(numpy_dtype):
     "data,grid",
     [
         pytest.param(
-            datasets_structured["ds_2d_left"]["data_g"], Grid(datasets_structured["ds_2d_left"]), id="ds_2d_left"
+            datasets_structured["ds_2d_left"]["data_g"],
+            XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])),
+            id="ds_2d_left",
         ),
     ],
 )

tests/v4/test_fieldset.py

@@ -5,19 +5,20 @@
 import pytest
 import xarray as xr
 
+from parcels import xgcm
 from parcels._datasets.structured.generic import T as T_structured
 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.xgcm import Grid
+from parcels.xgrid import XGrid
 
 ds = datasets_structured["ds_2d_left"]
 
 
 @pytest.fixture
 def fieldset() -> FieldSet:
     """Fixture to create a FieldSet object for testing."""
-    grid = Grid(ds)
+    grid = XGrid(xgcm.Grid(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)
@@ -51,7 +52,7 @@ def test_fieldset_add_constant_field(fieldset):
 
 
 def test_fieldset_add_field(fieldset):
-    grid = Grid(ds)
+    grid = XGrid(xgcm.Grid(ds))
     field = Field("test_field", ds["U (A grid)"], grid, mesh_type="flat")
     fieldset.add_field(field)
     assert fieldset.test_field == field
@@ -64,7 +65,7 @@ def test_fieldset_add_field_wrong_type(fieldset):
 
 
 def test_fieldset_add_field_already_exists(fieldset):
-    grid = Grid(ds)
+    grid = XGrid(xgcm.Grid(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'"):
@@ -77,12 +78,12 @@ def test_fieldset_gridset_size(fieldset):
 
 
 def test_fieldset_time_interval():
-    grid1 = Grid(ds)
+    grid1 = XGrid(xgcm.Grid(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 = Grid(ds2)
+    grid2 = XGrid(xgcm.Grid(ds2))
     field2 = Field("field2", ds2["U (A grid)"], grid2, mesh_type="flat")
 
     fieldset = FieldSet([field1, field2])
@@ -96,14 +97,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 = Grid(ds1)
+    grid = XGrid(xgcm.Grid(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 = Grid(ds2)
+    grid2 = XGrid(xgcm.Grid(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"):
@@ -113,7 +114,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 = Grid(ds_test)
+    grid = XGrid(xgcm.Grid(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"):