Skip to content

Commit

Permalink
Preserve attrs with coarsen (pydata#4360)
Browse files Browse the repository at this point in the history
* pass **kwargs into
_coarsen_reshape

* change _replace to copy in coarsen return

* copy back to replace, variable=copy

* take out self.copy

* update pre-commit config (tags not branches)

* add test that coarsen maintains OG object

* del comment

* check global for keep attrs

* black reformatter

* line break

* variable._attrs to variable.attrs

* if not keep_attrs

* {} to None

* set_options goes in a with block

* remove test dependency on netcdf

* add bug fix to whats-new.rst

* Update doc/whats-new.rst

Co-authored-by: Deepak Cherian <[email protected]>

* go back to v0.1.2 of blackdock

* add test_coarsen_keep_attrs to test_dataarray.py

* fix tests

* black test_dataarray

* xr.set_options

* move keep_attrs to coarsen from _reshape_coarsen

* flake8

* clean up

* move keep_attrs to fx signature

* remove kwarg check for keep_attrs

* black on variable.py

* fix test_variable

* Format with black

* fix test

* check for global attribute

* black variable.py

* black test_variable.py

* format w black

* Update xarray/core/variable.py

Co-authored-by: Deepak Cherian <[email protected]>

Co-authored-by: Deepak Cherian <[email protected]>
Co-authored-by: Maximilian Roos <[email protected]>
  • Loading branch information
3 people authored Sep 9, 2020
1 parent 66259d1 commit 572a528
Show file tree
Hide file tree
Showing 7 changed files with 55 additions and 8 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ repos:
hooks:
- id: isort
# https://github.com/python/black#version-control-integration
- repo: https://github.com/python/black
- repo: https://github.com/psf/black
rev: 20.8b1
hooks:
- id: black
Expand Down
1 change: 1 addition & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ Bug fixes
and :py:meth:`DataArray.str.wrap` (:issue:`4334`). By `Mathias Hauser <https://github.com/mathause>`_.
- Fixed overflow issue causing incorrect results in computing means of :py:class:`cftime.datetime`
arrays (:issue:`4341`). By `Spencer Clark <https://github.com/spencerkclark>`_.
- Fixed :py:meth:`Dataset.coarsen`, :py:meth:`DataArray.coarsen` dropping attributes on original object (:issue:`4120`, :pull:`4360`). by `Julia Kent <https://github.com/jukent>`_.
- fix the signature of the plot methods. (:pull:`4359`) By `Justus Magin <https://github.com/keewis>`_.
- Fix :py:func:`xarray.apply_ufunc` with ``vectorize=True`` and ``exclude_dims`` (:issue:`3890`).
By `Mathias Hauser <https://github.com/mathause>`_.
Expand Down
3 changes: 2 additions & 1 deletion xarray/core/rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -690,7 +690,7 @@ def wrapped_func(self, **kwargs):
from .dataarray import DataArray

reduced = self.obj.variable.coarsen(
self.windows, func, self.boundary, self.side, **kwargs
self.windows, func, self.boundary, self.side, self.keep_attrs, **kwargs
)
coords = {}
for c, v in self.obj.coords.items():
Expand All @@ -703,6 +703,7 @@ def wrapped_func(self, **kwargs):
self.coord_func[c],
self.boundary,
self.side,
self.keep_attrs,
**kwargs,
)
else:
Expand Down
18 changes: 13 additions & 5 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -1993,21 +1993,32 @@ def rolling_window(
),
)

def coarsen(self, windows, func, boundary="exact", side="left", **kwargs):
def coarsen(
self, windows, func, boundary="exact", side="left", keep_attrs=None, **kwargs
):
"""
Apply reduction function.
"""
windows = {k: v for k, v in windows.items() if k in self.dims}
if not windows:
return self.copy()

if keep_attrs is None:
keep_attrs = _get_keep_attrs(default=False)

if keep_attrs:
_attrs = self.attrs
else:
_attrs = None

reshaped, axes = self._coarsen_reshape(windows, boundary, side)
if isinstance(func, str):
name = func
func = getattr(duck_array_ops, name, None)
if func is None:
raise NameError(f"{name} is not a valid method.")
return self._replace(data=func(reshaped, axis=axes, **kwargs))

return self._replace(data=func(reshaped, axis=axes, **kwargs), attrs=_attrs)

def _coarsen_reshape(self, windows, boundary, side):
"""
Expand Down Expand Up @@ -2072,9 +2083,6 @@ def _coarsen_reshape(self, windows, boundary, side):
else:
shape.append(variable.shape[i])

keep_attrs = _get_keep_attrs(default=False)
variable.attrs = variable._attrs if keep_attrs else {}

return variable.data.reshape(shape), tuple(axes)

@property
Expand Down
29 changes: 29 additions & 0 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -6186,6 +6186,35 @@ def test_isin(da):
assert_equal(result, expected)


def test_coarsen_keep_attrs():
_attrs = {"units": "test", "long_name": "testing"}

da = xr.DataArray(
np.linspace(0, 364, num=364),
dims="time",
coords={"time": pd.date_range("15/12/1999", periods=364)},
attrs=_attrs,
)

da2 = da.copy(deep=True)

# Test dropped attrs
dat = da.coarsen(time=3, boundary="trim").mean()
assert dat.attrs == {}

# Test kept attrs using dataset keyword
dat = da.coarsen(time=3, boundary="trim", keep_attrs=True).mean()
assert dat.attrs == _attrs

# Test kept attrs using global option
with xr.set_options(keep_attrs=True):
dat = da.coarsen(time=3, boundary="trim").mean()
assert dat.attrs == _attrs

# Test kept attrs in original object
xr.testing.assert_identical(da, da2)


@pytest.mark.filterwarnings("error:Mean of empty slice")
@pytest.mark.parametrize("da", (1, 2), indirect=True)
def test_rolling_iter(da):
Expand Down
5 changes: 5 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -5956,6 +5956,8 @@ def test_coarsen_keep_attrs():
attrs=_attrs,
)

ds2 = ds.copy(deep=True)

# Test dropped attrs
dat = ds.coarsen(coord=5).mean()
assert dat.attrs == {}
Expand All @@ -5969,6 +5971,9 @@ def test_coarsen_keep_attrs():
dat = ds.coarsen(coord=5).mean()
assert dat.attrs == _attrs

# Test kept attrs in original object
xr.testing.assert_identical(ds, ds2)


def test_rolling_keep_attrs():
_attrs = {"units": "test", "long_name": "testing"}
Expand Down
5 changes: 4 additions & 1 deletion xarray/tests/test_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -1949,7 +1949,10 @@ def test_coarsen_keep_attrs(self, operation="mean"):
# Test kept attrs
with set_options(keep_attrs=True):
new = Variable(["coord"], np.linspace(1, 10, 100), attrs=_attrs).coarsen(
windows={"coord": 1}, func=test_func, boundary="exact", side="left"
windows={"coord": 1},
func=test_func,
boundary="exact",
side="left",
)
assert new.attrs == _attrs

Expand Down

0 comments on commit 572a528

Please sign in to comment.