From 24cffafc0d545002e319c4907a2591d06311aefc Mon Sep 17 00:00:00 2001 From: Mathias Hauser Date: Mon, 21 Oct 2019 13:40:36 +0200 Subject: [PATCH 1/4] enable xr.ALL_DIMS in xr.dot --- doc/whats-new.rst | 4 ++++ xarray/core/computation.py | 17 ++++++++++++++--- xarray/core/dataarray.py | 2 +- xarray/tests/test_computation.py | 17 +++++++++++++++++ xarray/tests/test_dataarray.py | 10 ++++++++++ 5 files changed, 46 insertions(+), 4 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index abe472cc6bb..f2528699879 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -31,6 +31,10 @@ New Features ``pip install git+https://github.com/andrewgsavage/pint.git@refs/pull/6/head)``. Even with it, interaction with non-numpy array libraries, e.g. dask or sparse, is broken. +- :py:func:`~xarray.dot`, and :py:func:`~xarray.DataArray.dot` now support the + ``dims=xarray.ALL_DIMS`` option to sum over the union of dimensions of all + arrays (:issue:`3423`) by `Mathias Hauser `_. + Bug fixes ~~~~~~~~~ - Fix regression introduced in v0.14.0 that would cause a crash if dask is installed diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 1393d76f283..3bc0bd591e4 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -25,6 +25,7 @@ from . import duck_array_ops, utils from .alignment import deep_align +from .common import ALL_DIMS from .merge import merge_coordinates_without_align from .pycompat import dask_array_type from .utils import is_dict_like @@ -1055,7 +1056,7 @@ def dot(*arrays, dims=None, **kwargs): ---------- arrays: DataArray (or Variable) objects Arrays to compute. - dims: str or tuple of strings, optional + dims: xarray.ALL_DIMS, str or tuple of strings, optional Which dimensions to sum over. If not speciified, then all the common dimensions are summed over. **kwargs: dict @@ -1070,7 +1071,7 @@ def dot(*arrays, dims=None, **kwargs): -------- >>> import numpy as np - >>> import xarray as xp + >>> import xarray as xr >>> da_a = xr.DataArray(np.arange(3 * 2).reshape(3, 2), dims=['a', 'b']) >>> da_b = xr.DataArray(np.arange(3 * 2 * 2).reshape(3, 2, 2), ... dims=['a', 'b', 'c']) @@ -1117,6 +1118,14 @@ def dot(*arrays, dims=None, **kwargs): [273, 446, 619]]) Dimensions without coordinates: a, d + >>> xr.dot(da_a, da_b) + + array([110, 125]) + Dimensions without coordinates: c + + >>> xr.dot(da_a, da_b, dims=xr.ALL_DIMS) + + array(235) """ from .dataarray import DataArray from .variable import Variable @@ -1141,7 +1150,9 @@ def dot(*arrays, dims=None, **kwargs): einsum_axes = "abcdefghijklmnopqrstuvwxyz" dim_map = {d: einsum_axes[i] for i, d in enumerate(all_dims)} - if dims is None: + if dims is ALL_DIMS: + dims = all_dims + elif dims is None: # find dimensions that occur more than one times dim_counts = Counter() for arr in arrays: diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 5fccb9236e8..e9e2c0f9055 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -2747,7 +2747,7 @@ def dot( ---------- other : DataArray The other array with which the dot product is performed. - dims: hashable or sequence of hashables, optional + dims: xarray.ALL_DIMS, hashable or sequence of hashables, optional Along which dimensions to be summed over. Default all the common dimensions are summed over. diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 383427b479b..7ec0b7656b5 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -998,6 +998,23 @@ def test_dot(use_dask): assert actual.dims == ("b",) assert (actual.data == np.zeros(actual.shape)).all() + # xr.ALL_DIMS + actual = xr.dot(da_a, da_b, dims=xr.ALL_DIMS) + assert actual.dims == () + assert (actual.data == np.einsum("ij,ijk->", a, b)).all() + + actual = xr.dot(da_a, da_b, da_c, dims=xr.ALL_DIMS) + assert actual.dims == () + assert (actual.data == np.einsum("ij,ijk,kl-> ", a, b, c)).all() + + actual = xr.dot(da_a, dims=xr.ALL_DIMS) + assert actual.dims == () + assert (actual.data == np.einsum("ij-> ", a)).all() + + actual = xr.dot(da_a.sel(a=[]), da_a.sel(a=[]), dims=xr.ALL_DIMS) + assert actual.dims == () + assert (actual.data == np.zeros(actual.shape)).all() + # Invalid cases if not use_dask: with pytest.raises(TypeError): diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index d05a02ae705..e28d129500e 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -3930,6 +3930,16 @@ def test_dot(self): expected = DataArray(expected_vals, coords=[x, j], dims=["x", "j"]) assert_equal(expected, actual) + # xr.ALL_DIMS: all shared dims + actual = da.dot(da, dims=xr.ALL_DIMS) + expected = da.dot(da) + assert_equal(expected, actual) + + # xr.ALL_DIMS: multiple shared dims + actual = da.dot(dm, dims=xr.ALL_DIMS) + expected = da.dot(dm, dims=("j", "x", "y", "z")) + assert_equal(expected, actual) + with pytest.raises(NotImplementedError): da.dot(dm.to_dataset(name="dm")) with pytest.raises(TypeError): From fb01a13c731532469bfde4fd009f54f95242d934 Mon Sep 17 00:00:00 2001 From: Mathias Hauser Date: Mon, 21 Oct 2019 13:43:39 +0200 Subject: [PATCH 2/4] trailing whitespace --- xarray/core/computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 3bc0bd591e4..ff010ad5f43 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1118,7 +1118,7 @@ def dot(*arrays, dims=None, **kwargs): [273, 446, 619]]) Dimensions without coordinates: a, d - >>> xr.dot(da_a, da_b) + >>> xr.dot(da_a, da_b) array([110, 125]) Dimensions without coordinates: c From 7544251b97153dd2466bc31bde01d159fc11da9c Mon Sep 17 00:00:00 2001 From: Mathias Hauser Date: Tue, 29 Oct 2019 18:09:35 +0100 Subject: [PATCH 3/4] move whats new to other ellipsis work --- doc/whats-new.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 09c4160492c..7a97e262ac9 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -34,6 +34,9 @@ New Features to reduce over all dimensions. While we have no plans to remove `xr.ALL_DIMS`, we suggest using `...`. By `Maximilian Roos `_ +- :py:func:`~xarray.dot`, and :py:func:`~xarray.DataArray.dot` now support the + `dims=...` option to sum over the union of dimensions of all input arrays + (:issue:`3423`) by `Mathias Hauser `_. - Added integration tests against `pint `_. (:pull:`3238`) by `Justus Magin `_. @@ -45,9 +48,6 @@ New Features ``pip install git+https://github.com/andrewgsavage/pint.git@refs/pull/6/head)``. Even with it, interaction with non-numpy array libraries, e.g. dask or sparse, is broken. -- :py:func:`~xarray.dot`, and :py:func:`~xarray.DataArray.dot` now support the - ``dims=xarray.ALL_DIMS`` option to sum over the union of dimensions of all - arrays (:issue:`3423`) by `Mathias Hauser `_. - Added new :py:meth:`Dataset._repr_html_` and :py:meth:`DataArray._repr_html_` to improve representation of objects in jupyter. By default this feature is turned off for now. Enable it with :py:meth:`xarray.set_options(display_style="html")`. From 74ef2d15b8fe42cc9d01e9d0778ea9da8a5db794 Mon Sep 17 00:00:00 2001 From: Mathias Hauser Date: Tue, 29 Oct 2019 18:26:24 +0100 Subject: [PATCH 4/4] xr.ALL_DIMS -> Ellipsis --- xarray/core/computation.py | 11 +++++------ xarray/core/dataarray.py | 6 +++--- xarray/tests/test_computation.py | 10 +++++----- xarray/tests/test_dataarray.py | 8 ++++---- 4 files changed, 17 insertions(+), 18 deletions(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index ff010ad5f43..96992bd2a3a 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -25,7 +25,6 @@ from . import duck_array_ops, utils from .alignment import deep_align -from .common import ALL_DIMS from .merge import merge_coordinates_without_align from .pycompat import dask_array_type from .utils import is_dict_like @@ -1056,9 +1055,9 @@ def dot(*arrays, dims=None, **kwargs): ---------- arrays: DataArray (or Variable) objects Arrays to compute. - dims: xarray.ALL_DIMS, str or tuple of strings, optional - Which dimensions to sum over. - If not speciified, then all the common dimensions are summed over. + dims: '...', str or tuple of strings, optional + Which dimensions to sum over. Ellipsis ('...') sums over all dimensions. + If not specified, then all the common dimensions are summed over. **kwargs: dict Additional keyword arguments passed to numpy.einsum or dask.array.einsum @@ -1123,7 +1122,7 @@ def dot(*arrays, dims=None, **kwargs): array([110, 125]) Dimensions without coordinates: c - >>> xr.dot(da_a, da_b, dims=xr.ALL_DIMS) + >>> xr.dot(da_a, da_b, dims=...) array(235) """ @@ -1150,7 +1149,7 @@ def dot(*arrays, dims=None, **kwargs): einsum_axes = "abcdefghijklmnopqrstuvwxyz" dim_map = {d: einsum_axes[i] for i, d in enumerate(all_dims)} - if dims is ALL_DIMS: + if dims is ...: dims = all_dims elif dims is None: # find dimensions that occur more than one times diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 80fa0312928..2e59af57260 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -2742,9 +2742,9 @@ def dot( ---------- other : DataArray The other array with which the dot product is performed. - dims: xarray.ALL_DIMS, hashable or sequence of hashables, optional - Along which dimensions to be summed over. Default all the common - dimensions are summed over. + dims: '...', hashable or sequence of hashables, optional + Which dimensions to sum over. Ellipsis ('...') sums over all dimensions. + If not specified, then all the common dimensions are summed over. Returns ------- diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 7ec0b7656b5..1f2634cc9b0 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -998,20 +998,20 @@ def test_dot(use_dask): assert actual.dims == ("b",) assert (actual.data == np.zeros(actual.shape)).all() - # xr.ALL_DIMS - actual = xr.dot(da_a, da_b, dims=xr.ALL_DIMS) + # Ellipsis (...) sums over all dimensions + actual = xr.dot(da_a, da_b, dims=...) assert actual.dims == () assert (actual.data == np.einsum("ij,ijk->", a, b)).all() - actual = xr.dot(da_a, da_b, da_c, dims=xr.ALL_DIMS) + actual = xr.dot(da_a, da_b, da_c, dims=...) assert actual.dims == () assert (actual.data == np.einsum("ij,ijk,kl-> ", a, b, c)).all() - actual = xr.dot(da_a, dims=xr.ALL_DIMS) + actual = xr.dot(da_a, dims=...) assert actual.dims == () assert (actual.data == np.einsum("ij-> ", a)).all() - actual = xr.dot(da_a.sel(a=[]), da_a.sel(a=[]), dims=xr.ALL_DIMS) + actual = xr.dot(da_a.sel(a=[]), da_a.sel(a=[]), dims=...) assert actual.dims == () assert (actual.data == np.zeros(actual.shape)).all() diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index b442b37844e..38fe7384df2 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -3925,13 +3925,13 @@ def test_dot(self): expected = DataArray(expected_vals, coords=[x, j], dims=["x", "j"]) assert_equal(expected, actual) - # xr.ALL_DIMS: all shared dims - actual = da.dot(da, dims=xr.ALL_DIMS) + # Ellipsis: all dims are shared + actual = da.dot(da, dims=...) expected = da.dot(da) assert_equal(expected, actual) - # xr.ALL_DIMS: multiple shared dims - actual = da.dot(dm, dims=xr.ALL_DIMS) + # Ellipsis: not all dims are shared + actual = da.dot(dm, dims=...) expected = da.dot(dm, dims=("j", "x", "y", "z")) assert_equal(expected, actual)