From 5c99707eff20bccf74d666ee7de30c33bfd6f72a Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Sat, 5 Sep 2026 12:19:11 +0800 Subject: [PATCH 01/12] Add a couple of extra guards. Think the only issue now is the missing dims are missing in the test data --- pixi.toml | 3 +++ xarray/computation/weighted.py | 17 ++++++++++++++--- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/pixi.toml b/pixi.toml index 90fd51a975a..0fa6595b300 100644 --- a/pixi.toml +++ b/pixi.toml @@ -477,3 +477,6 @@ build-package = { features = [ "release", ], no-default-feature = true } test-nightly = { features = ["nightly"], no-default-feature = true } + +[pypi-dependencies] +pdbpp = ">=0.12.1, <0.13" diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index b311290aabf..7f6471b249c 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -201,12 +201,23 @@ def _check_dim(self, dim: Dims): else: dims = list(dim) all_dims = set(self.obj.dims).union(set(self.weights.dims)) - missing_dims = set(dims) - all_dims - if missing_dims: + if missing_obj_dims := (set(dims) - all_dims): raise ValueError( - f"Dimensions {tuple(missing_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" + f"Dimensions {tuple(missing_obj_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" ) + if set(self.obj.dims) == set(self.weights.dims): + return + + if not dims: + return + + if missing_weightdims := (set(all_dims) - set(dims)): + exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} + exp_dims = {k: v for k, v in exp_dims.items() if v is not None} + self.weights = self.weights.expand_dims(dims=exp_dims) + # ^expand_dims - sequence of hashable - should this be a list, not a tuple + @staticmethod def _reduce( da: T_DataArray, From 0ebf26e6de8e6b6b836bbe29ee6f8acc170c062c Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Sat, 5 Sep 2026 13:12:33 +0800 Subject: [PATCH 02/12] WIP --- xarray/tests/test_weighted.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py index 5e913c00629..592ed738aff 100644 --- a/xarray/tests/test_weighted.py +++ b/xarray/tests/test_weighted.py @@ -8,13 +8,9 @@ import xarray as xr from xarray import DataArray, Dataset -from xarray.tests import ( - assert_allclose, - assert_equal, - raise_if_dask_computes, - requires_cftime, - requires_dask, -) +from xarray.tests import (assert_allclose, assert_equal, + raise_if_dask_computes, requires_cftime, + requires_dask) @pytest.mark.parametrize("as_dataset", (True, False)) @@ -527,6 +523,7 @@ def expected_weighted(da, weights, dim, skipna, operation): Generate expected result using ``*`` and ``sum``. This is checked against the result of da.weighted which uses ``dot`` """ + breakpoint() weighted_sum = (da * weights).sum(dim=dim, skipna=skipna) From 5ba6ce444d716e3313d43332031054eb8745f78e Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Sat, 5 Sep 2026 21:25:38 +0800 Subject: [PATCH 03/12] WIP --- xarray/computation/weighted.py | 41 +++++++++++++++++++++------------- xarray/tests/test_weighted.py | 11 +++++---- 2 files changed, 32 insertions(+), 20 deletions(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index 7f6471b249c..9a25840e738 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -195,28 +195,21 @@ def _weight_check(w): def _check_dim(self, dim: Dims): """raise an error if any dimension is missing""" + dims, all_dims = self._all_dims(dim) + missing_dims = set(dims) - all_dims + if missing_dims: + raise ValueError( + f"Dimensions {tuple(missing_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" + ) + + def _all_dims(self, dim: Dims) -> tuple[list[Hashable], set[Hashable]]: dims: list[Hashable] if isinstance(dim, str) or not isinstance(dim, Iterable): dims = [dim] if dim else [] else: dims = list(dim) all_dims = set(self.obj.dims).union(set(self.weights.dims)) - if missing_obj_dims := (set(dims) - all_dims): - raise ValueError( - f"Dimensions {tuple(missing_obj_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" - ) - - if set(self.obj.dims) == set(self.weights.dims): - return - - if not dims: - return - - if missing_weightdims := (set(all_dims) - set(dims)): - exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} - exp_dims = {k: v for k, v in exp_dims.items() if v is not None} - self.weights = self.weights.expand_dims(dims=exp_dims) - # ^expand_dims - sequence of hashable - should this be a list, not a tuple + return dims, all_dims @staticmethod def _reduce( @@ -559,8 +552,24 @@ def _implementation(self, func, dim, **kwargs) -> DataArray: class DatasetWeighted(Weighted["Dataset"]): + def _expand_weights(self, dim: Dims): + """expand weights to include any missing dimensions""" + dims, all_dims = self._all_dims(dim) + + if set(self.obj.dims) == set(self.weights.dims): + return + + if not dims: + return + + if missing_weightdims := (set(all_dims) - set(dims)): + exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} + exp_dims = {k: v for k, v in exp_dims.items() if v is not None} + self.weights = self.weights.expand_dims(dims=exp_dims) + def _implementation(self, func, dim, **kwargs) -> Dataset: self._check_dim(dim) + self._expand_weights(dim) return self.obj.map(func, dim=dim, **kwargs) diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py index 592ed738aff..5e913c00629 100644 --- a/xarray/tests/test_weighted.py +++ b/xarray/tests/test_weighted.py @@ -8,9 +8,13 @@ import xarray as xr from xarray import DataArray, Dataset -from xarray.tests import (assert_allclose, assert_equal, - raise_if_dask_computes, requires_cftime, - requires_dask) +from xarray.tests import ( + assert_allclose, + assert_equal, + raise_if_dask_computes, + requires_cftime, + requires_dask, +) @pytest.mark.parametrize("as_dataset", (True, False)) @@ -523,7 +527,6 @@ def expected_weighted(da, weights, dim, skipna, operation): Generate expected result using ``*`` and ``sum``. This is checked against the result of da.weighted which uses ``dot`` """ - breakpoint() weighted_sum = (da * weights).sum(dim=dim, skipna=skipna) From e4c9cc6b824fc22ba837400eea76255febfda933 Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Sat, 5 Sep 2026 22:57:37 +0800 Subject: [PATCH 04/12] Restructure guard claude & relabel dims -> dim. Down to 13 failures --- xarray/computation/weighted.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index 9a25840e738..ca76101ccc9 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -562,10 +562,12 @@ def _expand_weights(self, dim: Dims): if not dims: return - if missing_weightdims := (set(all_dims) - set(dims)): - exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} - exp_dims = {k: v for k, v in exp_dims.items() if v is not None} - self.weights = self.weights.expand_dims(dims=exp_dims) + if not (missing_weightdims := (set(all_dims) - set(dims))): + return + + exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} + exp_dims = {k: v for k, v in exp_dims.items() if v is not None} + self.weights = self.weights.expand_dims(dim=exp_dims) def _implementation(self, func, dim, **kwargs) -> Dataset: self._check_dim(dim) From 994cf18db88efa0054b10af024b2db1a207506ab Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Sat, 5 Sep 2026 23:11:02 +0800 Subject: [PATCH 05/12] Now just one pandas roundtripping property test. Clean up changes a lot too --- xarray/computation/weighted.py | 24 +++++++----------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index ca76101ccc9..1e90e1f640e 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -195,21 +195,18 @@ def _weight_check(w): def _check_dim(self, dim: Dims): """raise an error if any dimension is missing""" - dims, all_dims = self._all_dims(dim) - missing_dims = set(dims) - all_dims - if missing_dims: - raise ValueError( - f"Dimensions {tuple(missing_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" - ) - - def _all_dims(self, dim: Dims) -> tuple[list[Hashable], set[Hashable]]: dims: list[Hashable] if isinstance(dim, str) or not isinstance(dim, Iterable): dims = [dim] if dim else [] else: dims = list(dim) all_dims = set(self.obj.dims).union(set(self.weights.dims)) - return dims, all_dims + + missing_dims = set(dims) - all_dims + if missing_dims: + raise ValueError( + f"Dimensions {tuple(missing_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" + ) @staticmethod def _reduce( @@ -554,15 +551,8 @@ def _implementation(self, func, dim, **kwargs) -> DataArray: class DatasetWeighted(Weighted["Dataset"]): def _expand_weights(self, dim: Dims): """expand weights to include any missing dimensions""" - dims, all_dims = self._all_dims(dim) - - if set(self.obj.dims) == set(self.weights.dims): - return - - if not dims: - return - if not (missing_weightdims := (set(all_dims) - set(dims))): + if not (missing_weightdims := set(self.obj.dims) - set(self.weights.dims)): return exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} From 294d9e5f9419ef4538d06d69c83cc0c50adb20d4 Mon Sep 17 00:00:00 2001 From: Charles Turner <52199577+charles-turner-1@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:12:11 +0800 Subject: [PATCH 06/12] Update xarray/computation/weighted.py --- xarray/computation/weighted.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index 1e90e1f640e..35b60a68fee 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -201,7 +201,6 @@ def _check_dim(self, dim: Dims): else: dims = list(dim) all_dims = set(self.obj.dims).union(set(self.weights.dims)) - missing_dims = set(dims) - all_dims if missing_dims: raise ValueError( From 931d772a73119961a47817435ae6e6aede7930cb Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Mon, 7 Sep 2026 14:36:36 +0800 Subject: [PATCH 07/12] Fix implementation --- xarray/computation/weighted.py | 40 ++++++++++++++++++++++------------ xarray/tests/test_weighted.py | 28 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 14 deletions(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index 35b60a68fee..3d57162a5b6 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -195,11 +195,7 @@ def _weight_check(w): def _check_dim(self, dim: Dims): """raise an error if any dimension is missing""" - dims: list[Hashable] - if isinstance(dim, str) or not isinstance(dim, Iterable): - dims = [dim] if dim else [] - else: - dims = list(dim) + dims = self._dims_to_list(dim) all_dims = set(self.obj.dims).union(set(self.weights.dims)) missing_dims = set(dims) - all_dims if missing_dims: @@ -207,6 +203,15 @@ def _check_dim(self, dim: Dims): f"Dimensions {tuple(missing_dims)} not found in {self.__class__.__name__} dimensions {tuple(all_dims)}" ) + def _dims_to_list(self, dim: Dims) -> list[Hashable]: + dims: list[Hashable] + if isinstance(dim, str) or not isinstance(dim, Iterable): + dims = [dim] if dim else [] + else: + dims = list(dim) + + return dims + @staticmethod def _reduce( da: T_DataArray, @@ -548,20 +553,27 @@ def _implementation(self, func, dim, **kwargs) -> DataArray: class DatasetWeighted(Weighted["Dataset"]): - def _expand_weights(self, dim: Dims): - """expand weights to include any missing dimensions""" + def _restore_dims(self, dim: Dims, ds: Dataset) -> Dataset: + """self.obj.map will drop orphaned coordinates & dims that are not in + the weights DataArray. This restores them. + """ + if dim is None: + return ds + + dims = self._dims_to_list(dim) - if not (missing_weightdims := set(self.obj.dims) - set(self.weights.dims)): - return + existing_dims = {dim for v in ds.variables.values() for dim in v.dims} + dims_to_restore = set(self.obj.dims) - set(dims) - existing_dims - exp_dims = {k: self.obj.sizes.get(k, None) for k in missing_weightdims} - exp_dims = {k: v for k, v in exp_dims.items() if v is not None} - self.weights = self.weights.expand_dims(dim=exp_dims) + dims_to_restore = {d: self.obj.coords[d] for d in dims_to_restore} + ds = ds.assign_coords(dims_to_restore) + return ds def _implementation(self, func, dim, **kwargs) -> Dataset: self._check_dim(dim) - self._expand_weights(dim) - return self.obj.map(func, dim=dim, **kwargs) + res = self.obj.map(func, dim=dim, **kwargs) + res = self._restore_dims(dim, res) + return res def _inject_docstring(cls, cls_name): diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py index 5e913c00629..abbdf55b116 100644 --- a/xarray/tests/test_weighted.py +++ b/xarray/tests/test_weighted.py @@ -808,3 +808,31 @@ def test_weighted_bad_dim(operation, as_dataset): ), ): getattr(data.weighted(weights), operation)(**kwargs) + + +def test_dataset_weighted_mean_preserves_orphaned_dims(): + """ + Can't test against `check_weighted_operations` - this is an edge case where + the dataset has a dimension not present in the weights dataarray or on any + of it's variables. + """ + data = Dataset( + { + "a": (("x", "y"), np.random.randn(2, 2)), + }, + coords={ + "x": ("x", [0, 1]), + "y": ("y", [0, 1]), + "t": ("t", [0, 1]), + }, + ) + weights = DataArray(np.random.randn(2), dims="x") + + weighted_mean = data.weighted(weights).mean(dim="x") + mean = data.mean(dim="x") + + # Don't promote dims *on the dataarray* + assert weighted_mean["a"].dims == mean["a"].dims + # Retain orphaned dims on the dataset + assert set(weighted_mean.dims) == set(mean.dims) + assert set(weighted_mean.dims) == set(mean.dims) From 63798dd7f7bd49e1a17e8b56627e478ecc4a7648 Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Mon, 7 Sep 2026 14:47:04 +0800 Subject: [PATCH 08/12] Rename var & update changelog --- doc/whats-new.rst | 3 +++ xarray/computation/weighted.py | 6 +++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6d67d234020..f5473ba4cc7 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -117,6 +117,9 @@ Bug Fixes type error when applying reduction methods, due to the reduction methods being dynamically generated (:issue:`8136`). By `Andrew Scherer `_. +- Fix DatasetWeighted dropping coordinates/dimensions which are not affected by + the reduction nor in the weights array (:issue:`11560`, :pull:`11562`). + By `Charles Turner` .. _`pandas-dev/pandas#64793`: https://github.com/pandas-dev/pandas/pull/64793 diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index 3d57162a5b6..e2ebdb53da1 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -571,9 +571,9 @@ def _restore_dims(self, dim: Dims, ds: Dataset) -> Dataset: def _implementation(self, func, dim, **kwargs) -> Dataset: self._check_dim(dim) - res = self.obj.map(func, dim=dim, **kwargs) - res = self._restore_dims(dim, res) - return res + mapped_ds = self.obj.map(func, dim=dim, **kwargs) + mapped_ds = self._restore_dims(dim, mapped_ds) + return mapped_ds def _inject_docstring(cls, cls_name): From 7b90a6fe172a04a66334bc457fd0079f35c611ff Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 7 Sep 2026 06:47:34 +0000 Subject: [PATCH 09/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- 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 f5473ba4cc7..343914db1b9 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -117,7 +117,7 @@ Bug Fixes type error when applying reduction methods, due to the reduction methods being dynamically generated (:issue:`8136`). By `Andrew Scherer `_. -- Fix DatasetWeighted dropping coordinates/dimensions which are not affected by +- Fix DatasetWeighted dropping coordinates/dimensions which are not affected by the reduction nor in the weights array (:issue:`11560`, :pull:`11562`). By `Charles Turner` From ad6b20e8d7eb6f0948071dd84d55e334dae09c2f Mon Sep 17 00:00:00 2001 From: Charles Turner <52199577+charles-turner-1@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:50:30 +1000 Subject: [PATCH 10/12] Update pixi.toml --- pixi.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/pixi.toml b/pixi.toml index 0fa6595b300..90fd51a975a 100644 --- a/pixi.toml +++ b/pixi.toml @@ -477,6 +477,3 @@ build-package = { features = [ "release", ], no-default-feature = true } test-nightly = { features = ["nightly"], no-default-feature = true } - -[pypi-dependencies] -pdbpp = ">=0.12.1, <0.13" From 6c34e01e3fe2dccdf86bf8746c7db5ded7427d7d Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Mon, 7 Sep 2026 14:51:58 +0800 Subject: [PATCH 11/12] Fix changelog entry --- 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 f5473ba4cc7..5184ebd74ce 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -119,7 +119,7 @@ Bug Fixes By `Andrew Scherer `_. - Fix DatasetWeighted dropping coordinates/dimensions which are not affected by the reduction nor in the weights array (:issue:`11560`, :pull:`11562`). - By `Charles Turner` + By `Charles Turner `_. .. _`pandas-dev/pandas#64793`: https://github.com/pandas-dev/pandas/pull/64793 From 6ef78443e1abe343017fa035831a14cd7beac3a5 Mon Sep 17 00:00:00 2001 From: Charles Turner Date: Mon, 7 Sep 2026 15:06:46 +0800 Subject: [PATCH 12/12] Clean up typing --- xarray/computation/weighted.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index e2ebdb53da1..91b7165432b 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -565,15 +565,13 @@ def _restore_dims(self, dim: Dims, ds: Dataset) -> Dataset: existing_dims = {dim for v in ds.variables.values() for dim in v.dims} dims_to_restore = set(self.obj.dims) - set(dims) - existing_dims - dims_to_restore = {d: self.obj.coords[d] for d in dims_to_restore} - ds = ds.assign_coords(dims_to_restore) - return ds + coords = {d: self.obj.coords[d] for d in dims_to_restore} + return ds.assign_coords(coords) def _implementation(self, func, dim, **kwargs) -> Dataset: self._check_dim(dim) mapped_ds = self.obj.map(func, dim=dim, **kwargs) - mapped_ds = self._restore_dims(dim, mapped_ds) - return mapped_ds + return self._restore_dims(dim, mapped_ds) def _inject_docstring(cls, cls_name):