Removing default mesh in XGrid and UxGrid

This addresses #2408

Author: erikvansebille

src/parcels/_core/fieldset.py

@@ -10,7 +10,6 @@
 import xarray as xr
 import xgcm
 
-from parcels._core.converters import Geographic, GeographicPolar
 from parcels._core.field import Field, VectorField
 from parcels._core.utils.string import _assert_str_and_python_varname
 from parcels._core.utils.time import get_datetime_type_calendar
@@ -277,15 +276,13 @@ def from_fesom2(ds: ux.UxDataset):
             raise ValueError(
                 f"Dataset missing one of the required dimensions 'time', 'nz', or 'nz1'. Found dimensions {ds_dims}"
             )
-        grid = UxGrid(ds.uxgrid, z=ds.coords["nz"])
+        grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="spherical")
         ds = _discover_fesom2_U_and_V(ds)
 
         fields = {}
         if "U" in ds.data_vars and "V" in ds.data_vars:
             fields["U"] = Field("U", ds["U"], grid, _select_uxinterpolator(ds["U"]))
             fields["V"] = Field("V", ds["V"], grid, _select_uxinterpolator(ds["U"]))
-            fields["U"].units = GeographicPolar()
-            fields["V"].units = Geographic()
 
             if "W" in ds.data_vars:
                 fields["W"] = Field("W", ds["W"], grid, _select_uxinterpolator(ds["U"]))

src/parcels/_core/uxgrid.py

@@ -18,7 +18,7 @@ class UxGrid(BaseGrid):
     for interpolation on unstructured grids.
     """
 
-    def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh="flat") -> UxGrid:
+    def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh) -> UxGrid:
         """
         Initializes the UxGrid with a uxarray grid and vertical coordinate array.
 
@@ -30,8 +30,8 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh="flat") -> UxGrid
             A 1D array of vertical coordinates (depths) associated with the layer interface heights (not the mid-layer depths).
             While uxarray allows nz to be spatially and temporally varying, the parcels.UxGrid class considers the case where
             the vertical coordinate is constant in time and space. This implies flat bottom topography and no moving ALE vertical grid.
-        mesh : str, optional
-            The type of mesh used for the grid. Either "flat" (default) or "spherical".
+        mesh : str
+            The type of mesh used for the grid. Either "flat" or "spherical".
         """
         self.uxgrid = grid
         if not isinstance(z, ux.UxDataArray):

src/parcels/_core/xgrid.py

@@ -102,7 +102,7 @@ class XGrid(BaseGrid):
 
     """
 
-    def __init__(self, grid: xgcm.Grid, mesh="flat"):
+    def __init__(self, grid: xgcm.Grid, mesh):
         self.xgcm_grid = grid
         self._mesh = mesh
         self._spatialhash = None

tests/test_field.py

@@ -55,6 +55,7 @@ def test_field_init_param_types():
             UxGrid(
                 datasets_unstructured["stommel_gyre_delaunay"].uxgrid,
                 z=datasets_unstructured["stommel_gyre_delaunay"].coords["nz"],
+                mesh="flat",
             ),
             id="xarray-uxgrid",
         ),
@@ -194,7 +195,7 @@ def test_field_unstructured_z_linear():
     for k, z in enumerate(ds.coords["nz"]):
         ds["W"].values[:, k, :] = z
 
-    grid = UxGrid(ds.uxgrid, z=ds.coords["nz"])
+    grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat")
     # Note that the vertical coordinate is required to be the position of the layer interfaces ("nz"), not the mid-layers ("nz1")
     P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseConstantFace)
 
@@ -232,7 +233,7 @@ def test_field_unstructured_z_linear():
 def test_field_constant_in_time():
     """Tests field evaluation for a field with no time interval (i.e., constant in time)."""
     ds = datasets_unstructured["stommel_gyre_delaunay"]
-    grid = UxGrid(ds.uxgrid, z=ds.coords["nz"])
+    grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat")
     # Note that the vertical coordinate is required to be the position of the layer interfaces ("nz"), not the mid-layers ("nz1")
     P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseConstantFace)
 

tests/test_uxarray_fieldset.py

@@ -38,13 +38,13 @@ def uv_fesom_channel(ds_fesom_channel) -> VectorField:
         U=Field(
             name="U",
             data=ds_fesom_channel.U,
-            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]),
+            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"),
             interp_method=UXPiecewiseConstantFace,
         ),
         V=Field(
             name="V",
             data=ds_fesom_channel.V,
-            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]),
+            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"),
             interp_method=UXPiecewiseConstantFace,
         ),
     )
@@ -58,19 +58,19 @@ def uvw_fesom_channel(ds_fesom_channel) -> VectorField:
         U=Field(
             name="U",
             data=ds_fesom_channel.U,
-            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]),
+            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"),
             interp_method=UXPiecewiseConstantFace,
         ),
         V=Field(
             name="V",
             data=ds_fesom_channel.V,
-            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]),
+            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"),
             interp_method=UXPiecewiseConstantFace,
         ),
         W=Field(
             name="W",
             data=ds_fesom_channel.W,
-            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]),
+            grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"),
             interp_method=UXPiecewiseLinearNode,
         ),
     )
@@ -112,13 +112,14 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval():
     Since the underlying data is constant, we can check that the values are as expected.
     """
     ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"]
+    grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat")
     UVW = VectorField(
         name="UVW",
-        U=Field(name="U", data=ds.U, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace),
-        V=Field(name="V", data=ds.V, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace),
-        W=Field(name="W", data=ds.W, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode),
+        U=Field(name="U", data=ds.U, grid=grid, interp_method=UXPiecewiseConstantFace),
+        V=Field(name="V", data=ds.V, grid=grid, interp_method=UXPiecewiseConstantFace),
+        W=Field(name="W", data=ds.W, grid=grid, interp_method=UXPiecewiseLinearNode),
     )
-    P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode)
+    P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseLinearNode)
     fieldset = FieldSet([UVW, P, UVW.U, UVW.V, UVW.W])
 
     assert fieldset.U.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0

tests/test_uxgrid.py

@@ -6,18 +6,25 @@
 
 @pytest.mark.parametrize("uxds", [pytest.param(uxds, id=key) for key, uxds in uxdatasets.items()])
 def test_uxgrid_init_on_generic_datasets(uxds):
-    UxGrid(uxds.uxgrid, z=uxds.coords["nz"])
+    UxGrid(uxds.uxgrid, z=uxds.coords["nz"], mesh="flat")
 
 
 @pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]])
 def test_uxgrid_axes(uxds):
-    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"])
+    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"], mesh="flat")
     assert grid.axes == ["Z", "FACE"]
 
 
+@pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]])
+@pytest.mark.parametrize("mesh", ["flat", "spherical"])
+def test_uxgrid_mesh(uxds, mesh):
+    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"], mesh=mesh)
+    assert grid._mesh == mesh
+
+
 @pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]])
 def test_xgrid_get_axis_dim(uxds):
-    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"])
+    grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"], mesh="flat")
 
     assert grid.get_axis_dim("FACE") == 721
     assert grid.get_axis_dim("Z") == 2

tests/test_xgrid.py

@@ -65,12 +65,18 @@ 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.from_dataset(ds)
     assert grid.axes == ["Z", "Y", "X"]
 
 
+@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
+@pytest.mark.parametrize("mesh", ["flat", "spherical"])
+def test_uxgrid_mesh(ds, mesh):
+    grid = XGrid.from_dataset(ds, mesh=mesh)
+    assert grid._mesh == mesh
+
+
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_transpose_xfield_data_to_tzyx(ds):
     da = ds["data_g"]