xarray icon indicating copy to clipboard operation
xarray copied to clipboard

Fix assert_equal with check_dim_order=False for mixed dimension orders

Open max-sixty opened this issue 4 months ago • 0 comments

Summary

  • Fixes #10704
  • assert_equal with check_dim_order=False now works correctly for Datasets containing variables with different dimension orders

The Bug

When comparing Datasets with check_dim_order=False, the comparison would fail if individual variables had different dimension orders, even when comparing a Dataset to itself:

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

# This would fail, but shouldn't:
xr.testing.assert_equal(dataset, dataset, check_dim_order=False)

The Fix

The solution transposes both objects to a canonical dimension order using the intersection of their dimensions. The ellipsis (...) handles any dimensions unique to either object:

common_dims = set(a.dims) & set(b.dims)
if common_dims:
    canonical_order = sorted(common_dims) + [...]
    return a.transpose(*canonical_order), b.transpose(*canonical_order)

This elegant approach works uniformly for Variables, DataArrays, and Datasets without special casing.

Test plan

  • [x] Added test case test_assert_equal_dataset_check_dim_order that reproduces the issue
  • [x] All existing tests pass
  • [x] Verified fix works for both assert_equal and assert_allclose

🤖 Generated with Claude Code

max-sixty avatar Sep 07 '25 17:09 max-sixty