Removing default mesh from XGrid.from_dataset()

Author: erikvansebille

src/parcels/_core/xgrid.py

@@ -124,7 +124,7 @@ def __init__(self, grid: xgcm.Grid, mesh):
         self._ds = ds
 
     @classmethod
-    def from_dataset(cls, ds: xr.Dataset, mesh="flat", xgcm_kwargs=None):
+    def from_dataset(cls, ds: xr.Dataset, mesh, 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 = {}

tests/test_advection.py

@@ -84,7 +84,7 @@ def test_advection_zonal_periodic():
     halo.XG.values = ds.XG.values[1] + 2
     ds = xr.concat([ds, halo], dim="XG")
 
-    grid = XGrid.from_dataset(ds)
+    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)
@@ -103,7 +103,7 @@ def test_horizontal_advection_in_3D_flow(npart=10):
     """Flat 2D zonal flow that increases linearly with z from 0 m/s to 1 m/s."""
     ds = simple_UV_dataset(mesh="flat")
     ds["U"].data[:] = 1.0
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     U = Field("U", ds["U"], grid, interp_method=XLinear)
     U.data[:, 0, :, :] = 0.0  # Set U to 0 at the surface
     V = Field("V", ds["V"], grid, interp_method=XLinear)
@@ -121,7 +121,7 @@ def test_horizontal_advection_in_3D_flow(npart=10):
 @pytest.mark.parametrize("wErrorThroughSurface", [True, False])
 def test_advection_3D_outofbounds(direction, wErrorThroughSurface):
     ds = simple_UV_dataset(mesh="flat")
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     U = Field("U", ds["U"], grid, interp_method=XLinear)
     U.data[:] = 0.01  # Set U to small value (to avoid horizontal out of bounds)
     V = Field("V", ds["V"], grid, interp_method=XLinear)
@@ -202,7 +202,7 @@ def test_length1dimensions(u, v, w):  # TODO: Refactor this test to be more read
     if w:
         ds["W"] = (["time", "depth", "YG", "XG"], W)
 
-    grid = XGrid.from_dataset(ds)
+    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)]
@@ -263,7 +263,7 @@ def test_radialrotation(npart=10):
 )
 def test_moving_eddy(kernel, rtol):
     ds = moving_eddy_dataset()
-    grid = XGrid.from_dataset(ds)
+    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]:
@@ -315,7 +315,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)
+    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)
@@ -363,7 +363,7 @@ def truth_moving(x_0, y_0, t):
 def test_stommelgyre_fieldset(kernel, rtol, grid_type):
     npart = 2
     ds = stommel_gyre_dataset(grid_type=grid_type)
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     vector_interp_method = None 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)
@@ -404,7 +404,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)
+    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)

tests/test_field.py

@@ -14,7 +14,7 @@
 
 def test_field_init_param_types():
     data = datasets_structured["ds_2d_left"]
-    grid = XGrid.from_dataset(data)
+    grid = XGrid.from_dataset(data, mesh="flat")
 
     with pytest.raises(TypeError, match="Expected a string for variable name, got int instead."):
         Field(name=123, data=data["data_g"], grid=grid, interp_method=XLinear)
