Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 18 additions & 11 deletions xarray/testing/assertions.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,24 +85,31 @@ def assert_isomorphic(a: DataTree, b: DataTree):


def maybe_transpose_dims(a, b, check_dim_order: bool):
"""Helper for assert_equal/allclose/identical"""
"""Helper for assert_equal/allclose/identical

Returns (a, b) tuple with dimensions transposed to canonical order if needed.
"""

__tracebackhide__ = True

def _maybe_transpose_dims(a, b):
if not isinstance(a, Variable | DataArray | Dataset):
return b
if 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)
return b
return a, b

# Find common dimensions and transpose both to canonical order
common_dims = set(a.dims) & set(b.dims)
if common_dims:
# Use sorted order as canonical, with ellipsis for any unique dims
canonical_order = sorted(common_dims) + [...]
return a.transpose(*canonical_order), b.transpose(*canonical_order)
return a, b

if check_dim_order:
return b
return a, b

if isinstance(a, DataTree):
return map_over_datasets(_maybe_transpose_dims, a, b)
# DataTree case needs special handling - only transpose b
return a, map_over_datasets(lambda a, b: _maybe_transpose_dims(a, b)[1], a, b)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this is correct.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok thanks, sorry for outsourcing finding this. we should find a better balance of using AI agents; I do worry about this sort of OSS interaction

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for example — if an AI agent writes the code, it needs to have 100% test coverage before a human is asked to review it


return _maybe_transpose_dims(a, b)

Expand Down Expand Up @@ -139,7 +146,7 @@ def assert_equal(a, b, check_dim_order: bool = True):
assert type(a) is type(b) or (
isinstance(a, Coordinates) and isinstance(b, Coordinates)
)
b = maybe_transpose_dims(a, b, check_dim_order)
a, 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):
Expand Down Expand Up @@ -227,7 +234,7 @@ def assert_allclose(
"""
__tracebackhide__ = True
assert type(a) is type(b)
b = maybe_transpose_dims(a, b, check_dim_order)
a, b = maybe_transpose_dims(a, b, check_dim_order)

equiv = functools.partial(
_data_allclose_or_equiv, rtol=rtol, atol=atol, decode_bytes=decode_bytes
Expand Down
34 changes: 34 additions & 0 deletions xarray/tests/test_assertions.py
Original file line number Diff line number Diff line change
Expand Up @@ -220,3 +220,37 @@ def __array__(
getattr(xr.testing, func)(a, b)

assert len(w) == 0


def test_assert_equal_dataset_check_dim_order():
"""Test for issue #10704 - check_dim_order=False with Datasets containing mixed dimension orders."""
import numpy as np

# Dataset with variables having different dimension orders
dataset_1 = xr.Dataset(
{
"foo": xr.DataArray(np.zeros([4, 5]), dims=("a", "b")),
"bar": xr.DataArray(np.ones([5, 4]), dims=("b", "a")),
}
)

dataset_2 = xr.Dataset(
{
"foo": xr.DataArray(np.zeros([5, 4]), dims=("b", "a")),
"bar": xr.DataArray(np.ones([4, 5]), dims=("a", "b")),
}
)

# These should be equal when ignoring dimension order
xr.testing.assert_equal(dataset_1, dataset_2, check_dim_order=False)
xr.testing.assert_allclose(dataset_1, dataset_2, check_dim_order=False)

# Should also work when comparing dataset to itself
xr.testing.assert_equal(dataset_1, dataset_1, check_dim_order=False)
xr.testing.assert_allclose(dataset_1, dataset_1, check_dim_order=False)

# But should fail with check_dim_order=True
with pytest.raises(AssertionError):
xr.testing.assert_equal(dataset_1, dataset_2, check_dim_order=True)
with pytest.raises(AssertionError):
xr.testing.assert_allclose(dataset_1, dataset_2, check_dim_order=True)
Loading