From 572a52883c11c53fa0d82c86b95fbf2be5647b90 Mon Sep 17 00:00:00 2001 From: Julia Kent <46687291+jukent@users.noreply.github.com> Date: Wed, 9 Sep 2020 12:56:35 -0600 Subject: [PATCH] Preserve attrs with coarsen (#4360) * 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 * 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 Co-authored-by: Deepak Cherian Co-authored-by: Maximilian Roos --- .pre-commit-config.yaml | 2 +- doc/whats-new.rst | 1 + xarray/core/rolling.py | 3 ++- xarray/core/variable.py | 18 +++++++++++++----- xarray/tests/test_dataarray.py | 29 +++++++++++++++++++++++++++++ xarray/tests/test_dataset.py | 5 +++++ xarray/tests/test_variable.py | 5 ++++- 7 files changed, 55 insertions(+), 8 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7265bbac155..161652888c9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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 diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 316254a97d7..b948679b637 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -77,6 +77,7 @@ Bug fixes and :py:meth:`DataArray.str.wrap` (:issue:`4334`). By `Mathias Hauser `_. - Fixed overflow issue causing incorrect results in computing means of :py:class:`cftime.datetime` arrays (:issue:`4341`). By `Spencer Clark `_. +- Fixed :py:meth:`Dataset.coarsen`, :py:meth:`DataArray.coarsen` dropping attributes on original object (:issue:`4120`, :pull:`4360`). by `Julia Kent `_. - fix the signature of the plot methods. (:pull:`4359`) By `Justus Magin `_. - Fix :py:func:`xarray.apply_ufunc` with ``vectorize=True`` and ``exclude_dims`` (:issue:`3890`). By `Mathias Hauser `_. diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 0c4614e0b57..3d7be955642 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -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(): @@ -703,6 +703,7 @@ def wrapped_func(self, **kwargs): self.coord_func[c], self.boundary, self.side, + self.keep_attrs, **kwargs, ) else: diff --git a/xarray/core/variable.py b/xarray/core/variable.py index ce6df5282c5..075d79043b2 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -1993,7 +1993,9 @@ 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. """ @@ -2001,13 +2003,22 @@ def coarsen(self, windows, func, boundary="exact", side="left", **kwargs): 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): """ @@ -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 diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 8fd5644646b..a22ed58c9bc 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -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): diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 16b1cc330f1..40e2bdfc6de 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -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 == {} @@ -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"} diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 169105bb4d0..efebe09e2ec 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -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