Replace deprecated xarray .drop() and Dataset.dims[...] (#2734)

Clears the FutureWarnings that originate in Parcels' own code for two
deprecated xarray patterns, ahead of enabling fail-on-warning (#2413):

- .drop([...]) / .drop("...") -> .drop_vars(...)
- Dataset.dims[name] / .dims.items() / .dims.get() -> .sizes...

Leaves positional DataArray.dims[0] access untouched, since only the
Dataset.dims mapping-by-name is deprecated.

Towards #2682.

Signed-off-by: Mike German <mike@stepsventures.com>

Author: steps-re

src/parcels/_datasets/remote.py

@@ -184,12 +184,12 @@ def open_dataset(self) -> xr.Dataset:
 def _preprocess_drop_time_from_mesh1(ds: xr.Dataset) -> xr.Dataset:
     # For some reason on the mesh "NemoNorthSeaORCA025-N006_data/coordinates.nc" there are two time dimensions (of length 1). These dimension also has broken cf-time metadata
     # this fixes that
-    return ds.isel(time=0).drop(["time", "time_steps"])
+    return ds.isel(time=0).drop_vars(["time", "time_steps"])
 
 
 def _preprocess_drop_time_from_mesh2(ds: xr.Dataset) -> xr.Dataset:
     # For some reason on the mesh "NemoCurvilinear_data_zonal/mesh_mask" there is a time dimension.
-    return ds.isel(time=0).drop(["time"])
+    return ds.isel(time=0).drop_vars(["time"])
 
 
 def _preprocess_set_cf_calendar_360_day(ds: xr.Dataset) -> xr.Dataset:

src/parcels/_datasets/utils.py

@@ -123,7 +123,7 @@ def verbose_print(*args, **kwargs):
     for dim_name in ds1_dims.intersection(ds2_dims):
         verbose_print(f"  Dimension '{dim_name}':")
         # Sizes will differ due to DIM_SIZE, so we don't strictly compare them.
-        verbose_print(f"    {ds1_name} size: {ds1.dims[dim_name]}, {ds2_name} size: {ds2.dims[dim_name]}")
+        verbose_print(f"    {ds1_name} size: {ds1.sizes[dim_name]}, {ds2_name} size: {ds2.sizes[dim_name]}")
         # Check if coordinates associated with dimensions are sorted (increasing)
         if dim_name in ds1.coords and dim_name in ds2.coords:
             check_val = (

src/parcels/_sgrid/accessor.py

@@ -98,11 +98,11 @@ def assert_metadata_ds_consistency(ds: xr.Dataset, metadata: SGrid2DMetadata):
         face, node, padding = obj.face, obj.node, obj.padding
 
         try:
-            n_nodes = ds.dims[node]
+            n_nodes = ds.sizes[node]
         except KeyError:  # node dimension is not in this dataset
             continue
         try:
-            n_faces = ds.dims[face]
+            n_faces = ds.sizes[face]
         except KeyError:  # face dimension is not in this dataset
             continue
 

src/parcels/convert.py

@@ -319,7 +319,7 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da
     if "time" in coords.dims:
         if coords.sizes["time"] != 1:
             raise ValueError("Time dimension in coords must be length 1 (i.e., no time-varying grid).")
-        coords = coords.isel(time=0).drop("time")
+        coords = coords.isel(time=0).drop_vars("time")
 
     if (
         len(coords.dims) == 3

tests/sgrid/test_accessor.py

@@ -151,10 +151,10 @@ def test_isel(ds, indexer, node_dim, face_dim):
     assert_metadata_ds_consistency(ds_trimmed, metadata)
 
     # Assert that other dims haven't been affected
-    for dim, size_before in ds.dims.items():
+    for dim, size_before in ds.sizes.items():
         if dim in (node_dim, face_dim):
             continue
-        size_after = ds_trimmed.dims[dim]
+        size_after = ds_trimmed.sizes[dim]
         assert size_before == size_after
 
 
@@ -167,8 +167,8 @@ def test_isel_drop_dim(ds):
     assert "node_dimension1" not in ds_trimmed.dims
     assert "face_dimension1" not in ds_trimmed.dims
 
-    for dim, size_after in ds_trimmed.dims.items():
-        size_before = ds.dims[dim]
+    for dim, size_after in ds_trimmed.sizes.items():
+        size_before = ds.sizes[dim]
         assert size_before == size_after
 
 
@@ -201,7 +201,7 @@ def test_isel_p1_consistency_invariant(ds, s):
     fnp = metadata.face_dimensions[0]
     node_dim = fnp.node
     assume(node_dim in ds.dims)
-    n_nodes = ds.dims[node_dim]
+    n_nodes = ds.sizes[node_dim]
     assume(len(range(*s.indices(n_nodes))) > 0)  # exclude empty selections
 
     result = ds.sgrid.isel({node_dim: s})
@@ -219,7 +219,7 @@ def test_isel_p2_data_correctness(ds, s):
     fnp = metadata.face_dimensions[0]
     node_dim, face_dim = fnp.node, fnp.face
     assume(node_dim in ds.dims)
-    n_nodes = ds.dims[node_dim]
+    n_nodes = ds.sizes[node_dim]
     assume(len(range(*s.indices(n_nodes))) > 0)
 
     result = ds.sgrid.isel({node_dim: s})
@@ -249,14 +249,14 @@ def test_isel_p3_specification_symmetry(ds, s):
     fnp = metadata.face_dimensions[0]
     node_dim, face_dim = fnp.node, fnp.face
     assume(node_dim in ds.dims and face_dim in ds.dims)
-    n_nodes = ds.dims[node_dim]
+    n_nodes = ds.sizes[node_dim]
     assume(len(range(*s.indices(n_nodes))) > 0)
 
     from parcels._sgrid.accessor import _derive_paired_indexer
 
     _, face_indexer = _derive_paired_indexer(s, indexer_is_node=True, padding=fnp.padding, dim_size=n_nodes)
 
-    n_faces = ds.dims.get(face_dim)
+    n_faces = ds.sizes.get(face_dim)
     if n_faces is not None:
         assume(len(range(*face_indexer.indices(n_faces))) > 0)
 

tests/test_xgrid.py

@@ -189,7 +189,7 @@ def test_dim_with_duplicate_axis():
 @pytest.mark.skip("remove: see comment above")
 @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]])
 def test_vertical1D_field(ds):
-    ds = ds.drop(set(ds.data_vars) - {"grid"})
+    ds = ds.drop_vars(set(ds.data_vars) - {"grid"})
     ds["depth"] = (["ZG"], np.linspace(0, 1, ds["depth"].size), {"axis": "Z"})
     ds["z1d"] = xr.DataArray(np.linspace(0, 10, ds["depth"].size), dims=("ZG",))
     ds = ds.reset_coords("z1d")