@@ -47,7 +47,7 @@ def test_field_init_param_types():
     [
         pytest.param(
             ux.UxDataArray(),
-            XGrid.from_dataset(datasets_structured["ds_2d_left"]),
+            XGrid.from_dataset(datasets_structured["ds_2d_left"], mesh="flat"),
             id="uxdata-grid",
         ),
         pytest.param(
@@ -76,7 +76,7 @@ def test_field_incompatible_combination(data, grid):
     [
         pytest.param(
             datasets_structured["ds_2d_left"]["data_g"],
-            XGrid.from_dataset(datasets_structured["ds_2d_left"]),
+            XGrid.from_dataset(datasets_structured["ds_2d_left"], mesh="flat"),
             id="ds_2d_left",
         ),  # TODO: Perhaps this test should be expanded to cover more datasets?
     ],
@@ -107,7 +107,7 @@ def test_field_init_fail_on_float_time_dim():
     )
 
     data = ds["data_g"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     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\?",
@@ -125,7 +125,7 @@ def test_field_init_fail_on_float_time_dim():
     [
         pytest.param(
             datasets_structured["ds_2d_left"]["data_g"],
-            XGrid.from_dataset(datasets_structured["ds_2d_left"]),
+            XGrid.from_dataset(datasets_structured["ds_2d_left"], mesh="flat"),
             id="ds_2d_left",
         ),
     ],
@@ -144,7 +144,7 @@ def test_vectorfield_init_different_time_intervals():
 
 def test_field_invalid_interpolator():
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
 
     def invalid_interpolator_wrong_signature(particle_positions, grid_positions, invalid):
         return 0.0
@@ -161,7 +161,7 @@ def invalid_interpolator_wrong_signature(particle_positions, grid_positions, inv
 
 def test_vectorfield_invalid_interpolator():
     ds = datasets_structured["ds_2d_left"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
 
     def invalid_interpolator_wrong_signature(particle_positions, grid_positions, invalid):
         return 0.0

tests/test_index_search.py

@@ -13,7 +13,7 @@
 @pytest.fixture
 def field_cone():
     ds = datasets["2d_left_unrolled_cone"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     field = Field(
         name="test_field",
         data=ds["data_g"],

tests/test_interpolation.py

@@ -52,7 +52,7 @@ def field():
             "y": (["y"], [0.5, 1.5, 2.5, 3.5], {"axis": "Y"}),
         },
     )
-    return Field("U", ds["U"], XGrid.from_dataset(ds), interp_method=XLinear)
+    return Field("U", ds["U"], XGrid.from_dataset(ds, mesh="flat"), interp_method=XLinear)
 
 
 @pytest.mark.parametrize(

tests/test_particlefile.py

@@ -32,7 +32,7 @@
 def fieldset() -> FieldSet:  # TODO v4: Move into a `conftest.py` file and remove duplicates
     """Fixture to create a FieldSet object for testing."""
     ds = datasets["ds_2d_left"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     U = Field("U", ds["U_A_grid"], grid, XLinear)
     V = Field("V", ds["V_A_grid"], grid, XLinear)
     UV = VectorField("UV", U, V)
@@ -73,7 +73,7 @@ def test_pfile_array_write_zarr_memorystore(fieldset):
 def test_write_fieldset_without_time(tmp_zarrfile):
     ds = peninsula_dataset()  # DataSet without time
     assert "time" not in ds.dims
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     fieldset = FieldSet([Field("U", ds["U"], grid, XLinear)])
 
     pset = ParticleSet(fieldset, pclass=Particle, lon=0, lat=0)

tests/test_spatialhash.py

@@ -6,14 +6,14 @@
 
 def test_spatialhash_init():
     ds = datasets["2d_left_rotated"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     spatialhash = grid.get_spatial_hash()
     assert spatialhash is not None
 
 
 def test_invalid_positions():
     ds = datasets["2d_left_rotated"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
 
     j, i, coords = grid.get_spatial_hash().query([np.nan, np.inf], [np.nan, np.inf])
     assert np.all(j == -3)
@@ -22,7 +22,7 @@ def test_invalid_positions():
 
 def test_mixed_positions():
     ds = datasets["2d_left_rotated"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     lat = grid.lat.mean()
     lon = grid.lon.mean()
     y = [lat, np.nan]

tests/test_xgrid.py

@@ -55,18 +55,19 @@ def test_grid_init_param_types(ds):
 
 @pytest.mark.parametrize("ds, attr, expected", test_cases)
 def test_xgrid_properties_ground_truth(ds, attr, expected):
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     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_from_dataset_on_generic_datasets(ds):
-    XGrid.from_dataset(ds)
+    XGrid.from_dataset(ds, mesh="flat")
 
 
+@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_xgrid_axes(ds):
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     assert grid.axes == ["Z", "Y", "X"]
 
 
@@ -80,7 +81,7 @@ def test_uxgrid_mesh(ds, mesh):
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_transpose_xfield_data_to_tzyx(ds):
     da = ds["data_g"]
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
 
     all_combinations = (itertools.combinations(da.dims, n) for n in range(len(da.dims)))
     all_combinations = itertools.chain(*all_combinations)
@@ -93,7 +94,7 @@ def test_transpose_xfield_data_to_tzyx(ds):
 
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_xgrid_get_axis_dim(ds):
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     assert grid.get_axis_dim("Z") == Z - 1
     assert grid.get_axis_dim("Y") == Y - 1
     assert grid.get_axis_dim("X") == X - 1
@@ -113,15 +114,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.from_dataset(ds)
+        XGrid.from_dataset(ds, mesh="flat")
 
     ds = datasets["ds_2d_left"].copy()
     ds["lon"], _ = xr.broadcast(ds["YG"], ds["XG"])
     with pytest.raises(
         ValueError,
         match=".*have different dimensionalities\.",
     ):
-        XGrid.from_dataset(ds)
+        XGrid.from_dataset(ds, mesh="flat")
 
     ds = datasets["ds_2d_left"].copy()
     ds["lon"], ds["lat"] = xr.broadcast(ds["YG"], ds["XG"])
@@ -131,20 +132,20 @@ 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.from_dataset(ds)
+        XGrid.from_dataset(ds, mesh="flat")
 
 
 def test_invalid_depth():
     ds = datasets["ds_2d_left"].copy()
     ds = ds.reindex({"ZG": ds.ZG[::-1]})
 
     with pytest.raises(ValueError, match="Depth DataArray .* must be strictly increasing*"):
-        XGrid.from_dataset(ds)
+        XGrid.from_dataset(ds, mesh="flat")
 
 
 def test_dim_without_axis():
     ds = xr.Dataset({"z1d": (["depth"], [0])}, coords={"depth": [0]})
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     with pytest.raises(ValueError, match='Dimension "depth" has no axis attribute*'):
         Field("z1d", ds["z1d"], grid, XLinear)
 
@@ -155,7 +156,7 @@ def test_vertical1D_field():
         {"z1d": (["depth"], np.linspace(0, 10, nz))},
         coords={"depth": (["depth"], np.linspace(0, 1, nz), {"axis": "Z"})},
     )
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     field = Field("z1d", ds["z1d"], grid, XLinear)
 
     assert field.eval(np.timedelta64(0, "s"), 0.45, 0, 0) == 4.5
@@ -167,7 +168,7 @@ def test_time1D_field():
         {"t1d": (["time"], np.arange(0, len(timerange)))},
         coords={"time": (["time"], timerange, {"axis": "T"})},
     )
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     field = Field("t1d", ds["t1d"], grid, XLinear)
 
     assert field.eval(np.datetime64("2000-01-10T12:00:00"), -20, 5, 6) == 9.5
@@ -181,7 +182,7 @@ def test_time1D_field():
     ],
 )  # for key, ds in datasets.items()])
 def test_xgrid_search_cpoints(ds):
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     lat_array, lon_array = get_2d_fpoint_mesh(grid)
     lat_array, lon_array = corner_to_cell_center_points(lat_array, lon_array)
 
@@ -299,7 +300,7 @@ def test_search_1d_array_some_out_of_bounds(array, x, expected_xi):
 )
 def test_xgrid_localize_zero_position(ds, da_name, expected):
     """Test localize function using left and right datasets."""
-    grid = XGrid.from_dataset(ds)
+    grid = XGrid.from_dataset(ds, mesh="flat")
     da = ds[da_name]
     position = grid.search(0, 0, 0)