Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 7551a7a

Browse files
benbovydcherianpre-commit-ci[bot]
authored
Alignment of n-dimensional indexes with partially excluded dims (#10293)
Co-authored-by: Deepak Cherian <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent bb076e3 commit 7551a7a

File tree

7 files changed

+330
-82
lines changed

7 files changed

+330
-82
lines changed

‎doc/whats-new.rst

+5-1
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,11 @@ v2025.05.0 (unreleased)
2424

2525
New Features
2626
~~~~~~~~~~~~
27-
27+
- Allow an Xarray index that uses multiple dimensions checking equality with another
28+
index for only a subset of those dimensions (i.e., ignoring the dimensions
29+
that are excluded from alignment).
30+
(:issue:`10243`, :pull:`10293`)
31+
By `Benoit Bovy <https://github.com/benbovy>`_.
2832

2933
Breaking changes
3034
~~~~~~~~~~~~~~~~

‎xarray/core/coordinate_transform.py

+27-3
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1+
from __future__ import annotations
2+
13
from collections.abc import Hashable, Iterable, Mapping
2-
from typing import Any
4+
from typing import Any, overload
35

46
import numpy as np
57

@@ -64,8 +66,30 @@ def reverse(self, coord_labels: dict[Hashable, Any]) -> dict[str, Any]:
6466
"""
6567
raise NotImplementedError
6668

67-
def equals(self, other: "CoordinateTransform") -> bool:
68-
"""Check equality with another CoordinateTransform of the same kind."""
69+
@overload
70+
def equals(self, other: CoordinateTransform) -> bool: ...
71+
72+
@overload
73+
def equals(
74+
self, other: CoordinateTransform, *, exclude: frozenset[Hashable] | None = None
75+
) -> bool: ...
76+
77+
def equals(self, other: CoordinateTransform, **kwargs) -> bool:
78+
"""Check equality with another CoordinateTransform of the same kind.
79+
80+
Parameters
81+
----------
82+
other : CoordinateTransform
83+
The other Index object to compare with this object.
84+
exclude : frozenset of hashable, optional
85+
Dimensions excluded from checking. It is None by default, (i.e.,
86+
when this method is not called in the context of alignment). For a
87+
n-dimensional transform this option allows a CoordinateTransform to
88+
optionally ignore any dimension in ``exclude`` when comparing
89+
``self`` with ``other``. For a 1-dimensional transform this kwarg
90+
can be safely ignored, as this method is not called when all of the
91+
transform's dimensions are also excluded from alignment.
92+
"""
6993
raise NotImplementedError
7094

7195
def generate_coords(

‎xarray/core/indexes.py

+62-7
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22

33
import collections.abc
44
import copy
5+
import inspect
56
from collections import defaultdict
6-
from collections.abc import Hashable, Iterable, Iterator, Mapping, Sequence
7-
from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
7+
from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
8+
from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, overload
89

910
import numpy as np
1011
import pandas as pd
@@ -348,7 +349,15 @@ def reindex_like(self, other: Self) -> dict[Hashable, Any]:
348349
"""
349350
raise NotImplementedError(f"{self!r} doesn't support re-indexing labels")
350351

351-
def equals(self, other: Index) -> bool:
352+
@overload
353+
def equals(self, other: Index) -> bool: ...
354+
355+
@overload
356+
def equals(
357+
self, other: Index, *, exclude: frozenset[Hashable] | None = None
358+
) -> bool: ...
359+
360+
def equals(self, other: Index, **kwargs) -> bool:
352361
"""Compare this index with another index of the same type.
353362
354363
Implementation is optional but required in order to support alignment.
@@ -357,11 +366,22 @@ def equals(self, other: Index) -> bool:
357366
----------
358367
other : Index
359368
The other Index object to compare with this object.
369+
exclude : frozenset of hashable, optional
370+
Dimensions excluded from checking. It is None by default, (i.e.,
371+
when this method is not called in the context of alignment). For a
372+
n-dimensional index this option allows an Index to optionally ignore
373+
any dimension in ``exclude`` when comparing ``self`` with ``other``.
374+
For a 1-dimensional index this kwarg can be safely ignored, as this
375+
method is not called when all of the index's dimensions are also
376+
excluded from alignment (note: the index's dimensions correspond to
377+
the union of the dimensions of all coordinate variables associated
378+
with this index).
360379
361380
Returns
362381
-------
363382
is_equal : bool
364383
``True`` if the indexes are equal, ``False`` otherwise.
384+
365385
"""
366386
raise NotImplementedError()
367387

@@ -863,7 +883,7 @@ def sel(
863883

864884
return IndexSelResult({self.dim: indexer})
865885

866-
def equals(self, other: Index):
886+
def equals(self, other: Index, *, exclude: frozenset[Hashable] | None = None):
867887
if not isinstance(other, PandasIndex):
868888
return False
869889
return self.index.equals(other.index) and self.dim == other.dim
@@ -1542,10 +1562,12 @@ def sel(
15421562

15431563
return IndexSelResult(results)
15441564

1545-
def equals(self, other: Index) -> bool:
1565+
def equals(
1566+
self, other: Index, *, exclude: frozenset[Hashable] | None = None
1567+
) -> bool:
15461568
if not isinstance(other, CoordinateTransformIndex):
15471569
return False
1548-
return self.transform.equals(other.transform)
1570+
return self.transform.equals(other.transform, exclude=exclude)
15491571

15501572
def rename(
15511573
self,
@@ -1925,6 +1947,36 @@ def default_indexes(
19251947
return indexes
19261948

19271949

1950+
def _wrap_index_equals(
1951+
index: Index,
1952+
) -> Callable[[Index, frozenset[Hashable]], bool]:
1953+
# TODO: remove this Index.equals() wrapper (backward compatibility)
1954+
1955+
sig = inspect.signature(index.equals)
1956+
1957+
if len(sig.parameters) == 1:
1958+
index_cls_name = type(index).__module__ + "." + type(index).__qualname__
1959+
emit_user_level_warning(
1960+
f"the signature ``{index_cls_name}.equals(self, other)`` is deprecated. "
1961+
f"Please update it to "
1962+
f"``{index_cls_name}.equals(self, other, *, exclude=None)`` "
1963+
"or kindly ask the maintainers of ``{index_cls_name}`` to do it. "
1964+
"See documentation of xarray.Index.equals() for more info.",
1965+
FutureWarning,
1966+
)
1967+
exclude_kwarg = False
1968+
else:
1969+
exclude_kwarg = True
1970+
1971+
def equals_wrapper(other: Index, exclude: frozenset[Hashable]) -> bool:
1972+
if exclude_kwarg:
1973+
return index.equals(other, exclude=exclude)
1974+
else:
1975+
return index.equals(other)
1976+
1977+
return equals_wrapper
1978+
1979+
19281980
def indexes_equal(
19291981
index: Index,
19301982
other_index: Index,
@@ -1966,6 +2018,7 @@ def indexes_equal(
19662018

19672019
def indexes_all_equal(
19682020
elements: Sequence[tuple[Index, dict[Hashable, Variable]]],
2021+
exclude_dims: frozenset[Hashable],
19692022
) -> bool:
19702023
"""Check if indexes are all equal.
19712024
@@ -1990,9 +2043,11 @@ def check_variables():
19902043

19912044
same_type = all(type(indexes[0]) is type(other_idx) for other_idx in indexes[1:])
19922045
if same_type:
2046+
index_equals_func = _wrap_index_equals(indexes[0])
19932047
try:
19942048
not_equal = any(
1995-
not indexes[0].equals(other_idx) for other_idx in indexes[1:]
2049+
not index_equals_func(other_idx, exclude_dims)
2050+
for other_idx in indexes[1:]
19962051
)
19972052
except NotImplementedError:
19982053
not_equal = check_variables()

‎xarray/indexes/range_index.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,9 @@ def reverse(self, coord_labels: dict[Hashable, Any]) -> dict[str, Any]:
6565
positions = (labels - self.start) / self.step
6666
return {self.dim: positions}
6767

68-
def equals(self, other: CoordinateTransform) -> bool:
68+
def equals(
69+
self, other: CoordinateTransform, exclude: frozenset[Hashable] | None = None
70+
) -> bool:
6971
if not isinstance(other, RangeCoordinateTransform):
7072
return False
7173

0 commit comments

Comments
 (0)