Cleanup tests/test_advection.py fieldset ingestion to use SGRID

docs/development/posting-issues.md

@@ -46,8 +46,8 @@ As a user with access to your dataset, you would do:
 
 # Generate an example dataset to zip. The user would use their own.
 import xarray as xr
-from parcels._datasets.structured.generic import datasets
-datasets['ds_2d_left'].to_netcdf("my_dataset.nc")
+from parcels._datasets.structured.generic import datasets_comodo
+datasets_comodo['ds_2d_left'].to_netcdf("my_dataset.nc")
 ```
 
 ```{code-cell}

src/parcels/_datasets/structured/generated.py

@@ -75,6 +75,19 @@ def radial_rotation_dataset(xdim=200, ydim=200):  # Define 2D flat, square field
             "lat": (["YG"], lat, {"axis": "Y", "c_grid_axis_shift": 0.5}),
             "lon": (["XG"], lon, {"axis": "X", "c_grid_axis_shift": -0.5}),
         },
+    ).pipe(
+        sgrid._attach_sgrid_metadata,
+        sgrid.SGrid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("XG", "YG"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
+                sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
+            ),
+            vertical_dimensions=(sgrid.FaceNodePadding("ZC", "depth", sgrid.Padding.BOTH),),
+        ),
     )
 
 
@@ -111,6 +124,19 @@ def moving_eddy_dataset(xdim=2, ydim=2):  # TODO check if this also works with x
             "u_g": u_g,
             "f": f,
         },
+    ).pipe(
+        sgrid._attach_sgrid_metadata,
+        sgrid.SGrid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("XG", "YG"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
+                sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
+            ),
+            vertical_dimensions=(sgrid.FaceNodePadding("ZC", "depth", sgrid.Padding.BOTH),),
+        ),
     )
 
 
@@ -161,6 +187,19 @@ def decaying_moving_eddy_dataset(xdim=2, ydim=2):
             "gamma": gamma,
             "gamma_g": gamma_g,
         },
+    ).pipe(
+        sgrid._attach_sgrid_metadata,
+        sgrid.SGrid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("XG", "YG"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
+                sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
+            ),
+            vertical_dimensions=(sgrid.FaceNodePadding("ZC", "depth", sgrid.Padding.BOTH),),
+        ),
     )
 
 
@@ -244,6 +283,18 @@ def peninsula_dataset(xdim=100, ydim=50, mesh="flat", grid_type="A"):
             "lat": (["YG"], lat, {"axis": "Y", "c_grid_axis_shift": 0.5}),
             "lon": (["XG"], lon, {"axis": "X", "c_grid_axis_shift": -0.5}),
         },
+    ).pipe(
+        sgrid._attach_sgrid_metadata,
+        sgrid.SGrid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("XG", "YG"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
+                sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
+            ),
+        ),
     )
 
 
@@ -300,4 +351,16 @@ def stommel_gyre_dataset(xdim=200, ydim=200, grid_type="A"):
             "lat": (["YG"], lat, {"axis": "Y", "c_grid_axis_shift": 0.5}),
             "lon": (["XG"], lon, {"axis": "X", "c_grid_axis_shift": -0.5}),
         },
+    ).pipe(
+        sgrid._attach_sgrid_metadata,
+        sgrid.SGrid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("XG", "YG"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW),
+                sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
+            ),
+        ),
     )

src/parcels/_datasets/structured/generic.py

@@ -5,7 +5,7 @@
 
 from . import T, X, Y, Z
 
-__all__ = ["T", "X", "Y", "Z", "datasets"]
+__all__ = ["T", "X", "Y", "Z", "datasets_comodo", "datasets_sgrid"]
 
 TIME = xr.date_range("2000", "2001", T)
 
@@ -139,7 +139,7 @@ def _unrolled_cone_curvilinear_grid():
     )
 
 
