diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6d67d234020..72d9c7c16ea 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 b311290aabf..91b7165432b 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,9 +553,25 @@ def _implementation(self, func, dim, **kwargs) -> DataArray: class DatasetWeighted(Weighted["Dataset"]): + 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) + + 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 + + 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) - return self.obj.map(func, dim=dim, **kwargs) + mapped_ds = self.obj.map(func, dim=dim, **kwargs) + return self._restore_dims(dim, mapped_ds) 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)