diff --git a/doc/reshaping.rst b/doc/reshaping.rst index 51202f9be41..455a24f9216 100644 --- a/doc/reshaping.rst +++ b/doc/reshaping.rst @@ -18,12 +18,14 @@ Reordering dimensions --------------------- To reorder dimensions on a :py:class:`~xarray.DataArray` or across all variables -on a :py:class:`~xarray.Dataset`, use :py:meth:`~xarray.DataArray.transpose`: +on a :py:class:`~xarray.Dataset`, use :py:meth:`~xarray.DataArray.transpose`. An +ellipsis (`...`) can be use to represent all other dimensions: .. ipython:: python ds = xr.Dataset({'foo': (('x', 'y', 'z'), [[[42]]]), 'bar': (('y', 'z'), [[24]])}) ds.transpose('y', 'z', 'x') + ds.transpose(..., 'x') # equivalent ds.transpose() # reverses all dimensions Expand and squeeze dimensions diff --git a/doc/whats-new.rst b/doc/whats-new.rst index dea110b5e46..cced7276ff3 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -25,6 +25,10 @@ Breaking changes New Features ~~~~~~~~~~~~ +- :py:meth:`Dataset.transpose` and :py:meth:`DataArray.transpose` now support an ellipsis (`...`) + to represent all 'other' dimensions. For example, to move one dimension to the front, + use `.transpose('x', ...)`. (:pull:`3421`) + By `Maximilian Roos `_ - Changed `xr.ALL_DIMS` to equal python's `Ellipsis` (`...`), and changed internal usages to use `...` directly. As before, you can use this to instruct a `groupby` operation to reduce over all dimensions. While we have no plans to remove `xr.ALL_DIMS`, we suggest diff --git a/setup.cfg b/setup.cfg index eee8b2477b2..fec2ca6bbe4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -117,4 +117,7 @@ tag_prefix = v parentdir_prefix = xarray- [aliases] -test = pytest \ No newline at end of file +test = pytest + +[pytest-watch] +nobeep = True \ No newline at end of file diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 5fccb9236e8..33dcad13204 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1863,12 +1863,7 @@ def transpose(self, *dims: Hashable, transpose_coords: bool = None) -> "DataArra Dataset.transpose """ if dims: - if set(dims) ^ set(self.dims): - raise ValueError( - "arguments to transpose (%s) must be " - "permuted array dimensions (%s)" % (dims, tuple(self.dims)) - ) - + dims = tuple(utils.infix_dims(dims, self.dims)) variable = self.variable.transpose(*dims) if transpose_coords: coords: Dict[Hashable, Variable] = {} diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 55ac0bc6135..2a0464515c6 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3712,14 +3712,14 @@ def transpose(self, *dims: Hashable) -> "Dataset": DataArray.transpose """ if dims: - if set(dims) ^ set(self.dims): + if set(dims) ^ set(self.dims) and ... not in dims: raise ValueError( "arguments to transpose (%s) must be " "permuted dataset dimensions (%s)" % (dims, tuple(self.dims)) ) ds = self.copy() for name, var in self._variables.items(): - var_dims = tuple(dim for dim in dims if dim in var.dims) + var_dims = tuple(dim for dim in dims if dim in (var.dims + (...,))) ds._variables[name] = var.transpose(*var_dims) return ds diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 6befe0b5efc..492c595a887 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -10,6 +10,7 @@ AbstractSet, Any, Callable, + Collection, Container, Dict, Hashable, @@ -660,6 +661,30 @@ def __len__(self) -> int: return len(self._data) - num_hidden +def infix_dims(dims_supplied: Collection, dims_all: Collection) -> Iterator: + """ + Resolves a supplied list containing an ellispsis representing other items, to + a generator with the 'realized' list of all items + """ + if ... in dims_supplied: + if len(set(dims_all)) != len(dims_all): + raise ValueError("Cannot use ellipsis with repeated dims") + if len([d for d in dims_supplied if d == ...]) > 1: + raise ValueError("More than one ellipsis supplied") + other_dims = [d for d in dims_all if d not in dims_supplied] + for d in dims_supplied: + if d == ...: + yield from other_dims + else: + yield d + else: + if set(dims_supplied) ^ set(dims_all): + raise ValueError( + f"{dims_supplied} must be a permuted list of {dims_all}, unless `...` is included" + ) + yield from dims_supplied + + def get_temp_dimname(dims: Container[Hashable], new_dim: Hashable) -> Hashable: """ Get an new dimension name based on new_dim, that is not used in dims. If the same name exists, we add an underscore(s) in the head. diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 93ad1eafb97..7d03fd58d39 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -25,6 +25,7 @@ OrderedSet, decode_numpy_dict_values, either_dict_or_kwargs, + infix_dims, ensure_us_time_resolution, ) @@ -1228,6 +1229,7 @@ def transpose(self, *dims) -> "Variable": """ if len(dims) == 0: dims = self.dims[::-1] + dims = tuple(infix_dims(dims, self.dims)) axes = self.get_axis_num(dims) if len(dims) < 2: # no need to transpose if only one dimension return self.copy(deep=False) diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 88476e5e730..f85a33f7a3c 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -158,18 +158,21 @@ def source_ndarray(array): def assert_equal(a, b): + __tracebackhide__ = True xarray.testing.assert_equal(a, b) xarray.testing._assert_internal_invariants(a) xarray.testing._assert_internal_invariants(b) def assert_identical(a, b): + __tracebackhide__ = True xarray.testing.assert_identical(a, b) xarray.testing._assert_internal_invariants(a) xarray.testing._assert_internal_invariants(b) def assert_allclose(a, b, **kwargs): + __tracebackhide__ = True xarray.testing.assert_allclose(a, b, **kwargs) xarray.testing._assert_internal_invariants(a) xarray.testing._assert_internal_invariants(b) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 101bb44660c..ad474d533be 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -2068,6 +2068,10 @@ def test_transpose(self): ) assert_equal(expected, actual) + # same as previous but with ellipsis + actual = da.transpose("z", ..., "x", transpose_coords=True) + assert_equal(expected, actual) + with pytest.raises(ValueError): da.transpose("x", "y") diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index b3ffdf68e3f..647eb733adb 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -4675,6 +4675,10 @@ def test_dataset_transpose(self): ) assert_identical(expected, actual) + actual = ds.transpose(...) + expected = ds + assert_identical(expected, actual) + actual = ds.transpose("x", "y") expected = ds.apply(lambda x: x.transpose("x", "y", transpose_coords=True)) assert_identical(expected, actual) @@ -4690,13 +4694,32 @@ 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 raises_regex(ValueError, "arguments to transpose"): + # same as above but with ellipsis + new_order = ("dim2", "dim3", "dim1", "time") + actual = ds.transpose("dim2", "dim3", ...) + for k in ds.variables: + expected_dims = tuple(d for d in new_order if d in ds[k].dims) + assert actual[k].dims == expected_dims + + with raises_regex(ValueError, "permuted"): ds.transpose("dim1", "dim2", "dim3") - with raises_regex(ValueError, "arguments to transpose"): + with raises_regex(ValueError, "permuted"): ds.transpose("dim1", "dim2", "dim3", "time", "extra_dim") assert "T" not in dir(ds) + def test_dataset_ellipsis_transpose_different_ordered_vars(self): + # https://github.com/pydata/xarray/issues/1081#issuecomment-544350457 + ds = Dataset( + dict( + a=(("w", "x", "y", "z"), np.ones((2, 3, 4, 5))), + b=(("x", "w", "y", "z"), np.zeros((3, 2, 4, 5))), + ) + ) + result = ds.transpose(..., "z", "y") + assert list(result["a"].dims) == list("wxzy") + assert list(result["b"].dims) == list("xwzy") + def test_dataset_retains_period_index_on_transpose(self): ds = create_test_data() diff --git a/xarray/tests/test_utils.py b/xarray/tests/test_utils.py index c36e8a1775d..5bb9deaf240 100644 --- a/xarray/tests/test_utils.py +++ b/xarray/tests/test_utils.py @@ -275,3 +275,27 @@ def test_either_dict_or_kwargs(): with pytest.raises(ValueError, match=r"foo"): result = either_dict_or_kwargs(dict(a=1), dict(a=1), "foo") + + +@pytest.mark.parametrize( + ["supplied", "all_", "expected"], + [ + (list("abc"), list("abc"), list("abc")), + (["a", ..., "c"], list("abc"), list("abc")), + (["a", ...], list("abc"), list("abc")), + (["c", ...], list("abc"), list("cab")), + ([..., "b"], list("abc"), list("acb")), + ([...], list("abc"), list("abc")), + ], +) +def test_infix_dims(supplied, all_, expected): + result = list(utils.infix_dims(supplied, all_)) + assert result == expected + + +@pytest.mark.parametrize( + ["supplied", "all_"], [([..., ...], list("abc")), ([...], list("aac"))] +) +def test_infix_dims_errors(supplied, all_): + with pytest.raises(ValueError): + list(utils.infix_dims(supplied, all_)) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 78723eda013..528027ed149 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -1280,6 +1280,9 @@ def test_transpose(self): w2 = Variable(["d", "b", "c", "a"], np.einsum("abcd->dbca", x)) assert w2.shape == (5, 3, 4, 2) assert_identical(w2, w.transpose("d", "b", "c", "a")) + assert_identical(w2, w.transpose("d", ..., "a")) + assert_identical(w2, w.transpose("d", "b", "c", ...)) + assert_identical(w2, w.transpose(..., "b", "c", "a")) assert_identical(w, w2.transpose("a", "b", "c", "d")) w3 = Variable(["b", "c", "d", "a"], np.einsum("abcd->bcda", x)) assert_identical(w, w3.transpose("a", "b", "c", "d"))