-datasets = {
+datasets_comodo = {
     "2d_left_rotated": _rotated_curvilinear_grid(),
     "ds_2d_left": xr.Dataset(  # MITgcm indexing style
         {
@@ -238,7 +238,7 @@ def _unrolled_cone_curvilinear_grid():
 }
 datasets_sgrid = {
     "ds_2d_padded_high": (
-        datasets["ds_2d_left"]
+        datasets_comodo["ds_2d_left"]
         .pipe(
             sgrid._attach_sgrid_metadata,
             sgrid.SGrid2DMetadata(
@@ -258,7 +258,7 @@ def _unrolled_cone_curvilinear_grid():
         )
     ),
     "ds_2d_padded_low": (
-        datasets["ds_2d_right"]
+        datasets_comodo["ds_2d_right"]
         .pipe(
             sgrid._attach_sgrid_metadata,
             sgrid.SGrid2DMetadata(

tests/datasets/test_structured.py

@@ -1,12 +1,12 @@
 import xgcm
 
 from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS
-from parcels._datasets.structured.generic import datasets
+from parcels._datasets.structured.generic import datasets_comodo
 
 
 def test_left_indexed_dataset():
     """Checks that 'ds_2d_left' is right indexed on all variables."""
-    ds = datasets["ds_2d_left"]
+    ds = datasets_comodo["ds_2d_left"]
     grid = xgcm.Grid(ds, **_DEFAULT_XGCM_KWARGS)
 
     for _axis_name, axis in grid.axes.items():
@@ -16,7 +16,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"]
+    ds = datasets_comodo["ds_2d_right"]
     grid = xgcm.Grid(ds, **_DEFAULT_XGCM_KWARGS)
     for _axis_name, axis in grid.axes.items():
         for pos, _dim_name in axis.coords.items():

tests/datasets/test_utils.py

@@ -3,7 +3,7 @@
 import xarray as xr
 
 from parcels._datasets import utils
-from parcels._datasets.structured.generic import datasets
+from parcels._datasets.structured.generic import datasets_comodo
 
 
 @pytest.fixture
@@ -28,7 +28,7 @@ def nonzero_ds():
     )
 
 
-@pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets.items()])
+@pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets_comodo.items()])
 @pytest.mark.parametrize("except_for", [None, "coords"])
 def test_replace_arrays_with_zeros(ds, except_for):
     # make sure doesn't error with range of datasets

tests/test_advection.py

@@ -6,15 +6,12 @@
 import parcels
 import parcels.tutorial
 from parcels import (
-    Field,
     FieldSet,
     Particle,
     ParticleFile,
     ParticleSet,
     StatusCode,
     Variable,
-    VectorField,
-    XGrid,
     convert,
 )
 from parcels._core.utils.time import timedelta_to_float
@@ -26,7 +23,7 @@
     simple_UV_dataset,
     stommel_gyre_dataset,
 )
-from parcels.interpolators import CGrid_Velocity, XLinear, XLinear_Velocity
+from parcels._datasets.structured.generic import datasets_sgrid
 from parcels.kernels import (
     AdvectionDiffusionEM,
     AdvectionDiffusionM1,
@@ -188,70 +185,49 @@ def SubmergeParticle(particles, fieldset):  # pragma: no cover
         assert len(pset) == 0
 
 
-@pytest.mark.parametrize("u", [-0.3, np.array(0.2)])
-@pytest.mark.parametrize("v", [0.2, np.array(1)])
-@pytest.mark.parametrize("w", [None, -0.2, np.array(0.7)])
-def test_length1dimensions(u, v, w):  # TODO: Refactor this test to be more readable (and isolate test setup)
-    (lon, xdim) = (np.linspace(-10, 10, 21), 21) if isinstance(u, np.ndarray) else (np.array([0]), 1)
-    (lat, ydim) = (np.linspace(-15, 15, 31), 31) if isinstance(v, np.ndarray) else (np.array([-4]), 1)
-    (depth, zdim) = (
-        (np.linspace(-5, 5, 11), 11) if (isinstance(w, np.ndarray) and w is not None) else (np.array([3]), 1)
-    )
+@pytest.mark.parametrize(
+    "u_value, x_slice", [(-0.03, slice(0, 1)), (0.02, slice(None))], ids=["single_u_layer", "full_u"]
+)
+@pytest.mark.parametrize(
+    "v_value, y_slice", [(0.02, slice(0, 1)), (0.1, slice(None))], ids=["single_v_layer", "full_v"]
+)
+@pytest.mark.parametrize(
+    "w_value, z_slice",
+    [(None, None), (-0.02, slice(0, 1)), (0.07, slice(None))],
+    ids=["no_vertical", "single_w_layer", "full_w"],
+)
+def test_length1dimensions(u_value, x_slice, v_value, y_slice, w_value, z_slice):
+    ds = datasets_sgrid["ds_2d_padded_high"].copy()[["U_A_grid", "grid"]]
+    ds = ds.isel(time=slice(2))  # TODO make this also work for length-1 time dimensions
+    ds = ds.rename({"U_A_grid": "U"})
+    ds["U"] = xr.full_like(ds["U"], u_value)
+    ds["V"] = xr.full_like(ds["U"], v_value)
+    if w_value is not None:
+        ds["W"] = xr.full_like(ds["U"], w_value)
+
+    # Slice dataset
+    indexers = {"node_dimension1": x_slice, "node_dimension2": y_slice}
+    if w_value:
+        indexers.update({"vertical_dimensions_dim1": z_slice})
+    ds = ds.sgrid.isel(indexers)
 
-    tdim = 2  # TODO make this also work for length-1 time dimensions
-    dims = (tdim, zdim, ydim, xdim)
-    U = u * np.ones(dims, dtype=np.float32)
-    V = v * np.ones(dims, dtype=np.float32)
-    if w is not None:
-        W = w * np.ones(dims, dtype=np.float32)
-
-    ds = xr.Dataset(
-        {
-            "U": (["time", "depth", "YG", "XG"], U),
-            "V": (["time", "depth", "YG", "XG"], V),
-        },
-        coords={
-            "time": (["time"], [np.timedelta64(0, "s"), np.timedelta64(10, "s")], {"axis": "T"}),
-            "depth": (["depth"], depth, {"axis": "Z"}),
-            "YC": (["YC"], np.arange(ydim) + 0.5, {"axis": "Y"}),
-            "YG": (["YG"], np.arange(ydim), {"axis": "Y", "c_grid_axis_shift": -0.5}),
-            "XC": (["XC"], np.arange(xdim) + 0.5, {"axis": "X"}),
-            "XG": (["XG"], np.arange(xdim), {"axis": "X", "c_grid_axis_shift": -0.5}),
-            "lat": (["YG"], lat, {"axis": "Y", "c_grid_axis_shift": 0.5}),
-            "lon": (["XG"], lon, {"axis": "X", "c_grid_axis_shift": -0.5}),
-        },
-    )
-    if w:
-        ds["W"] = (["time", "depth", "YG", "XG"], W)
-
-    grid = XGrid.from_dataset(ds, mesh="flat")
-    U = Field("U", ds["U"], grid, interp_method=XLinear)
-    V = Field("V", ds["V"], grid, interp_method=XLinear)
-    fields = [U, V, VectorField("UV", U, V, vector_interp_method=XLinear_Velocity)]
-    if w:
-        W = Field("W", ds["W"], grid, interp_method=XLinear)
-        fields.append(VectorField("UVW", U, V, W, vector_interp_method=XLinear_Velocity))
-    fieldset = FieldSet(fields)
-
-    x0, y0, z0 = 2, 8, -4
+    fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
+
+    x0, y0, z0 = 3, 3, 20
     pset = ParticleSet(fieldset, lon=x0, lat=y0, z=z0)
-    kernel = AdvectionRK4 if w is None else AdvectionRK4_3D
+    kernel = AdvectionRK4 if w_value is None else AdvectionRK4_3D
     pset.execute(kernel, runtime=np.timedelta64(4, "s"), dt=np.timedelta64(1, "s"))
 
     assert len(pset.lon) == len([p.lon for p in pset])
-    np.testing.assert_allclose(np.array([p.lon - x0 for p in pset]), 4 * u, atol=1e-6)
-    np.testing.assert_allclose(np.array([p.lat - y0 for p in pset]), 4 * v, atol=1e-6)
-    if w:
-        np.testing.assert_allclose(np.array([p.z - z0 for p in pset]), 4 * w, atol=1e-6)
+    np.testing.assert_allclose(np.array([p.lon - x0 for p in pset]), 4 * u_value, atol=1e-5)
+    np.testing.assert_allclose(np.array([p.lat - y0 for p in pset]), 4 * v_value, atol=1e-5)
+    if w_value:
+        np.testing.assert_allclose(np.array([p.z - z0 for p in pset]), 4 * w_value, atol=1e-5)
 
 
 def test_radialrotation(npart=10):
     ds = radial_rotation_dataset()
-    grid = XGrid.from_dataset(ds, mesh="flat")
-    U = parcels.Field("U", ds["U"], grid, interp_method=XLinear)
-    V = parcels.Field("V", ds["V"], grid, interp_method=XLinear)
-    UV = parcels.VectorField("UV", U, V, vector_interp_method=XLinear_Velocity)
-    fieldset = parcels.FieldSet([U, V, UV])
+    fieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh="flat")
 
     dt = np.timedelta64(30, "s")
     lon = np.linspace(32, 50, npart)
@@ -285,22 +261,18 @@ def test_radialrotation(npart=10):
 )
 def test_moving_eddy(kernel, rtol):
     ds = moving_eddy_dataset()
-    grid = XGrid.from_dataset(ds, mesh="flat")
-    U = Field("U", ds["U"], grid, interp_method=XLinear)
-    V = Field("V", ds["V"], grid, interp_method=XLinear)
     if kernel in [AdvectionRK2_3D, AdvectionRK4_3D]:
         # Using W to test 3D advection (assuming same velocity as V)
-        W = Field("W", ds["V"], grid, interp_method=XLinear)
-        UVW = VectorField("UVW", U, V, W, vector_interp_method=XLinear_Velocity)
-        fieldset = FieldSet([U, V, W, UVW])
-    else:
-        UV = VectorField("UV", U, V, vector_interp_method=XLinear_Velocity)
-        fieldset = FieldSet([U, V, UV])
+        ds["W"] = ds["V"]
+
     if kernel in [AdvectionDiffusionEM, AdvectionDiffusionM1]:
         # Add zero diffusivity field for diffusion kernels
-        ds["Kh"] = (["time", "depth", "YG", "XG"], np.full(ds["U"].shape, 0))
-        fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_zonal")
-        fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_meridional")
+        ds["Kh_zonal"] = (["time", "depth", "YG", "XG"], np.full(ds["U"].shape, 0))
+        ds["Kh_meridional"] = ds["Kh_zonal"]
+
+    fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
+
+    if kernel in [AdvectionDiffusionEM, AdvectionDiffusionM1]:
         fieldset.add_constant("dres", 0.1)
 
     start_lon, start_lat, start_z = 12000, 12500, 12500
@@ -339,11 +311,7 @@ def truth_moving(x_0, y_0, t):
 )
 def test_decaying_moving_eddy(kernel, rtol):
     ds = decaying_moving_eddy_dataset()
