apply_ufunc drops passed DataArray's attrs when keep_attrs=False

[charlesgauthier-udm]

What happened?

Coming from xESMF/issues/286. When passing a DataArray to apply_ufunc with the kwarg keep_attrs = False, the attrs on the input DataArray are dropped.

Documentation states:

‘drop’ or False: empty attrs on returned xarray object.

It does not state that the input argument's attrs will be dropped.

What did you expect to happen?

I expected the input DataArray to keep its attrs when passed to apply_ufunc.

Minimal Complete Verifiable Example

import xarray as xr
import numpy as np

def magnitude(a, b):
    func = lambda x, y: np.sqrt(x ** 2 + y ** 2)
    return xr.apply_ufunc(func, a, b, keep_attrs=False)

array = xr.DataArray([1, 2, 3], coords=[("x", [0.1, 0.2, 0.3], {'axis': 'X'})])

before = array['x'].attrs
magnitude(array, -array)
after = array['x'].attrs

print(before)
print(after)

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.

Relevant log output

No response

Anything else we need to know?

No response

Environment

Details

INSTALLED VERSIONS

commit: None
python: 3.10.11 | packaged by conda-forge | (main, May 10 2023, 18:58:44) [GCC 11.3.0]
python-bits: 64
OS: Linux
OS-release: 5.15.0-76-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: fr_CA.UTF-8
LOCALE: ('fr_CA', 'UTF-8')
libhdf5: 1.14.0
libnetcdf: 4.9.2
xarray: 2023.5.0
pandas: 2.0.2
numpy: 1.24.3
scipy: 1.10.1
netCDF4: 1.6.4
pydap: installed
h5netcdf: 1.2.0
h5py: 3.9.0
Nio: None
zarr: None
cftime: 1.6.2
nc_time_axis: None
PseudoNetCDF: None
iris: None
bottleneck: None
dask: 2023.6.0
distributed: 2023.6.0
matplotlib: 3.7.1
cartopy: None
seaborn: None
numbagg: None
fsspec: 2023.6.0
cupy: None
pint: None
sparse: 0.14.0
flox: None
numpy_groupies: None
setuptools: 67.7.2
pip: 23.1.2
conda: None
pytest: 7.3.2
mypy: None
IPython: 8.14.0
sphinx: None

[welcome]

Thanks for opening your first issue here at xarray! Be sure to follow the issue template!
If you have an idea for a solution, we would really welcome a Pull Request with proposed changes.
See the Contributing Guide for more.
It may take us a while to respond here, but we really value your contribution. Contributors like you help make xarray better.
Thank you!

[jhamman]

Thanks for the report @charlesgauthier-udm. This does seem to be surprising behavior and is likely a bug. There is no reason that Xarray should be modifying the inputs to apply_ufunc.

[keewis]

This appears to be limited to dimension coordinates (not sure if other coordinates with indexes would be affected, I didn't try), with the cause being these lines:

xarray/xarray/core/merge.py

Lines 281 to 285 in fb6bbbe

	merged_vars[name] = variable
	merged_vars[name].attrs = merge_attrs(
	[var.attrs for var, _ in indexed_elements],
	combine_attrs=combine_attrs,
	)

Changing the first line to merged_vars[name] = variable.copy(deep=False) seems to fix this, but something like

merged_attrs = merge_attrs(...)
merged_vars[name] = variable.assign_attrs(merged_attrs)

might be cleaner (if Variable.assign_attrs existed)?

cc @benbovy

[benbovy]

not sure if other coordinates with indexes would be affected

Shallow copy of indexed coordinate variables should be OK I think. Although generally it would be safer if all coordinate variables bound to a same index are handled together.

There might be some remaining (temporary) hacks in xarray internals were we rely on variable ._data identity for extracting and wrapping pandas multi-index objects. I'm not 100% sure though. If that's the case, even a shallow copy will break those hacks.

[yt87]

Using where with keep_attrs=True also removes attributes. Here is a reworked example from #8480:

import xarray as xr
import numpy as np

values = xr.DataArray(np.linspace(0, 2), dims='x')
values = values.assign_coords(x=values.x.assign_attrs(attr=10))
print('attr' in values.x.attrs)  # True
zero = values == 0
print('attr' in values.x.attrs)  # True
xr.where(zero, np.nan, values, keep_attrs=True)
print('attr' in values.x.attrs)  # False

keep_attrs=True keeps attributes of the first option (x). So, if there are none, the argument is ignored.
I think it would make sense to make it symmetric and allow either one, e.g. keep_attrs='y'. For example, in the expression xr.where(u>0.5, 1, 2*u) one cannot simply swap inequality sign and the arguments when u can contain NaNs (I know I could use clip instead).