From 2a9740ddd0caca81f2a033f33742eef9f28ff02f Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 19 Oct 2019 19:18:05 -0400 Subject: [PATCH 1/9] infix_dims function --- setup.cfg | 5 ++++- xarray/core/utils.py | 22 ++++++++++++++++++++++ xarray/tests/test_utils.py | 25 +++++++++++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) 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/utils.py b/xarray/core/utils.py index c9c5866654d..6e4f1f9b5fa 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -12,6 +12,7 @@ Callable, Container, Dict, + Generator, Hashable, Iterable, Iterator, @@ -660,6 +661,27 @@ def __len__(self) -> int: return len(self._data) - num_hidden +def infix_dims(dims_supplied: list, dims_all: list) -> Generator: + """ + 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: + # we could check if all dims are present, as a future feature + 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/tests/test_utils.py b/xarray/tests/test_utils.py index c36e8a1775d..bcc7940a2c3 100644 --- a/xarray/tests/test_utils.py +++ b/xarray/tests/test_utils.py @@ -275,3 +275,28 @@ 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_)) From 70ae24293e4baaea483a44c7b4ef7294f394c82e Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 20 Oct 2019 18:09:07 -0400 Subject: [PATCH 2/9] implement transpose with ellipsis --- xarray/core/dataarray.py | 5 +++-- xarray/core/dataset.py | 7 +++++-- xarray/core/utils.py | 3 ++- xarray/core/variable.py | 2 ++ xarray/tests/test_dataarray.py | 4 ++++ xarray/tests/test_dataset.py | 7 +++++++ xarray/tests/test_utils.py | 1 - xarray/tests/test_variable.py | 3 +++ 8 files changed, 26 insertions(+), 6 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 4a48f13b86d..6d17aa97aa9 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1863,10 +1863,11 @@ def transpose(self, *dims: Hashable, transpose_coords: bool = None) -> "DataArra Dataset.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 array dimensions (%s)" % (dims, tuple(self.dims)) + "permuted array dimensions (%s) unless `...` is included" + % (dims, tuple(self.dims)) ) variable = self.variable.transpose(*dims) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 6123b42b77e..1b81bd5b748 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3710,12 +3710,15 @@ 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)) + "permuted array dimensions (%s) unless `...` is included" + % (dims, tuple(self.dims)) ) ds = self.copy() + # infix here so the filter `var_dims` below works correctly + dims = tuple(utils.infix_dims(dims, self.dims)) for name, var in self._variables.items(): var_dims = tuple(dim for dim in dims if dim in var.dims) ds._variables[name] = var.transpose(*var_dims) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 6e4f1f9b5fa..84137798f7d 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -10,6 +10,7 @@ AbstractSet, Any, Callable, + Collection, Container, Dict, Generator, @@ -661,7 +662,7 @@ def __len__(self) -> int: return len(self._data) - num_hidden -def infix_dims(dims_supplied: list, dims_all: list) -> Generator: +def infix_dims(dims_supplied: Collection, dims_all: Collection) -> Generator: """ Resolves a supplied list containing an ellispsis representing other items, to a generator with the 'realized' list of all items diff --git a/xarray/core/variable.py b/xarray/core/variable.py index b17597df580..c8e41059461 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/test_dataarray.py b/xarray/tests/test_dataarray.py index 1398e936f37..0b57a551c3e 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", ..., 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 a5c9920f1d9..901befa4798 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -4679,6 +4679,13 @@ 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 + # 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, "arguments to transpose"): ds.transpose("dim1", "dim2", "dim3") with raises_regex(ValueError, "arguments to transpose"): diff --git a/xarray/tests/test_utils.py b/xarray/tests/test_utils.py index bcc7940a2c3..5bb9deaf240 100644 --- a/xarray/tests/test_utils.py +++ b/xarray/tests/test_utils.py @@ -289,7 +289,6 @@ def test_either_dict_or_kwargs(): ], ) def test_infix_dims(supplied, all_, expected): - result = list(utils.infix_dims(supplied, all_)) assert result == expected diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index eb6101fe37d..87c021e707b 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")) From 77c4cd38e8af493d8f9254d6592a01e65adbd9da Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 20 Oct 2019 19:27:26 -0400 Subject: [PATCH 3/9] also infix in dataarray --- xarray/core/dataarray.py | 3 ++- xarray/tests/__init__.py | 3 +++ xarray/tests/test_dataarray.py | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 6d17aa97aa9..9327a975222 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1869,7 +1869,8 @@ def transpose(self, *dims: Hashable, transpose_coords: bool = None) -> "DataArra "permuted array dimensions (%s) unless `...` is included" % (dims, tuple(self.dims)) ) - + # infix here so the filter `coord_dims` below works correctly + dims = tuple(utils.infix_dims(dims, self.dims)) variable = self.variable.transpose(*dims) if transpose_coords: coords: Dict[Hashable, Variable] = {} diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index acf8b67effa..d2d4409f041 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 0b57a551c3e..afc926f4ceb 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -2069,7 +2069,7 @@ def test_transpose(self): assert_equal(expected, actual) # same as previous but with ellipsis - actual = da.transpose("z", ..., transpose_coords=True) + actual = da.transpose("z", ..., "x", transpose_coords=True) assert_equal(expected, actual) with pytest.raises(ValueError): From 2ed122855b7a5efa7307d743e6be78f13ab0f3c3 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 20 Oct 2019 22:08:43 -0400 Subject: [PATCH 4/9] check errors centrally, remove boilerplate from transpose methods --- xarray/core/dataarray.py | 9 +-------- xarray/core/dataset.py | 8 +------- xarray/core/utils.py | 4 ++++ xarray/tests/test_dataset.py | 4 ++-- 4 files changed, 8 insertions(+), 17 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 3c3d0cc0f20..33dcad13204 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1863,14 +1863,7 @@ def transpose(self, *dims: Hashable, transpose_coords: bool = None) -> "DataArra Dataset.transpose """ if dims: - if set(dims) ^ set(self.dims) and ... not in dims: - raise ValueError( - "arguments to transpose (%s) must be " - "permuted array dimensions (%s) unless `...` is included" - % (dims, tuple(self.dims)) - ) - # infix here so the filter `coord_dims` below works correctly - dims = tuple(utils.infix_dims(dims, 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 7b244469ddf..90b21052ac9 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3708,15 +3708,9 @@ def transpose(self, *dims: Hashable) -> "Dataset": DataArray.transpose """ if dims: - if set(dims) ^ set(self.dims) and ... not in dims: - raise ValueError( - "arguments to transpose (%s) must be " - "permuted array dimensions (%s) unless `...` is included" - % (dims, tuple(self.dims)) - ) + dims = tuple(utils.infix_dims(dims, self.dims)) ds = self.copy() # infix here so the filter `var_dims` below works correctly - dims = tuple(utils.infix_dims(dims, self.dims)) for name, var in self._variables.items(): var_dims = tuple(dim for dim in dims if dim in var.dims) ds._variables[name] = var.transpose(*var_dims) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 98c754e13f0..2b409bc6f4e 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -679,6 +679,10 @@ def infix_dims(dims_supplied: Collection, dims_all: Collection) -> Generator: 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" + ) # we could check if all dims are present, as a future feature yield from dims_supplied diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index c7ec1b1b477..8dd256fc0a5 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -4685,9 +4685,9 @@ 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"): + 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) From ca4202bf8bc327af7adc7012515c39a320dfdc65 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 20 Oct 2019 22:14:14 -0400 Subject: [PATCH 5/9] whatsnew --- doc/whats-new.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index ab9a0adc101..9d2fb573003 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -20,6 +20,10 @@ v0.14.1 (unreleased) 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, + pass `.transpose('x', ...)`. (:pull:`3421`) + By `Maximilian Roos `_ - Added integration tests against `pint `_. (:pull:`3238`) by `Justus Magin `_. From 65eddca4afcbaf0de60951acdbf3f15dfb98c687 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 20 Oct 2019 22:18:32 -0400 Subject: [PATCH 6/9] docs --- doc/reshaping.rst | 4 +++- doc/whats-new.rst | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) 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 9d2fb573003..a5a911f2dd2 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,7 +22,7 @@ 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, - pass `.transpose('x', ...)`. (:pull:`3421`) + use `.transpose('x', ...)`. (:pull:`3421`) By `Maximilian Roos `_ - Added integration tests against `pint `_. (:pull:`3238`) by `Justus Magin `_. From 87ac7ac797e2c6d7fcfc54366d315709f0fb03ba Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 20 Oct 2019 23:21:23 -0400 Subject: [PATCH 7/9] remove old comments --- xarray/core/dataset.py | 1 - xarray/core/utils.py | 1 - 2 files changed, 2 deletions(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 90b21052ac9..004280606a5 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3710,7 +3710,6 @@ def transpose(self, *dims: Hashable) -> "Dataset": if dims: dims = tuple(utils.infix_dims(dims, self.dims)) ds = self.copy() - # infix here so the filter `var_dims` below works correctly for name, var in self._variables.items(): var_dims = tuple(dim for dim in dims if dim in var.dims) ds._variables[name] = var.transpose(*var_dims) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 2b409bc6f4e..960cd1151f4 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -683,7 +683,6 @@ def infix_dims(dims_supplied: Collection, dims_all: Collection) -> Generator: raise ValueError( f"{dims_supplied} must be a permuted list of {dims_all}, unless `...` is included" ) - # we could check if all dims are present, as a future feature yield from dims_supplied From c21a67aac00ca371d6f0871baa99e53955f6c22e Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Mon, 21 Oct 2019 01:40:36 -0400 Subject: [PATCH 8/9] generator->iterator --- xarray/core/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 960cd1151f4..492c595a887 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -13,7 +13,6 @@ Collection, Container, Dict, - Generator, Hashable, Iterable, Iterator, @@ -662,7 +661,7 @@ def __len__(self) -> int: return len(self._data) - num_hidden -def infix_dims(dims_supplied: Collection, dims_all: Collection) -> Generator: +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 From 86ad50626f2b846e1fa06dc2f4f709b26813cccc Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sun, 27 Oct 2019 16:09:42 -0400 Subject: [PATCH 9/9] test for differently ordered dimensions --- xarray/core/dataset.py | 8 ++++++-- xarray/tests/test_dataset.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index e3c14f01cd5..2a0464515c6 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3712,10 +3712,14 @@ def transpose(self, *dims: Hashable) -> "Dataset": DataArray.transpose """ if dims: - dims = tuple(utils.infix_dims(dims, 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/tests/test_dataset.py b/xarray/tests/test_dataset.py index b42af4c0ed1..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) @@ -4704,6 +4708,18 @@ def test_dataset_transpose(self): 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()