Accept missing_dims in Variable.transpose and Dataset.transpose (#5586)

Author: gcaria

doc/whats-new.rst

@@ -21,6 +21,9 @@ v0.18.3 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- Allow passing argument ``missing_dims`` to :py:meth:`Variable.transpose` and :py:meth:`Dataset.transpose`
+  (:issue:`5550`, :pull:`5586`)
+  By `Giacomo Caria <https://github.com/gcaria>`_.
 - Allow passing a dictionary as coords to a :py:class:`DataArray` (:issue:`5527`,
   reverts :pull:`1539`, which had deprecated this due to python's inconsistent ordering in earlier versions).
   By `Sander van Rijn <https://github.com/sjvrijn>`_.

xarray/core/dataset.py

@@ -4543,7 +4543,11 @@ def drop_dims(
         drop_vars = {k for k, v in self._variables.items() if set(v.dims) & drop_dims}
         return self.drop_vars(drop_vars)
 
-    def transpose(self, *dims: Hashable) -> "Dataset":
+    def transpose(
+        self,
+        *dims: Hashable,
+        missing_dims: str = "raise",
+    ) -> "Dataset":
         """Return a new Dataset object with all array dimensions transposed.
 
         Although the order of dimensions on each array will change, the dataset
@@ -4554,6 +4558,12 @@ def transpose(self, *dims: Hashable) -> "Dataset":
         *dims : hashable, optional
             By default, reverse the dimensions on each array. Otherwise,
             reorder the dimensions to this order.
+        missing_dims : {"raise", "warn", "ignore"}, default: "raise"
+            What to do if dimensions that should be selected from are not present in the
+            Dataset:
+            - "raise": raise an exception
+            - "warn": raise a warning, and ignore the missing dimensions
+            - "ignore": ignore the missing dimensions
 
         Returns
         -------
@@ -4572,12 +4582,10 @@ def transpose(self, *dims: Hashable) -> "Dataset":
         numpy.transpose
         DataArray.transpose
         """
-        if dims:
-            if set(dims) ^ set(self.dims) and ... not in dims:
-                raise ValueError(
-                    f"arguments to transpose ({dims}) must be "
-                    f"permuted dataset dimensions ({tuple(self.dims)})"
-                )
+        # Use infix_dims to check once for missing dimensions
+        if len(dims) != 0:
+            _ = list(infix_dims(dims, self.dims, missing_dims))
+
         ds = self.copy()
         for name, var in self._variables.items():
             var_dims = tuple(dim for dim in dims if dim in (var.dims + (...,)))

xarray/core/variable.py

@@ -1378,14 +1378,24 @@ def roll(self, shifts=None, **shifts_kwargs):
             result = result._roll_one_dim(dim, count)
         return result
 
-    def transpose(self, *dims) -> "Variable":
+    def transpose(
+        self,
+        *dims,
+        missing_dims: str = "raise",
+    ) -> "Variable":
         """Return a new Variable object with transposed dimensions.
 
         Parameters
         ----------
         *dims : str, optional
             By default, reverse the dimensions. Otherwise, reorder the
             dimensions to this order.
+        missing_dims : {"raise", "warn", "ignore"}, default: "raise"
+            What to do if dimensions that should be selected from are not present in the
+            Variable:
+            - "raise": raise an exception
+            - "warn": raise a warning, and ignore the missing dimensions
+            - "ignore": ignore the missing dimensions
 
         Returns
         -------
@@ -1404,7 +1414,9 @@ def transpose(self, *dims) -> "Variable":
         """
         if len(dims) == 0:
             dims = self.dims[::-1]
-        dims = tuple(infix_dims(dims, self.dims))
+        else:
+            dims = tuple(infix_dims(dims, self.dims, missing_dims))
+
         if len(dims) < 2 or dims == self.dims:
             # no need to transpose if only one dimension
             # or dims are in same order

xarray/tests/test_dataset.py

@@ -5195,10 +5195,19 @@ def test_dataset_transpose(self):
             expected_dims = tuple(d for d in new_order if d in ds[k].dims)
             assert actual[k].dims == expected_dims
 
-        with pytest.raises(ValueError, match=r"permuted"):
-            ds.transpose("dim1", "dim2", "dim3")
-        with pytest.raises(ValueError, match=r"permuted"):
-            ds.transpose("dim1", "dim2", "dim3", "time", "extra_dim")
+        # test missing dimension, raise error
+        with pytest.raises(ValueError):
+            ds.transpose(..., "not_a_dim")
+
+        # test missing dimension, ignore error
+        actual = ds.transpose(..., "not_a_dim", missing_dims="ignore")
+        expected_ell = ds.transpose(...)
+        assert_identical(expected_ell, actual)
+
+        # test missing dimension, raise warning
+        with pytest.warns(UserWarning):
+            actual = ds.transpose(..., "not_a_dim", missing_dims="warn")
+            assert_identical(expected_ell, actual)
 
         assert "T" not in dir(ds)
 

xarray/tests/test_variable.py

@@ -1466,6 +1466,20 @@ def test_transpose(self):
         w3 = Variable(["b", "c", "d", "a"], np.einsum("abcd->bcda", x))
         assert_identical(w, w3.transpose("a", "b", "c", "d"))
 
+        # test missing dimension, raise error
+        with pytest.raises(ValueError):
+            v.transpose(..., "not_a_dim")
+
+        # test missing dimension, ignore error
+        actual = v.transpose(..., "not_a_dim", missing_dims="ignore")
+        expected_ell = v.transpose(...)
+        assert_identical(expected_ell, actual)
+
+        # test missing dimension, raise warning
+        with pytest.warns(UserWarning):
+            v.transpose(..., "not_a_dim", missing_dims="warn")
+            assert_identical(expected_ell, actual)
+
     def test_transpose_0d(self):
         for value in [
             3.5,