From cfd41abe58a1a5ea145bb7209a5f7c3c5f303390 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Wed, 1 May 2024 14:02:36 +0200 Subject: [PATCH 01/16] Add argument check_dims to assert_allclose to allow transposed inputs --- xarray/testing/assertions.py | 18 +++++++++++++++--- xarray/tests/test_assertions.py | 11 +++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 018874c169e..31797b5a192 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -213,7 +213,9 @@ def assert_identical(a, b, from_root=True): @ensure_warnings -def assert_allclose(a, b, rtol=1e-05, atol=1e-08, decode_bytes=True): +def assert_allclose( + a, b, rtol=1e-05, atol=1e-08, decode_bytes=True, check_dims="strict" +): """Like :py:func:`numpy.testing.assert_allclose`, but for xarray objects. Raises an AssertionError if two objects are not equal up to desired @@ -233,6 +235,10 @@ def assert_allclose(a, b, rtol=1e-05, atol=1e-08, decode_bytes=True): Whether byte dtypes should be decoded to strings as UTF-8 or not. This is useful for testing serialization methods on Python 3 that return saved strings as bytes. + check_dims : {"strict", "transpose"}, optional + To what degree the dimensions must match + * "strict": a and b must have the same dimensions in the same order + * "transpose": a and b must have the same dimensions, not necessarily in the same order See Also -------- @@ -249,8 +255,14 @@ def assert_allclose(a, b, rtol=1e-05, atol=1e-08, decode_bytes=True): def compat_variable(a, b): a = getattr(a, "variable", a) b = getattr(b, "variable", b) - - return a.dims == b.dims and (a._data is b._data or equiv(a.data, b.data)) + if check_dims == "strict": + return a.dims == b.dims and (a._data is b._data or equiv(a.data, b.data)) + elif check_dims == "transpose": + return set(a.dims) == set(b.dims) and ( + a._data is b._data or equiv(a.data, b.transpose(*a.dims).data) + ) + else: + raise ValueError(f"Invalid value for check_dims: {check_dims}") if isinstance(a, Variable): allclose = compat_variable(a, b) diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index 59861ef7981..866b0cf53c8 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -57,6 +57,17 @@ def test_allclose_regression() -> None: def test_assert_allclose(obj1, obj2) -> None: with pytest.raises(AssertionError): xr.testing.assert_allclose(obj1, obj2) + with pytest.raises(AssertionError): + xr.testing.assert_allclose(obj1, obj2, check_dims="transpose") + + +def test_assert_allclose_transpose() -> None: + """Transposed DataArray raises assertion unless check_dims="transpose".""" + obj1 = xr.DataArray([[0, 1, 2], [2, 3, 4]], dims=["a", "b"]) + obj2 = xr.DataArray([[0, 2], [1, 3], [2, 4]], dims=["b", "a"]) + with pytest.raises(AssertionError): + xr.testing.assert_allclose(obj1, obj2) + xr.testing.assert_allclose(obj1, obj2, check_dims="transpose") @pytest.mark.filterwarnings("error") From 49c318c4a7954ae770729d89bedd42c58f5c0615 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Wed, 1 May 2024 14:05:10 +0200 Subject: [PATCH 02/16] Update whats-new.rst --- doc/whats-new.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 8b05b729a31..c56885d6f15 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -29,6 +29,7 @@ New Features for example, will retain the object. However, one cannot do operations that are not possible on the `ExtensionArray` then, such as broadcasting. By `Ilan Gold `_. +- `xarray.testing.assert_allclose` now accepts transposed inputs when `check_dims="transpose"` Breaking changes ~~~~~~~~~~~~~~~~ From c4e39f9479aaf40002b4d33930ab66d4274417e3 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Wed, 1 May 2024 15:22:44 +0200 Subject: [PATCH 03/16] Add `check_dims` argument to assert_equal and assert_identical + tests --- doc/whats-new.rst | 2 +- xarray/testing/assertions.py | 26 +++++++++++++++----------- xarray/tests/test_assertions.py | 16 +++++++++++++--- 3 files changed, 29 insertions(+), 15 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c56885d6f15..26f30514315 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -29,7 +29,7 @@ New Features for example, will retain the object. However, one cannot do operations that are not possible on the `ExtensionArray` then, such as broadcasting. By `Ilan Gold `_. -- `xarray.testing.assert_allclose` now accepts transposed inputs when `check_dims="transpose"` +- `assert_allclose/assert_equal/assert_identical` now accept transposed inputs when `check_dims="transpose"` Breaking changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 31797b5a192..684990dd361 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -95,6 +95,14 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False): raise TypeError(f"{type(a)} not of type DataTree") +def align_dims(a, b, check_dims): + if check_dims not in ("strict", "transpose"): + raise ValueError(f"Invalid value for check_dims: {check_dims}") + if check_dims == "transpose" and isinstance(a, (Variable, DataArray, Dataset)): + return b.transpose(*a.dims) + return b + + @overload def assert_equal(a, b): ... @@ -104,7 +112,7 @@ def assert_equal(a: DataTree, b: DataTree, from_root: bool = True): ... @ensure_warnings -def assert_equal(a, b, from_root=True): +def assert_equal(a, b, from_root=True, check_dims="strict"): """Like :py:func:`numpy.testing.assert_array_equal`, but for xarray objects. @@ -137,6 +145,7 @@ def assert_equal(a, b, from_root=True): assert ( type(a) == type(b) or isinstance(a, Coordinates) and isinstance(b, Coordinates) ) + b = align_dims(a, b, check_dims) if isinstance(a, (Variable, DataArray)): assert a.equals(b), formatting.diff_array_repr(a, b, "equals") elif isinstance(a, Dataset): @@ -162,13 +171,13 @@ def assert_identical(a: DataTree, b: DataTree, from_root: bool = True): ... @ensure_warnings -def assert_identical(a, b, from_root=True): +def assert_identical(a, b, from_root=True, check_dims="strict"): """Like :py:func:`xarray.testing.assert_equal`, but also matches the objects' names and attributes. Raises an AssertionError if two objects are not identical. - For DataTree objects, assert_identical is mapped over all Datasets on each + For Dif ataTree objects, assert_identical is mapped over all Datasets on each node, with the DataTrees being identical if both are isomorphic and the corresponding Datasets at each node are themselves identical. @@ -191,6 +200,7 @@ def assert_identical(a, b, from_root=True): assert ( type(a) == type(b) or isinstance(a, Coordinates) and isinstance(b, Coordinates) ) + b = align_dims(a, b, check_dims) if isinstance(a, Variable): assert a.identical(b), formatting.diff_array_repr(a, b, "identical") elif isinstance(a, DataArray): @@ -246,6 +256,7 @@ def assert_allclose( """ __tracebackhide__ = True assert type(a) == type(b) + b = align_dims(a, b, check_dims) equiv = functools.partial( _data_allclose_or_equiv, rtol=rtol, atol=atol, decode_bytes=decode_bytes @@ -255,14 +266,7 @@ def assert_allclose( def compat_variable(a, b): a = getattr(a, "variable", a) b = getattr(b, "variable", b) - if check_dims == "strict": - return a.dims == b.dims and (a._data is b._data or equiv(a.data, b.data)) - elif check_dims == "transpose": - return set(a.dims) == set(b.dims) and ( - a._data is b._data or equiv(a.data, b.transpose(*a.dims).data) - ) - else: - raise ValueError(f"Invalid value for check_dims: {check_dims}") + return a.dims == b.dims and (a._data is b._data or equiv(a.data, b.data)) if isinstance(a, Variable): allclose = compat_variable(a, b) diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index 866b0cf53c8..f57987ef399 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -61,13 +61,16 @@ def test_assert_allclose(obj1, obj2) -> None: xr.testing.assert_allclose(obj1, obj2, check_dims="transpose") -def test_assert_allclose_transpose() -> None: +@pytest.mark.parametrize( + "func", ["assert_equal", "assert_allclose", "assert_identical"] +) +def test_assert_allclose_equal_identical_transpose(func) -> None: """Transposed DataArray raises assertion unless check_dims="transpose".""" obj1 = xr.DataArray([[0, 1, 2], [2, 3, 4]], dims=["a", "b"]) obj2 = xr.DataArray([[0, 2], [1, 3], [2, 4]], dims=["b", "a"]) with pytest.raises(AssertionError): - xr.testing.assert_allclose(obj1, obj2) - xr.testing.assert_allclose(obj1, obj2, check_dims="transpose") + getattr(xr.testing, func)(obj1, obj2) + getattr(xr.testing, func)(obj1, obj2, check_dims="transpose") @pytest.mark.filterwarnings("error") @@ -138,6 +141,13 @@ def test_assert_duckarray_equal(duckarray, obj1, obj2) -> None: xr.testing.assert_duckarray_equal(a, b) +def test_assert_equal_transpose(duckarray, obj1, obj2) -> None: + a = duckarray(obj1) + b = duckarray(obj2) + + xr.testing.assert_duckarray_equal(a, b) + + @pytest.mark.parametrize( "func", [ From 9ea914887e962d8e1f14f624717705cfd290a934 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Wed, 1 May 2024 15:31:00 +0200 Subject: [PATCH 04/16] Assert that dimensions match before transposing or comparing values --- xarray/testing/assertions.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 684990dd361..9996c7d46d9 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -98,8 +98,12 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False): def align_dims(a, b, check_dims): if check_dims not in ("strict", "transpose"): raise ValueError(f"Invalid value for check_dims: {check_dims}") - if check_dims == "transpose" and isinstance(a, (Variable, DataArray, Dataset)): + if not isinstance(a, (Variable, DataArray, Dataset)): + return b + if check_dims == "transpose": + assert set(a.dims) == set(b.dims), f"Dimensions differ: {a.dims} {b.dims}" return b.transpose(*a.dims) + assert a.dims == b.dims, f"Dimensions differ: {a.dims} {b.dims}" return b From 0ace755f19a71bf19c7147449e046ee10e3a1c05 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Wed, 1 May 2024 15:31:23 +0200 Subject: [PATCH 05/16] Add docstring for check_dims to assert_equal and assert_identical --- xarray/testing/assertions.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 9996c7d46d9..4f3e217d87d 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -139,6 +139,10 @@ def assert_equal(a, b, from_root=True, check_dims="strict"): Only used when comparing DataTree objects. Indicates whether or not to first traverse to the root of the trees before checking for isomorphism. If a & b have no parents then this has no effect. + check_dims : {"strict", "transpose"}, optional + To what degree the dimensions must match + * "strict": a and b must have the same dimensions in the same order + * "transpose": a and b must have the same dimensions, not necessarily in the same order See Also -------- @@ -195,6 +199,10 @@ def assert_identical(a, b, from_root=True, check_dims="strict"): Only used when comparing DataTree objects. Indicates whether or not to first traverse to the root of the trees before checking for isomorphism. If a & b have no parents then this has no effect. + check_dims : {"strict", "transpose"}, optional + To what degree the dimensions must match + * "strict": a and b must have the same dimensions in the same order + * "transpose": a and b must have the same dimensions, not necessarily in the same order See Also -------- From 413b8bf9f7e20fa80c65b9f77a79409a2a98f235 Mon Sep 17 00:00:00 2001 From: ignamv Date: Wed, 1 May 2024 19:01:03 +0200 Subject: [PATCH 06/16] Update doc/whats-new.rst Co-authored-by: Tom Nicholas --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 26f30514315..32a1caa49ab 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -29,7 +29,7 @@ New Features for example, will retain the object. However, one cannot do operations that are not possible on the `ExtensionArray` then, such as broadcasting. By `Ilan Gold `_. -- `assert_allclose/assert_equal/assert_identical` now accept transposed inputs when `check_dims="transpose"` +- :py:func:`testing.assert_allclose`/:py:func:`testing.assert_equal`/:py:func:`testing.assert_identical` now accept a new argument `check_dims="transpose"`, controlling whether a transposed array is considered equal. (:issue:`5733`, :pull:`8991`) Breaking changes ~~~~~~~~~~~~~~~~ From 97ea491d1eced0a36c98d45ad165a1b020f1b672 Mon Sep 17 00:00:00 2001 From: ignamv Date: Thu, 2 May 2024 06:23:32 +0200 Subject: [PATCH 07/16] Undo fat finger Co-authored-by: Tom Nicholas --- xarray/testing/assertions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 4f3e217d87d..6d196d0c558 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -185,7 +185,7 @@ def assert_identical(a, b, from_root=True, check_dims="strict"): Raises an AssertionError if two objects are not identical. - For Dif ataTree objects, assert_identical is mapped over all Datasets on each + For DataTree objects, assert_identical is mapped over all Datasets on each node, with the DataTrees being identical if both are isomorphic and the corresponding Datasets at each node are themselves identical. From ab677c22854f952133cfb6d56ea4dfd5a29f5f18 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Wed, 1 May 2024 19:03:19 +0200 Subject: [PATCH 08/16] Add attribution to whats-new.rst --- doc/whats-new.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 32a1caa49ab..7ce4d6aa539 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -30,6 +30,7 @@ New Features then, such as broadcasting. By `Ilan Gold `_. - :py:func:`testing.assert_allclose`/:py:func:`testing.assert_equal`/:py:func:`testing.assert_identical` now accept a new argument `check_dims="transpose"`, controlling whether a transposed array is considered equal. (:issue:`5733`, :pull:`8991`) + By `Ignacio Martinez Vazquez `_. Breaking changes ~~~~~~~~~~~~~~~~ From de2d105ba71ee0bc6900baf5f316f44f41ea2c7f Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Thu, 2 May 2024 06:27:37 +0200 Subject: [PATCH 09/16] Replace check_dims with bool argument check_dim_order, rename align_dims to maybe_transpose_dims --- xarray/testing/assertions.py | 37 +++++++++++++-------------------- xarray/tests/test_assertions.py | 6 +++--- 2 files changed, 18 insertions(+), 25 deletions(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 6d196d0c558..d398d4b14ca 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -95,12 +95,11 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False): raise TypeError(f"{type(a)} not of type DataTree") -def align_dims(a, b, check_dims): - if check_dims not in ("strict", "transpose"): - raise ValueError(f"Invalid value for check_dims: {check_dims}") +def maybe_transpose_dims(a, b, check_dim_order: bool): + """Helper for assert_equal/allclose/identical""" if not isinstance(a, (Variable, DataArray, Dataset)): return b - if check_dims == "transpose": + if not check_dim_order: assert set(a.dims) == set(b.dims), f"Dimensions differ: {a.dims} {b.dims}" return b.transpose(*a.dims) assert a.dims == b.dims, f"Dimensions differ: {a.dims} {b.dims}" @@ -116,7 +115,7 @@ def assert_equal(a: DataTree, b: DataTree, from_root: bool = True): ... @ensure_warnings -def assert_equal(a, b, from_root=True, check_dims="strict"): +def assert_equal(a, b, from_root=True, check_dim_order: bool = True): """Like :py:func:`numpy.testing.assert_array_equal`, but for xarray objects. @@ -139,10 +138,8 @@ def assert_equal(a, b, from_root=True, check_dims="strict"): Only used when comparing DataTree objects. Indicates whether or not to first traverse to the root of the trees before checking for isomorphism. If a & b have no parents then this has no effect. - check_dims : {"strict", "transpose"}, optional - To what degree the dimensions must match - * "strict": a and b must have the same dimensions in the same order - * "transpose": a and b must have the same dimensions, not necessarily in the same order + check_dim_order : bool, optional, default is True + Whether dimensions must be in the same order. See Also -------- @@ -153,7 +150,7 @@ def assert_equal(a, b, from_root=True, check_dims="strict"): assert ( type(a) == type(b) or isinstance(a, Coordinates) and isinstance(b, Coordinates) ) - b = align_dims(a, b, check_dims) + b = maybe_transpose_dims(a, b, check_dim_order) if isinstance(a, (Variable, DataArray)): assert a.equals(b), formatting.diff_array_repr(a, b, "equals") elif isinstance(a, Dataset): @@ -179,7 +176,7 @@ def assert_identical(a: DataTree, b: DataTree, from_root: bool = True): ... @ensure_warnings -def assert_identical(a, b, from_root=True, check_dims="strict"): +def assert_identical(a, b, from_root=True, check_dim_order: bool = True): """Like :py:func:`xarray.testing.assert_equal`, but also matches the objects' names and attributes. @@ -199,10 +196,8 @@ def assert_identical(a, b, from_root=True, check_dims="strict"): Only used when comparing DataTree objects. Indicates whether or not to first traverse to the root of the trees before checking for isomorphism. If a & b have no parents then this has no effect. - check_dims : {"strict", "transpose"}, optional - To what degree the dimensions must match - * "strict": a and b must have the same dimensions in the same order - * "transpose": a and b must have the same dimensions, not necessarily in the same order + check_dim_order : bool, optional, default is True + Whether dimensions must be in the same order. See Also -------- @@ -212,7 +207,7 @@ def assert_identical(a, b, from_root=True, check_dims="strict"): assert ( type(a) == type(b) or isinstance(a, Coordinates) and isinstance(b, Coordinates) ) - b = align_dims(a, b, check_dims) + b = maybe_transpose_dims(a, b, check_dim_order) if isinstance(a, Variable): assert a.identical(b), formatting.diff_array_repr(a, b, "identical") elif isinstance(a, DataArray): @@ -236,7 +231,7 @@ def assert_identical(a, b, from_root=True, check_dims="strict"): @ensure_warnings def assert_allclose( - a, b, rtol=1e-05, atol=1e-08, decode_bytes=True, check_dims="strict" + a, b, rtol=1e-05, atol=1e-08, decode_bytes=True, check_dim_order: bool = True ): """Like :py:func:`numpy.testing.assert_allclose`, but for xarray objects. @@ -257,10 +252,8 @@ def assert_allclose( Whether byte dtypes should be decoded to strings as UTF-8 or not. This is useful for testing serialization methods on Python 3 that return saved strings as bytes. - check_dims : {"strict", "transpose"}, optional - To what degree the dimensions must match - * "strict": a and b must have the same dimensions in the same order - * "transpose": a and b must have the same dimensions, not necessarily in the same order + check_dim_order : bool, optional, default is True + Whether dimensions must be in the same order. See Also -------- @@ -268,7 +261,7 @@ def assert_allclose( """ __tracebackhide__ = True assert type(a) == type(b) - b = align_dims(a, b, check_dims) + b = maybe_transpose_dims(a, b, check_dim_order) equiv = functools.partial( _data_allclose_or_equiv, rtol=rtol, atol=atol, decode_bytes=decode_bytes diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index f57987ef399..f5df50c6e7c 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -58,19 +58,19 @@ def test_assert_allclose(obj1, obj2) -> None: with pytest.raises(AssertionError): xr.testing.assert_allclose(obj1, obj2) with pytest.raises(AssertionError): - xr.testing.assert_allclose(obj1, obj2, check_dims="transpose") + xr.testing.assert_allclose(obj1, obj2, check_dim_order=False) @pytest.mark.parametrize( "func", ["assert_equal", "assert_allclose", "assert_identical"] ) def test_assert_allclose_equal_identical_transpose(func) -> None: - """Transposed DataArray raises assertion unless check_dims="transpose".""" + """Transposed DataArray raises assertion unless check_dim_order=False.""" obj1 = xr.DataArray([[0, 1, 2], [2, 3, 4]], dims=["a", "b"]) obj2 = xr.DataArray([[0, 2], [1, 3], [2, 4]], dims=["b", "a"]) with pytest.raises(AssertionError): getattr(xr.testing, func)(obj1, obj2) - getattr(xr.testing, func)(obj1, obj2, check_dims="transpose") + getattr(xr.testing, func)(obj1, obj2, check_dim_order=False) @pytest.mark.filterwarnings("error") From 35c1838f2c7e423827b92ec19b016dbaf300a44a Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Sat, 4 May 2024 16:33:50 +0200 Subject: [PATCH 10/16] Remove left-over half-made test --- xarray/tests/test_assertions.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index f5df50c6e7c..530bb145ef6 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -141,13 +141,6 @@ def test_assert_duckarray_equal(duckarray, obj1, obj2) -> None: xr.testing.assert_duckarray_equal(a, b) -def test_assert_equal_transpose(duckarray, obj1, obj2) -> None: - a = duckarray(obj1) - b = duckarray(obj2) - - xr.testing.assert_duckarray_equal(a, b) - - @pytest.mark.parametrize( "func", [ From 18d18a9ee4a66ef0dcb16ee3f0635adb9bc50561 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Sun, 5 May 2024 11:41:15 +0200 Subject: [PATCH 11/16] Remove check_dim_order argument from assert_identical --- xarray/testing/assertions.py | 3 +-- xarray/tests/test_assertions.py | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index d398d4b14ca..2087fd840f5 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -176,7 +176,7 @@ def assert_identical(a: DataTree, b: DataTree, from_root: bool = True): ... @ensure_warnings -def assert_identical(a, b, from_root=True, check_dim_order: bool = True): +def assert_identical(a, b, from_root=True): """Like :py:func:`xarray.testing.assert_equal`, but also matches the objects' names and attributes. @@ -207,7 +207,6 @@ def assert_identical(a, b, from_root=True, check_dim_order: bool = True): assert ( type(a) == type(b) or isinstance(a, Coordinates) and isinstance(b, Coordinates) ) - b = maybe_transpose_dims(a, b, check_dim_order) if isinstance(a, Variable): assert a.identical(b), formatting.diff_array_repr(a, b, "identical") elif isinstance(a, DataArray): diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index 530bb145ef6..75947dfccda 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -61,9 +61,7 @@ def test_assert_allclose(obj1, obj2) -> None: xr.testing.assert_allclose(obj1, obj2, check_dim_order=False) -@pytest.mark.parametrize( - "func", ["assert_equal", "assert_allclose", "assert_identical"] -) +@pytest.mark.parametrize("func", ["assert_equal", "assert_allclose"]) def test_assert_allclose_equal_identical_transpose(func) -> None: """Transposed DataArray raises assertion unless check_dim_order=False.""" obj1 = xr.DataArray([[0, 1, 2], [2, 3, 4]], dims=["a", "b"]) From 3f14b21408f562ed6fb6202970f159d222275afe Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Sun, 5 May 2024 11:42:36 +0200 Subject: [PATCH 12/16] assert_allclose/equal: emit full diff if dimensions don't match --- xarray/testing/assertions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 2087fd840f5..fd39d88e4dd 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -99,10 +99,10 @@ def maybe_transpose_dims(a, b, check_dim_order: bool): """Helper for assert_equal/allclose/identical""" if not isinstance(a, (Variable, DataArray, Dataset)): return b - if not check_dim_order: - assert set(a.dims) == set(b.dims), f"Dimensions differ: {a.dims} {b.dims}" + if not check_dim_order and set(a.dims) == set(b.dims): + # Ensure transpose won't fail if a dimension is missing + # If this is the case, the difference will be caught by the caller return b.transpose(*a.dims) - assert a.dims == b.dims, f"Dimensions differ: {a.dims} {b.dims}" return b From d558165263c8ed1a5eb81933c36691630ff4feea Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Sun, 5 May 2024 11:45:01 +0200 Subject: [PATCH 13/16] Rename check_dim_order test, test Dataset with different dim orders --- xarray/tests/test_assertions.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index 75947dfccda..d01f3c27b4a 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -62,13 +62,20 @@ def test_assert_allclose(obj1, obj2) -> None: @pytest.mark.parametrize("func", ["assert_equal", "assert_allclose"]) -def test_assert_allclose_equal_identical_transpose(func) -> None: +def test_assert_allclose_equal_transpose(func) -> None: """Transposed DataArray raises assertion unless check_dim_order=False.""" obj1 = xr.DataArray([[0, 1, 2], [2, 3, 4]], dims=["a", "b"]) obj2 = xr.DataArray([[0, 2], [1, 3], [2, 4]], dims=["b", "a"]) with pytest.raises(AssertionError): getattr(xr.testing, func)(obj1, obj2) getattr(xr.testing, func)(obj1, obj2, check_dim_order=False) + ds1 = obj1.to_dataset(name="varname") + ds1["var2"] = obj1 + ds2 = obj1.to_dataset(name="varname") + ds2["var2"] = obj1.transpose() + with pytest.raises(AssertionError): + getattr(xr.testing, func)(ds1, ds2) + getattr(xr.testing, func)(ds1, ds2, check_dim_order=False) @pytest.mark.filterwarnings("error") From bfa5e92aafee810d86768492c664183fbf61106a Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Sun, 5 May 2024 11:47:13 +0200 Subject: [PATCH 14/16] Update whats-new.rst --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 7ce4d6aa539..922850eb2bc 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -29,7 +29,7 @@ New Features for example, will retain the object. However, one cannot do operations that are not possible on the `ExtensionArray` then, such as broadcasting. By `Ilan Gold `_. -- :py:func:`testing.assert_allclose`/:py:func:`testing.assert_equal`/:py:func:`testing.assert_identical` now accept a new argument `check_dims="transpose"`, controlling whether a transposed array is considered equal. (:issue:`5733`, :pull:`8991`) +- :py:func:`testing.assert_allclose`/:py:func:`testing.assert_equal` now accept a new argument `check_dims="transpose"`, controlling whether a transposed array is considered equal. (:issue:`5733`, :pull:`8991`) By `Ignacio Martinez Vazquez `_. Breaking changes From 555a62667a62187236bd1e835a0027a23b4f9f16 Mon Sep 17 00:00:00 2001 From: ignamv Date: Sun, 5 May 2024 12:03:43 +0200 Subject: [PATCH 15/16] Hide maybe_transpose_dims from Pytest traceback Co-authored-by: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> --- xarray/testing/assertions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index fd39d88e4dd..e102f3d3329 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -97,6 +97,7 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False): def maybe_transpose_dims(a, b, check_dim_order: bool): """Helper for assert_equal/allclose/identical""" + __tracebackhide__ = True if not isinstance(a, (Variable, DataArray, Dataset)): return b if not check_dim_order and set(a.dims) == set(b.dims): From eb66fd2eb5654839351b91d9961435ce5c21cde1 Mon Sep 17 00:00:00 2001 From: Ignacio Martinez Vazquez Date: Sun, 5 May 2024 14:05:20 +0200 Subject: [PATCH 16/16] Ignore mypy error due to missing functools.partial.__name__ --- xarray/testing/assertions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index e102f3d3329..69885868f83 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -266,7 +266,7 @@ def assert_allclose( equiv = functools.partial( _data_allclose_or_equiv, rtol=rtol, atol=atol, decode_bytes=decode_bytes ) - equiv.__name__ = "allclose" + equiv.__name__ = "allclose" # type: ignore[attr-defined] def compat_variable(a, b): a = getattr(a, "variable", a)