Update test suite to use XGrid.from_dataset

Author: VeckoTheGecko

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

tests/v4/test_fieldset.py

@@ -5,7 +5,6 @@
 import pytest
 import xarray as xr
 
-from parcels import xgcm
 from parcels._datasets.structured.circulation_models import (
     datasets as datasets_circulation_models,  # noqa: F401
 )  # just making sure the import works. Will eventually be used in tests
@@ -21,7 +20,7 @@
 @pytest.fixture
 def fieldset() -> FieldSet:
     """Fixture to create a FieldSet object for testing."""
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(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)
@@ -55,7 +54,7 @@ def test_fieldset_add_constant_field(fieldset):
 
 
 def test_fieldset_add_field(fieldset):
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     field = Field("test_field", ds["U (A grid)"], grid, mesh_type="flat")
     fieldset.add_field(field)
     assert fieldset.test_field == field
@@ -68,7 +67,7 @@ def test_fieldset_add_field_wrong_type(fieldset):
 
 
 def test_fieldset_add_field_already_exists(fieldset):
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(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'"):
@@ -89,12 +88,12 @@ def test_fieldset_gridset_multiple_grids(): ...
 
 
 def test_fieldset_time_interval():
-    grid1 = XGrid(xgcm.Grid(ds))
+    grid1 = XGrid.from_dataset(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 = XGrid(xgcm.Grid(ds2))
+    grid2 = XGrid.from_dataset(ds2)
     field2 = Field("field2", ds2["U (A grid)"], grid2, mesh_type="flat")
 
     fieldset = FieldSet([field1, field2])
@@ -116,14 +115,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 = XGrid(xgcm.Grid(ds1))
+    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)
-    grid2 = XGrid(xgcm.Grid(ds2))
+    grid2 = XGrid.from_dataset(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"):
@@ -133,15 +132,15 @@ 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 = XGrid(xgcm.Grid(ds_test))
+    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]")
-    grid = XGrid(xgcm.Grid(ds_test))
+    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"):

tests/v4/test_index_search.py

@@ -1,7 +1,6 @@
 import numpy as np
 import pytest
 
-from parcels import xgcm
 from parcels._datasets.structured.generic import datasets
 from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.field import Field
@@ -13,7 +12,7 @@
 @pytest.fixture
 def field_cone():
     ds = datasets["2d_left_unrolled_cone"]
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     field = Field(
         name="test_field",
         data=ds["data_g"],

tests/v4/test_kernel.py

@@ -6,7 +6,6 @@
     Field,
     FieldSet,
     ParticleSet,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.xgrid import XGrid
@@ -16,7 +15,7 @@
 @pytest.fixture
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     return FieldSet([U, V])

tests/v4/test_particleset.py

@@ -13,7 +13,6 @@
     ParticleSet,
     ParticleSetWarning,
     Variable,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels.xgrid import XGrid
@@ -23,7 +22,7 @@
 @pytest.fixture
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     return FieldSet([U, V])

tests/v4/test_particleset_execute.py

@@ -10,7 +10,6 @@
     StatusCode,
     UXPiecewiseConstantFace,
     VectorField,
-    xgcm,
 )
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
@@ -22,7 +21,7 @@
 @pytest.fixture
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid(xgcm.Grid(ds))
+    grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
     V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
     return FieldSet([U, V])

tests/v4/test_xgrid.py

@@ -5,11 +5,10 @@
 import xarray as xr
 from numpy.testing import assert_allclose
 
-from parcels import xgcm
 from parcels._datasets.structured.generic import X, Y, Z, datasets
-from parcels.xgrid import XGrid, _search_1d_array
+from parcels.xgrid import XGrid, _drop_field_data, _search_1d_array
 
-GridTestCase = namedtuple("GridTestCase", ["Grid", "attr", "expected"])
+GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"])
 
 test_cases = [
     GridTestCase(datasets["ds_2d_left"], "lon", datasets["ds_2d_left"].XG.values),
@@ -34,25 +33,25 @@ def assert_equal(actual, expected):
 
 @pytest.mark.parametrize("ds, attr, expected", test_cases)
 def test_xgrid_properties_ground_truth(ds, attr, expected):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     actual = getattr(grid, attr)
     assert_equal(actual, expected)
 
 
 @pytest.mark.parametrize("ds", [pytest.param(ds, id=key) for key, ds in datasets.items()])
-def test_xgrid_init_on_generic_datasets(ds):
-    XGrid(xgcm.Grid(ds, periodic=False))
+def test_xgrid_from_dataset_on_generic_datasets(ds):
+    XGrid.from_dataset(ds)
 
 
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_xgrid_axes(ds):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     assert grid.axes == ["Z", "Y", "X"]
 
 
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_xgrid_get_axis_dim(ds):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     assert grid.get_axis_dim("Z") == Z - 1
     assert grid.get_axis_dim("Y") == Y - 1
     assert grid.get_axis_dim("X") == X - 1
@@ -72,15 +71,15 @@ def test_invalid_lon_lat():
         ValueError,
         match=".*is defined on the center of the grid, but must be defined on the F points\.",
     ):
-        XGrid(xgcm.Grid(ds, periodic=False))
+        XGrid.from_dataset(ds)
 
     ds = datasets["ds_2d_left"].copy()
     ds["lon"], _ = xr.broadcast(ds["YG"], ds["XG"])
     with pytest.raises(
         ValueError,
         match=".*have different dimensionalities\.",
     ):
-        XGrid(xgcm.Grid(ds, periodic=False))
+        XGrid.from_dataset(ds)
 
     ds = datasets["ds_2d_left"].copy()
     ds["lon"], ds["lat"] = xr.broadcast(ds["YG"], ds["XG"])
@@ -90,7 +89,7 @@ def test_invalid_lon_lat():
         ValueError,
         match=".*must be defined on the X and Y axes and transposed to have dimensions in order of Y, X\.",
     ):