-    grid = XGrid.from_dataset(ds, mesh="flat")
-    U = Field("U", ds["U"], grid, interp_method=XLinear)
-    V = Field("V", ds["V"], grid, interp_method=XLinear)
-    UV = VectorField("UV", U, V, vector_interp_method=XLinear_Velocity)
-    fieldset = FieldSet([U, V, UV])
+    fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
 
     start_lon, start_lat = 10000, 10000
     dt = np.timedelta64(60, "m")
@@ -388,14 +356,7 @@ def truth_moving(x_0, y_0, t):
 @pytest.mark.parametrize("grid_type", ["A", "C"])
 def test_stommelgyre_fieldset(kernel, rtol, grid_type):
     npart = 2
-    ds = stommel_gyre_dataset(grid_type=grid_type)
-    grid = XGrid.from_dataset(ds, mesh="flat")
-    vector_interp_method = XLinear_Velocity if grid_type == "A" else CGrid_Velocity
-    U = Field("U", ds["U"], grid, interp_method=XLinear)
-    V = Field("V", ds["V"], grid, interp_method=XLinear)
-    P = Field("P", ds["P"], grid, interp_method=XLinear)
-    UV = VectorField("UV", U, V, vector_interp_method=vector_interp_method)
-    fieldset = FieldSet([U, V, P, UV])
+    fieldset = FieldSet.from_sgrid_conventions(stommel_gyre_dataset(grid_type=grid_type), mesh="flat")
 
     dt = np.timedelta64(30, "m")
     runtime = np.timedelta64(1, "D")
@@ -430,12 +391,7 @@ def UpdateP(particles, fieldset):  # pragma: no cover
 def test_peninsula_fieldset(kernel, rtol, grid_type):
     npart = 2
     ds = peninsula_dataset(grid_type=grid_type)
-    grid = XGrid.from_dataset(ds, mesh="flat")
-    U = Field("U", ds["U"], grid, interp_method=XLinear)
-    V = Field("V", ds["V"], grid, interp_method=XLinear)
-    P = Field("P", ds["P"], grid, interp_method=XLinear)
-    UV = VectorField("UV", U, V, vector_interp_method=XLinear_Velocity)
-    fieldset = FieldSet([U, V, P, UV])
+    fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
 
     dt = np.timedelta64(30, "m")
     runtime = np.timedelta64(23, "h")

tests/test_field.py

@@ -7,7 +7,7 @@
 
 from parcels import Field, UxGrid, VectorField, XGrid
 from parcels._datasets.structured.generic import T as T_structured
-from parcels._datasets.structured.generic import datasets as datasets_structured
+from parcels._datasets.structured.generic import datasets_comodo as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.interpolators import (
     UxConstantFaceConstantZC,