diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index fc8edb1473c..b61408b8ae0 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4882,7 +4882,7 @@ def identical(self, other: Self) -> bool: def __array_wrap__(self, obj, context=None, return_scalar=False) -> Self: new_var = self.variable.__array_wrap__(obj, context, return_scalar) - return self._replace(new_var) + return self._replace_maybe_drop_dims(new_var) def __matmul__(self, obj: T_Xarray) -> T_Xarray: return self.dot(obj) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 1b5e0d4ff69..7bf40a3a5f8 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -2432,6 +2432,19 @@ def real(self) -> Variable: return self._new(data=self.data.real) def __array_wrap__(self, obj, context=None, return_scalar=False): + if obj.shape != self.shape: + # Shape changed during the numpy operation. + # Check if only the last two dims are transposed (e.g. np.linalg.pinv). + if ( + obj.ndim == self.ndim + and obj.ndim >= 2 + and obj.shape[:-2] == self.shape[:-2] + and obj.shape[-2:] == self.shape[-2:][::-1] + ): + new_dims = self.dims[:-2] + (self.dims[-1], self.dims[-2]) + return Variable(new_dims, obj) + # Fallback: use generic dim names since we can't reliably map dims. + return Variable(tuple(f"dim_{i}" for i in range(obj.ndim)), obj) return Variable(self.dims, obj) def _unary_op(self, f, *args, **kwargs): diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 377fd2f8a8b..d2ae2e351f3 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -2478,6 +2478,34 @@ def test_array_interface(self) -> None: bar = Variable(["x", "y"], np.zeros((10, 20))) assert_equal(self.dv, np.maximum(self.dv, bar)) + def test_array_wrap_drops_mismatched_coords(self) -> None: + array = DataArray( + np.arange(12).reshape(3, 4), + dims=("foo", "bar"), + coords={ + "foo": ["x", "y", "z"], + "bar": ["a", "b", "c", "d"], + "aux": (("foo", "bar"), np.ones((3, 4))), + "scalar": 1, + }, + ) + + actual = np.linalg.pinv(array) + # pinv transposes the last two dims, so dims are swapped from + # (foo: 3, bar: 4) to (bar: 4, foo: 3). All coords whose dims + # still exist in the result are preserved. + expected = DataArray( + np.linalg.pinv(array.data), + dims=("bar", "foo"), + coords={ + "foo": ["x", "y", "z"], + "bar": ["a", "b", "c", "d"], + "aux": (("foo", "bar"), np.ones((3, 4))), + "scalar": 1, + }, + ) + assert_identical(expected, actual) + def test_astype_attrs(self) -> None: # Split into two loops for mypy - Variable, DataArray, and Dataset # don't share a common base class, so mypy infers type object for v,