-        XGrid(xgcm.Grid(ds, periodic=False))
+        XGrid.from_dataset(ds)
 
 
 @pytest.mark.parametrize(
@@ -101,7 +100,7 @@ def test_invalid_lon_lat():
     ],
 )  # for key, ds in datasets.items()])
 def test_xgrid_search_cpoints(ds):
-    grid = XGrid(xgcm.Grid(ds, periodic=False))
+    grid = XGrid.from_dataset(ds)
     lat_array, lon_array = get_2d_fpoint_mesh(grid)
     lat_array, lon_array = corner_to_cell_center_points(lat_array, lon_array)
 
@@ -148,10 +147,10 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
 
 
 @pytest.mark.parametrize(
-    "grid, da_name, expected",
+    "ds, da_name, expected",
     [
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)),
+            datasets["ds_2d_left"],
             "U (C grid)",
             {
                 "XG": (np.int64(0), np.float64(0.0)),
@@ -161,7 +160,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
             id="MITgcm indexing style U (C grid)",
         ),
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)),
+            datasets["ds_2d_left"],
             "V (C grid)",
             {
                 "XC": (np.int64(-1), np.float64(0.5)),
@@ -171,7 +170,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
             id="MITgcm indexing style V (C grid)",
         ),
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)),
+            datasets["ds_2d_right"],
             "U (C grid)",
             {
                 "XG": (np.int64(0), np.float64(0.0)),
@@ -181,7 +180,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
             id="NEMO indexing style U (C grid)",
         ),
         pytest.param(
-            XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)),
+            datasets["ds_2d_right"],
             "V (C grid)",
             {
                 "XC": (np.int64(0), np.float64(0.5)),
@@ -192,10 +191,11 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
         ),
     ],
 )
-def test_xgrid_localize_zero_position(grid, da_name, expected):
+def test_xgrid_localize_zero_position(ds, da_name, expected):
     """Test localize function using left and right datasets."""
+    grid = XGrid.from_dataset(ds)
+    da = ds[da_name]
     position = grid.search(0, 0, 0)
-    da = grid.xgcm_grid._ds[da_name]
 
     local_position = grid.localize(position, da.dims)
     assert local_position == expected, f"Expected {expected}, got {local_position}"