Weighted reductions drop orphaned dimensions

[charles-turner-1]

What happened?

Weighted reductions drop orphaned dimensions - unlike their non-weighted versions. See the MRE, but TLDR;

********************  base array  ********************
<xarray.Dataset> Size: 56B
Dimensions:    (lat: 3, time: 1)
Coordinates:
  * lat        (lat) int64 24B 0 1 2
  * time       (time) int64 8B 0
Data variables:
    elevation  (lat) float64 24B 0.0 1.0 2.0
********************  weighted mean: base.weighted(weights).mean('lat') ********************
<xarray.Dataset> Size: 8B
Dimensions:    ()
Data variables:
    elevation  float64 8B 1.0
********************  regular mean: base.mean('lat') ********************
<xarray.Dataset> Size: 16B
Dimensions:    (time: 1)
Coordinates:
  * time       (time) int64 8B 0
Data variables:

I'll open a PR with a fix - assuming this is the desired behaviour

What did you expect to happen?

Given that operations like sel/isel don't drop the dimensions by default, I think it makes sense to keep the dimension by default here too.

Minimal Complete Verifiable Example

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "xarray[complete]@git+https://github.com/pydata/xarray.git@main",
# ]
# ///
#
# This script automatically imports the development branch of xarray to check for issues.
# Please delete this header if you have _not_ tested this script with `uv run`!

import numpy as np

import xarray as xr

xr.show_versions()

base = xr.Dataset(
    {"elevation": (("lat",), np.array([0.0, 1.0, 2.0]))},
    coords={"time": [0], "lat": [0, 1, 2]},
)
weights = xr.DataArray([1.0, 1.0, 1.0], dims="lat", coords={"lat": [0, 1, 2]})

print("*" * 20, " base array ", "*" * 20)
print(base)

print("*" * 20, " weighted mean: base.weighted(weights).mean('lat')", "*" * 20)
print(base.weighted(weights).mean("lat"))

print("*" * 20, " regular mean: base.mean('lat')", "*" * 20)
print(base.mean("lat"))

Steps to reproduce

Should work for any other weighted reduction

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

[charles-turner-1]

This diff:

diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py
index b311290a..313741db 100644
--- a/xarray/computation/weighted.py
+++ b/xarray/computation/weighted.py
@@ -201,10 +201,13 @@ class Weighted(Generic[T_Xarray]):
         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 missing_weightdims := (set(all_dims) - set(dims)):
+            raise ValueError(
+                f"Dimensions {tuple(missing_weightdims)} not found in weights array dimensions {tuple(dims)}"
             )

breaks 375/492 of the tests in xarray/tests/test_weighted.py: 375 failed, 47 passed, 72 skipped

Would be the simplest 'fix' but it looks like the behaviour is being relied on pretty heavily by the test suite. Looks a bit tough putting the orphaned dim back into a DataArray though.

[charles-turner-1]

This diff gives the expected result, keeping the orphaned dimension, but breaks 389 tests, mostly in the same test file:

diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py
index b311290a..d5adb041 100644
--- a/xarray/computation/weighted.py
+++ b/xarray/computation/weighted.py
@@ -201,11 +201,13 @@ class Weighted(Generic[T_Xarray]):
         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 missing_weightdims := (set(all_dims) - set(dims)):
+            self.weights = self.weights.expand_dims(dim=list(missing_weightdims))
+            # ^expand_dims - sequence of hashable - should this be a list, not a tuple

[dcherian]

But is that because the tests encode bad behavior?

[charles-turner-1]

Think that fix was a bit too crude tbh. The failures I'm looking at now seem like the arrays pinning behaviour need the dimensions which were dropped adding.