From 6f90b7c3549ce4d9a96daa6d69f7416507e48b57 Mon Sep 17 00:00:00 2001 From: Matthew Willson Date: Mon, 24 Nov 2025 12:42:50 +0000 Subject: [PATCH 1/5] Add an arithmetic_compat option to xr.set_options, which determines how non-index coordinates of the same name are compared for potential conflicts when performing binary operations. The default of compat='minimal' matches the previous behaviour. --- doc/whats-new.rst | 4 ++++ xarray/core/coordinates.py | 30 ++++++++++++++++++++------- xarray/core/dataarray.py | 8 +++++-- xarray/core/dataset.py | 2 +- xarray/core/options.py | 31 ++++++++++++++++++++++----- xarray/structure/merge.py | 3 ++- xarray/tests/test_dataarray.py | 38 ++++++++++++++++++++++++++++++++++ xarray/tests/test_dataset.py | 35 +++++++++++++++++++++++++++++++ 8 files changed, 135 insertions(+), 16 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 677b2194a55..6db173a9eee 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,6 +22,10 @@ New Features - :py:func:`combine_nested` now support :py:class:`DataTree` objects (:pull:`10849`). By `Stephan Hoyer `_. +- :py:func:`set_options` now supports an ``arithmetic_compat`` option which determines how non-index coordinates + of the same name are compared for potential conflicts when performing binary operations. The default for it is + ``arithmetic_compat='minimal'`` which matches the existing behaviour. + By `Matthew Willson `_. Breaking Changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index ab5bf7408f3..dd0fb68d803 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -21,7 +21,14 @@ assert_no_index_corrupted, create_default_index_implicit, ) -from xarray.core.types import DataVars, ErrorOptions, Self, T_DataArray, T_Xarray +from xarray.core.types import ( + CompatOptions, + DataVars, + ErrorOptions, + Self, + T_DataArray, + T_Xarray, +) from xarray.core.utils import ( Frozen, ReprObject, @@ -31,6 +38,7 @@ from xarray.core.variable import Variable, as_variable, calculate_dimensions from xarray.structure.alignment import Aligner from xarray.structure.merge import merge_coordinates_without_align, merge_coords +from xarray.util.deprecation_helpers import CombineKwargDefault if TYPE_CHECKING: from xarray.core.common import DataWithCoords @@ -499,18 +507,20 @@ def _drop_coords(self, coord_names): # redirect to DatasetCoordinates._drop_coords self._data.coords._drop_coords(coord_names) - def _merge_raw(self, other, reflexive): + def _merge_raw(self, other, reflexive, compat: CompatOptions | CombineKwargDefault): """For use with binary arithmetic.""" if other is None: variables = dict(self.variables) indexes = dict(self.xindexes) else: coord_list = [self, other] if not reflexive else [other, self] - variables, indexes = merge_coordinates_without_align(coord_list) + variables, indexes = merge_coordinates_without_align( + coord_list, compat=compat + ) return variables, indexes @contextmanager - def _merge_inplace(self, other): + def _merge_inplace(self, other, compat: CompatOptions | CombineKwargDefault): """For use with in-place binary arithmetic.""" if other is None: yield @@ -523,12 +533,16 @@ def _merge_inplace(self, other): if k not in self.xindexes } variables, indexes = merge_coordinates_without_align( - [self, other], prioritized + [self, other], prioritized, compat=compat ) yield self._update_coords(variables, indexes) - def merge(self, other: Mapping[Any, Any] | None) -> Dataset: + def merge( + self, + other: Mapping[Any, Any] | None, + compat: CompatOptions | CombineKwargDefault = "minimal", + ) -> Dataset: """Merge two sets of coordinates to create a new Dataset The method implements the logic used for joining coordinates in the @@ -545,6 +559,8 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset: other : dict-like, optional A :py:class:`Coordinates` object or any mapping that can be turned into coordinates. + compat : {"identical", "equals", "broadcast_equals", "no_conflicts", "override", "minimal"}, default: "minimal" + Compatibility checks to use between coordinate variables. Returns ------- @@ -559,7 +575,7 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset: if not isinstance(other, Coordinates): other = Dataset(coords=other).coords - coords, indexes = merge_coordinates_without_align([self, other]) + coords, indexes = merge_coordinates_without_align([self, other], compat=compat) coord_names = set(coords) return Dataset._construct_direct( variables=coords, coord_names=coord_names, indexes=indexes diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 6c8d0617038..eb17d06b82b 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4899,7 +4899,9 @@ def _binary_op( if not reflexive else f(other_variable_or_arraylike, self.variable) ) - coords, indexes = self.coords._merge_raw(other_coords, reflexive) + coords, indexes = self.coords._merge_raw( + other_coords, reflexive, compat=OPTIONS["arithmetic_compat"] + ) name = result_name([self, other]) return self._replace(variable, coords, name, indexes=indexes) @@ -4919,7 +4921,9 @@ def _inplace_binary_op(self, other: DaCompatible, f: Callable) -> Self: other_coords = getattr(other, "coords", None) other_variable = getattr(other, "variable", other) try: - with self.coords._merge_inplace(other_coords): + with self.coords._merge_inplace( + other_coords, compat=OPTIONS["arithmetic_compat"] + ): f(self.variable, other_variable) except MergeError as exc: raise MergeError( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 9c2c2f60db1..e3282049ca7 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7765,7 +7765,7 @@ def apply_over_both(lhs_data_vars, rhs_data_vars, lhs_vars, rhs_vars): return type(self)(new_data_vars) other_coords: Coordinates | None = getattr(other, "coords", None) - ds = self.coords.merge(other_coords) + ds = self.coords.merge(other_coords, compat=OPTIONS["arithmetic_compat"]) if isinstance(other, Dataset): new_vars = apply_over_both( diff --git a/xarray/core/options.py b/xarray/core/options.py index 451070ce7b4..016365b098b 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -2,14 +2,16 @@ import warnings from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Literal, TypedDict +from typing import TYPE_CHECKING, Any, Literal, TypedDict, get_args +from xarray.core.types import CompatOptions from xarray.core.utils import FrozenDict if TYPE_CHECKING: from matplotlib.colors import Colormap Options = Literal[ + "arithmetic_compat", "arithmetic_join", "chunk_manager", "cmap_divergent", @@ -40,6 +42,7 @@ class T_Options(TypedDict): arithmetic_broadcast: bool + arithmetic_compat: CompatOptions arithmetic_join: Literal["inner", "outer", "left", "right", "exact"] chunk_manager: str cmap_divergent: str | Colormap @@ -70,6 +73,7 @@ class T_Options(TypedDict): OPTIONS: T_Options = { "arithmetic_broadcast": True, + "arithmetic_compat": "minimal", "arithmetic_join": "inner", "chunk_manager": "dask", "cmap_divergent": "RdBu_r", @@ -109,6 +113,7 @@ def _positive_integer(value: Any) -> bool: _VALIDATORS = { "arithmetic_broadcast": lambda value: isinstance(value, bool), + "arithmetic_compat": get_args(CompatOptions).__contains__, "arithmetic_join": _JOIN_OPTIONS.__contains__, "display_max_children": _positive_integer, "display_max_rows": _positive_integer, @@ -178,8 +183,27 @@ class set_options: Parameters ---------- + arithmetic_broadcast: bool, default: True + Whether to perform automatic broadcasting in binary operations. + arithmetic_compat: {"identical", "equals", "broadcast_equals", "no_conflicts", "override", "minimal"}, default: "minimal" + How to compare non-index coordinates of the same name for potential + conflicts when performing binary operations. (For the alignment of index + coordinates in binary operations, see `arithmetic_join`.) + + - "identical": all values, dimensions and attributes of the coordinates + must be the same. + - "equals": all values and dimensions of the coordinates must be the + same. + - "broadcast_equals": all values of the coordinates must be equal after + broadcasting to ensure common dimensions. + - "no_conflicts": only values which are not null in both coordinates + must be equal. The returned coordinate then contains the combination + of all non-null values. + - "override": skip comparing and take the coordinates from the first + operand. + - "minimal": drop conflicting coordinates. arithmetic_join : {"inner", "outer", "left", "right", "exact"}, default: "inner" - DataArray/Dataset alignment in binary operations: + DataArray/Dataset index alignment in binary operations: - "outer": use the union of object indexes - "inner": use the intersection of object indexes @@ -187,9 +211,6 @@ class set_options: - "right": use indexes from the last object with each dimension - "exact": instead of aligning, raise `ValueError` when indexes to be aligned are not equal - - "override": if indexes are of same size, rewrite indexes to be - those of the first object with that dimension. Indexes for the same - dimension must have the same size in all objects. chunk_manager : str, default: "dask" Chunk manager to use for chunked array computations when multiple options are installed. diff --git a/xarray/structure/merge.py b/xarray/structure/merge.py index e5f3c0959bd..e44719ce3cd 100644 --- a/xarray/structure/merge.py +++ b/xarray/structure/merge.py @@ -433,6 +433,7 @@ def merge_coordinates_without_align( prioritized: Mapping[Any, MergeElement] | None = None, exclude_dims: AbstractSet = frozenset(), combine_attrs: CombineAttrsOptions = "override", + compat: CompatOptions | CombineKwargDefault = "minimal", ) -> tuple[dict[Hashable, Variable], dict[Hashable, Index]]: """Merge variables/indexes from coordinates without automatic alignments. @@ -457,7 +458,7 @@ def merge_coordinates_without_align( # TODO: indexes should probably be filtered in collected elements # before merging them merged_coords, merged_indexes = merge_collected( - filtered, prioritized, combine_attrs=combine_attrs + filtered, prioritized, compat=compat, combine_attrs=combine_attrs ) merged_indexes = filter_indexes_from_coords(merged_indexes, set(merged_coords)) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 5eec7b8a2fd..707ca973452 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -25,6 +25,7 @@ DataArray, Dataset, IndexVariable, + MergeError, Variable, align, broadcast, @@ -2516,6 +2517,43 @@ def test_math_with_coords(self) -> None: actual = alt + orig assert_identical(expected, actual) + def test_math_with_arithmetic_compat_options(self) -> None: + # Setting up a clash of non-index coordinate 'foo': + a = xr.DataArray( + data=[0, 0, 0], + dims=["x"], + coords={ + "x": [1, 2, 3], + "foo": (["x"], [1.0, 2.0, np.nan]), + }, + ) + b = xr.DataArray( + data=[0, 0, 0], + dims=["x"], + coords={ + "x": [1, 2, 3], + "foo": (["x"], [np.nan, 2.0, 3.0]), + }, + ) + + with xr.set_options(arithmetic_compat="minimal"): + assert_equal(a + b, a.drop_vars("foo")) + + with xr.set_options(arithmetic_compat="override"): + assert_equal(a + b, a) + assert_equal(b + a, b) + + with xr.set_options(arithmetic_compat="no_conflicts"): + expected = a.assign_coords(foo=(["x"], [1.0, 2.0, 3.0])) + assert_equal(a + b, expected) + assert_equal(b + a, expected) + + with xr.set_options(arithmetic_compat="equals"): + with pytest.raises(MergeError): + a + b + with pytest.raises(MergeError): + b + a + def test_index_math(self) -> None: orig = DataArray(range(3), dims="x", name="x") actual = orig + 1 diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index e677430dfbf..722fbb7a86c 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6834,6 +6834,41 @@ def test_binary_op_join_setting(self) -> None: actual = ds1 + ds2 assert_equal(actual, expected) + def test_binary_op_compat_setting(self) -> None: + # Setting up a clash of non-index coordinate 'foo': + a = xr.Dataset( + data_vars={"var": (["x"], [0, 0, 0])}, + coords={ + "x": [1, 2, 3], + "foo": (["x"], [1.0, 2.0, np.nan]), + }, + ) + b = xr.Dataset( + data_vars={"var": (["x"], [0, 0, 0])}, + coords={ + "x": [1, 2, 3], + "foo": (["x"], [np.nan, 2.0, 3.0]), + }, + ) + + with xr.set_options(arithmetic_compat="minimal"): + assert_equal(a + b, a.drop_vars("foo")) + + with xr.set_options(arithmetic_compat="override"): + assert_equal(a + b, a) + assert_equal(b + a, b) + + with xr.set_options(arithmetic_compat="no_conflicts"): + expected = a.assign_coords(foo=(["x"], [1.0, 2.0, 3.0])) + assert_equal(a + b, expected) + assert_equal(b + a, expected) + + with xr.set_options(arithmetic_compat="equals"): + with pytest.raises(MergeError): + a + b + with pytest.raises(MergeError): + b + a + @pytest.mark.parametrize( ["keep_attrs", "expected"], ( From f194cd8f3b3e90d2f669abff50aa6fdb0d463b32 Mon Sep 17 00:00:00 2001 From: Matthew Date: Fri, 28 Nov 2025 20:04:41 +0000 Subject: [PATCH 2/5] Make compat a kwarg-only argument to Coordinates.merge. Co-authored-by: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> --- xarray/core/coordinates.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index dd0fb68d803..c054b59a299 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -541,6 +541,7 @@ def _merge_inplace(self, other, compat: CompatOptions | CombineKwargDefault): def merge( self, other: Mapping[Any, Any] | None, + *, compat: CompatOptions | CombineKwargDefault = "minimal", ) -> Dataset: """Merge two sets of coordinates to create a new Dataset From 03b59cd3849488d12c837bcdbc65f1ae821b88f6 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 7 Dec 2025 15:55:14 +0100 Subject: [PATCH 3/5] proper numpydocs formatting --- xarray/core/options.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/options.py b/xarray/core/options.py index 016365b098b..9af6117d237 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -183,7 +183,7 @@ class set_options: Parameters ---------- - arithmetic_broadcast: bool, default: True + arithmetic_broadcast : bool, default: True Whether to perform automatic broadcasting in binary operations. arithmetic_compat: {"identical", "equals", "broadcast_equals", "no_conflicts", "override", "minimal"}, default: "minimal" How to compare non-index coordinates of the same name for potential From 6c80268c1c15b451ae64d44af3ce4a8f20a0c102 Mon Sep 17 00:00:00 2001 From: Matthew Willson Date: Mon, 24 Nov 2025 12:42:50 +0000 Subject: [PATCH 4/5] Add test for DataTree --- xarray/tests/test_datatree.py | 37 +++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 0cd888f5782..59db3446e3d 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -2423,6 +2423,43 @@ def test_arithmetic_inherited_coords(self) -> None: expected["/foo/bar"].data = np.array([8, 10, 12]) assert_identical(actual, expected) + def test_binary_op_compat_setting(self) -> None: + # Setting up a clash of non-index coordinate 'foo': + a = DataTree(xr.Dataset( + data_vars={"var": (["x"], [0, 0, 0])}, + coords={ + "x": [1, 2, 3], + "foo": (["x"], [1.0, 2.0, np.nan]), + }, + )) + b = DataTree(xr.Dataset( + data_vars={"var": (["x"], [0, 0, 0])}, + coords={ + "x": [1, 2, 3], + "foo": (["x"], [np.nan, 2.0, 3.0]), + }, + )) + + with xr.set_options(arithmetic_compat="minimal"): + expected = DataTree(a.dataset.drop_vars('foo')) + assert_equal(a + b, expected) + + with xr.set_options(arithmetic_compat="override"): + assert_equal(a + b, a) + assert_equal(b + a, b) + + with xr.set_options(arithmetic_compat="no_conflicts"): + expected = DataTree(a.dataset.assign_coords( + foo=(["x"], [1.0, 2.0, 3.0]))) + assert_equal(a + b, expected) + assert_equal(b + a, expected) + + with xr.set_options(arithmetic_compat="equals"): + with pytest.raises(xr.MergeError): + a + b + with pytest.raises(xr.MergeError): + b + a + def test_binary_op_commutativity_with_dataset(self) -> None: # regression test for #9365 From 51cc01f1b5e3589a7b816fab93c6d03ebb9d5c3c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 8 Dec 2025 13:43:43 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_datatree.py | 37 +++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 59db3446e3d..cb9e81da97a 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -2425,23 +2425,27 @@ def test_arithmetic_inherited_coords(self) -> None: def test_binary_op_compat_setting(self) -> None: # Setting up a clash of non-index coordinate 'foo': - a = DataTree(xr.Dataset( - data_vars={"var": (["x"], [0, 0, 0])}, - coords={ - "x": [1, 2, 3], - "foo": (["x"], [1.0, 2.0, np.nan]), - }, - )) - b = DataTree(xr.Dataset( - data_vars={"var": (["x"], [0, 0, 0])}, - coords={ - "x": [1, 2, 3], - "foo": (["x"], [np.nan, 2.0, 3.0]), - }, - )) + a = DataTree( + xr.Dataset( + data_vars={"var": (["x"], [0, 0, 0])}, + coords={ + "x": [1, 2, 3], + "foo": (["x"], [1.0, 2.0, np.nan]), + }, + ) + ) + b = DataTree( + xr.Dataset( + data_vars={"var": (["x"], [0, 0, 0])}, + coords={ + "x": [1, 2, 3], + "foo": (["x"], [np.nan, 2.0, 3.0]), + }, + ) + ) with xr.set_options(arithmetic_compat="minimal"): - expected = DataTree(a.dataset.drop_vars('foo')) + expected = DataTree(a.dataset.drop_vars("foo")) assert_equal(a + b, expected) with xr.set_options(arithmetic_compat="override"): @@ -2449,8 +2453,7 @@ def test_binary_op_compat_setting(self) -> None: assert_equal(b + a, b) with xr.set_options(arithmetic_compat="no_conflicts"): - expected = DataTree(a.dataset.assign_coords( - foo=(["x"], [1.0, 2.0, 3.0]))) + expected = DataTree(a.dataset.assign_coords(foo=(["x"], [1.0, 2.0, 3.0]))) assert_equal(a + b, expected) assert_equal(b + a, expected)