fix: Don't drop orphaned dims on weighted reductions

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 <https://github.com/andrew-s28>`_.
+- Fix DatasetWeighted dropping coordinates/dimensions which are not affected by
+  the reduction nor in the weights array (:issue:`11560`, :pull:`11562`).
+  By `Charles Turner <https://github.com/charles-turner-1>`_.
 
 .. _`pandas-dev/pandas#64793`: https://github.com/pandas-dev/pandas/pull/64793
 

xarray/computation/weighted.py

@@ -195,18 +195,23 @@ 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:
             raise ValueError(
                 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):

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)