From cbb2aeb6492ad5364694396fb10e3b86abfe0aa6 Mon Sep 17 00:00:00 2001 From: Andrew Huang Date: Wed, 15 Aug 2018 01:11:28 -0700 Subject: [PATCH] Add option to not roll coords (#2360) * Add option to not roll coords * Rename keyword arg and add tests * Add what's new * Fix passing None and add more tests * Revise from comments * Revise with cleaner version * Revisions based on comments * Fix either_dict_or_kwargs * Revisions from comments --- doc/whats-new.rst | 5 ++++ xarray/core/dataarray.py | 15 ++++++++---- xarray/core/dataset.py | 44 ++++++++++++++++++++++++---------- xarray/tests/test_dataarray.py | 19 +++++++++++++-- xarray/tests/test_dataset.py | 30 ++++++++++++++++++++--- 5 files changed, 92 insertions(+), 21 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 4552a4ca546..47d39b967e3 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -56,6 +56,11 @@ Enhancements (:issue:`1560`) By `Maximilian Maahn `_. +- You can now control whether or not to offset the coordinates when using + the ``roll`` method and the current behavior, coordinates rolled by default, + raises a deprecation warning unless explicitly setting the keyword argument. + (:issue:`1875`) + By `Andrew Huang `_. Bug fixes ~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index f215bc47df8..b1be994416e 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -2015,14 +2015,20 @@ def shift(self, **shifts): variable = self.variable.shift(**shifts) return self._replace(variable) - def roll(self, **shifts): + def roll(self, shifts=None, roll_coords=None, **shifts_kwargs): """Roll this array by an offset along one or more dimensions. - Unlike shift, roll rotates all variables, including coordinates. The - direction of rotation is consistent with :py:func:`numpy.roll`. + Unlike shift, roll may rotate all variables, including coordinates + if specified. The direction of rotation is consistent with + :py:func:`numpy.roll`. Parameters ---------- + roll_coords : bool + Indicates whether to roll the coordinates by the offset + The current default of roll_coords (None, equivalent to True) is + deprecated and will change to False in a future version. + Explicitly pass roll_coords to silence the warning. **shifts : keyword arguments of the form {dim: offset} Integer offset to rotate each of the given dimensions. Positive offsets roll to the right; negative offsets roll to the left. @@ -2046,7 +2052,8 @@ def roll(self, **shifts): Coordinates: * x (x) int64 2 0 1 """ - ds = self._to_temp_dataset().roll(**shifts) + ds = self._to_temp_dataset().roll( + shifts=shifts, roll_coords=roll_coords, **shifts_kwargs) return self._from_temp_dataset(ds) @property diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index e6bc2f8aeaf..37544aca372 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -2324,13 +2324,13 @@ def unstack(self, dim): 'a MultiIndex') full_idx = pd.MultiIndex.from_product(index.levels, names=index.names) - + # take a shortcut in case the MultiIndex was not modified. if index.equals(full_idx): obj = self else: obj = self.reindex(copy=False, **{dim: full_idx}) - + new_dim_names = index.names new_dim_sizes = [lev.size for lev in index.levels] @@ -3360,18 +3360,28 @@ def shift(self, **shifts): return self._replace_vars_and_dims(variables) - def roll(self, **shifts): + def roll(self, shifts=None, roll_coords=None, **shifts_kwargs): """Roll this dataset by an offset along one or more dimensions. - Unlike shift, roll rotates all variables, including coordinates. The - direction of rotation is consistent with :py:func:`numpy.roll`. + Unlike shift, roll may rotate all variables, including coordinates + if specified. The direction of rotation is consistent with + :py:func:`numpy.roll`. Parameters ---------- - **shifts : keyword arguments of the form {dim: offset} - Integer offset to rotate each of the given dimensions. Positive - offsets roll to the right; negative offsets roll to the left. + shifts : dict, optional + A dict with keys matching dimensions and values given + by integers to rotate each of the given dimensions. Positive + offsets roll to the right; negative offsets roll to the left. + roll_coords : bool + Indicates whether to roll the coordinates by the offset + The current default of roll_coords (None, equivalent to True) is + deprecated and will change to False in a future version. + Explicitly pass roll_coords to silence the warning. + **shifts_kwargs : {dim: offset, ...}, optional + The keyword arguments form of ``shifts``. + One of shifts or shifts_kwargs must be provided. Returns ------- rolled : Dataset @@ -3394,15 +3404,25 @@ def roll(self, **shifts): Data variables: foo (x) object 'd' 'e' 'a' 'b' 'c' """ + shifts = either_dict_or_kwargs(shifts, shifts_kwargs, 'roll') invalid = [k for k in shifts if k not in self.dims] if invalid: raise ValueError("dimensions %r do not exist" % invalid) + if roll_coords is None: + warnings.warn("roll_coords will be set to False in the future." + " Explicitly set roll_coords to silence warning.", + FutureWarning, stacklevel=2) + roll_coords = True + + unrolled_vars = () if roll_coords else self.coords + variables = OrderedDict() - for name, var in iteritems(self.variables): - var_shifts = dict((k, v) for k, v in shifts.items() - if k in var.dims) - variables[name] = var.roll(**var_shifts) + for k, v in iteritems(self.variables): + if k not in unrolled_vars: + variables[k] = v.roll(**shifts) + else: + variables[k] = v return self._replace_vars_and_dims(variables) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 2950e97cc75..1a115192fb4 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -3099,9 +3099,24 @@ def test_shift(self): actual = arr.shift(x=offset) assert_identical(expected, actual) - def test_roll(self): + def test_roll_coords(self): arr = DataArray([1, 2, 3], coords={'x': range(3)}, dims='x') - actual = arr.roll(x=1) + actual = arr.roll(x=1, roll_coords=True) + expected = DataArray([3, 1, 2], coords=[('x', [2, 0, 1])]) + assert_identical(expected, actual) + + def test_roll_no_coords(self): + arr = DataArray([1, 2, 3], coords={'x': range(3)}, dims='x') + actual = arr.roll(x=1, roll_coords=False) + expected = DataArray([3, 1, 2], coords=[('x', [0, 1, 2])]) + assert_identical(expected, actual) + + def test_roll_coords_none(self): + arr = DataArray([1, 2, 3], coords={'x': range(3)}, dims='x') + + with pytest.warns(FutureWarning): + actual = arr.roll(x=1, roll_coords=None) + expected = DataArray([3, 1, 2], coords=[('x', [2, 0, 1])]) assert_identical(expected, actual) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index c67183db1ec..e2a406b1e51 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -3862,18 +3862,42 @@ def test_shift(self): with raises_regex(ValueError, 'dimensions'): ds.shift(foo=123) - def test_roll(self): + def test_roll_coords(self): coords = {'bar': ('x', list('abc')), 'x': [-4, 3, 2]} attrs = {'meta': 'data'} ds = Dataset({'foo': ('x', [1, 2, 3])}, coords, attrs) - actual = ds.roll(x=1) + actual = ds.roll(x=1, roll_coords=True) ex_coords = {'bar': ('x', list('cab')), 'x': [2, -4, 3]} expected = Dataset({'foo': ('x', [3, 1, 2])}, ex_coords, attrs) assert_identical(expected, actual) with raises_regex(ValueError, 'dimensions'): - ds.roll(foo=123) + ds.roll(foo=123, roll_coords=True) + + def test_roll_no_coords(self): + coords = {'bar': ('x', list('abc')), 'x': [-4, 3, 2]} + attrs = {'meta': 'data'} + ds = Dataset({'foo': ('x', [1, 2, 3])}, coords, attrs) + actual = ds.roll(x=1, roll_coords=False) + + expected = Dataset({'foo': ('x', [3, 1, 2])}, coords, attrs) + assert_identical(expected, actual) + + with raises_regex(ValueError, 'dimensions'): + ds.roll(abc=321, roll_coords=False) + + def test_roll_coords_none(self): + coords = {'bar': ('x', list('abc')), 'x': [-4, 3, 2]} + attrs = {'meta': 'data'} + ds = Dataset({'foo': ('x', [1, 2, 3])}, coords, attrs) + + with pytest.warns(FutureWarning): + actual = ds.roll(x=1, roll_coords=None) + + ex_coords = {'bar': ('x', list('cab')), 'x': [2, -4, 3]} + expected = Dataset({'foo': ('x', [3, 1, 2])}, ex_coords, attrs) + assert_identical(expected, actual) def test_real_and_imag(self): attrs = {'foo': 'bar'}