Skip to content

Add option to not roll coords #2360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Aug 15, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ Enhancements
(:issue:`1560`)
By `Maximilian Maahn <https://github.com/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 <https://github.com/ahuang11>`_.

Bug fixes
~~~~~~~~~
Expand Down
15 changes: 11 additions & 4 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down
44 changes: 32 additions & 12 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down Expand Up @@ -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
Expand All @@ -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)

Expand Down
19 changes: 17 additions & 2 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
30 changes: 27 additions & 3 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'}
Expand Down