Skip to content

Commit cbb2aeb

Browse files
ahuang11fujiisoup
authored andcommitted
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
1 parent c27ca43 commit cbb2aeb

File tree

5 files changed

+92
-21
lines changed

5 files changed

+92
-21
lines changed

doc/whats-new.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,11 @@ Enhancements
5656
(:issue:`1560`)
5757
By `Maximilian Maahn <https://github.com/maahn>`_.
5858

59+
- You can now control whether or not to offset the coordinates when using
60+
the ``roll`` method and the current behavior, coordinates rolled by default,
61+
raises a deprecation warning unless explicitly setting the keyword argument.
62+
(:issue:`1875`)
63+
By `Andrew Huang <https://github.com/ahuang11>`_.
5964

6065
Bug fixes
6166
~~~~~~~~~

xarray/core/dataarray.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2015,14 +2015,20 @@ def shift(self, **shifts):
20152015
variable = self.variable.shift(**shifts)
20162016
return self._replace(variable)
20172017

2018-
def roll(self, **shifts):
2018+
def roll(self, shifts=None, roll_coords=None, **shifts_kwargs):
20192019
"""Roll this array by an offset along one or more dimensions.
20202020
2021-
Unlike shift, roll rotates all variables, including coordinates. The
2022-
direction of rotation is consistent with :py:func:`numpy.roll`.
2021+
Unlike shift, roll may rotate all variables, including coordinates
2022+
if specified. The direction of rotation is consistent with
2023+
:py:func:`numpy.roll`.
20232024
20242025
Parameters
20252026
----------
2027+
roll_coords : bool
2028+
Indicates whether to roll the coordinates by the offset
2029+
The current default of roll_coords (None, equivalent to True) is
2030+
deprecated and will change to False in a future version.
2031+
Explicitly pass roll_coords to silence the warning.
20262032
**shifts : keyword arguments of the form {dim: offset}
20272033
Integer offset to rotate each of the given dimensions. Positive
20282034
offsets roll to the right; negative offsets roll to the left.
@@ -2046,7 +2052,8 @@ def roll(self, **shifts):
20462052
Coordinates:
20472053
* x (x) int64 2 0 1
20482054
"""
2049-
ds = self._to_temp_dataset().roll(**shifts)
2055+
ds = self._to_temp_dataset().roll(
2056+
shifts=shifts, roll_coords=roll_coords, **shifts_kwargs)
20502057
return self._from_temp_dataset(ds)
20512058

20522059
@property

xarray/core/dataset.py

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2324,13 +2324,13 @@ def unstack(self, dim):
23242324
'a MultiIndex')
23252325

23262326
full_idx = pd.MultiIndex.from_product(index.levels, names=index.names)
2327-
2327+
23282328
# take a shortcut in case the MultiIndex was not modified.
23292329
if index.equals(full_idx):
23302330
obj = self
23312331
else:
23322332
obj = self.reindex(copy=False, **{dim: full_idx})
2333-
2333+
23342334
new_dim_names = index.names
23352335
new_dim_sizes = [lev.size for lev in index.levels]
23362336

@@ -3360,18 +3360,28 @@ def shift(self, **shifts):
33603360

33613361
return self._replace_vars_and_dims(variables)
33623362

3363-
def roll(self, **shifts):
3363+
def roll(self, shifts=None, roll_coords=None, **shifts_kwargs):
33643364
"""Roll this dataset by an offset along one or more dimensions.
33653365
3366-
Unlike shift, roll rotates all variables, including coordinates. The
3367-
direction of rotation is consistent with :py:func:`numpy.roll`.
3366+
Unlike shift, roll may rotate all variables, including coordinates
3367+
if specified. The direction of rotation is consistent with
3368+
:py:func:`numpy.roll`.
33683369
33693370
Parameters
33703371
----------
3371-
**shifts : keyword arguments of the form {dim: offset}
3372-
Integer offset to rotate each of the given dimensions. Positive
3373-
offsets roll to the right; negative offsets roll to the left.
33743372
3373+
shifts : dict, optional
3374+
A dict with keys matching dimensions and values given
3375+
by integers to rotate each of the given dimensions. Positive
3376+
offsets roll to the right; negative offsets roll to the left.
3377+
roll_coords : bool
3378+
Indicates whether to roll the coordinates by the offset
3379+
The current default of roll_coords (None, equivalent to True) is
3380+
deprecated and will change to False in a future version.
3381+
Explicitly pass roll_coords to silence the warning.
3382+
**shifts_kwargs : {dim: offset, ...}, optional
3383+
The keyword arguments form of ``shifts``.
3384+
One of shifts or shifts_kwargs must be provided.
33753385
Returns
33763386
-------
33773387
rolled : Dataset
@@ -3394,15 +3404,25 @@ def roll(self, **shifts):
33943404
Data variables:
33953405
foo (x) object 'd' 'e' 'a' 'b' 'c'
33963406
"""
3407+
shifts = either_dict_or_kwargs(shifts, shifts_kwargs, 'roll')
33973408
invalid = [k for k in shifts if k not in self.dims]
33983409
if invalid:
33993410
raise ValueError("dimensions %r do not exist" % invalid)
34003411

3412+
if roll_coords is None:
3413+
warnings.warn("roll_coords will be set to False in the future."
3414+
" Explicitly set roll_coords to silence warning.",
3415+
FutureWarning, stacklevel=2)
3416+
roll_coords = True
3417+
3418+
unrolled_vars = () if roll_coords else self.coords
3419+
34013420
variables = OrderedDict()
3402-
for name, var in iteritems(self.variables):
3403-
var_shifts = dict((k, v) for k, v in shifts.items()
3404-
if k in var.dims)
3405-
variables[name] = var.roll(**var_shifts)
3421+
for k, v in iteritems(self.variables):
3422+
if k not in unrolled_vars:
3423+
variables[k] = v.roll(**shifts)
3424+
else:
3425+
variables[k] = v
34063426

34073427
return self._replace_vars_and_dims(variables)
34083428

xarray/tests/test_dataarray.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3099,9 +3099,24 @@ def test_shift(self):
30993099
actual = arr.shift(x=offset)
31003100
assert_identical(expected, actual)
31013101

3102-
def test_roll(self):
3102+
def test_roll_coords(self):
31033103
arr = DataArray([1, 2, 3], coords={'x': range(3)}, dims='x')
3104-
actual = arr.roll(x=1)
3104+
actual = arr.roll(x=1, roll_coords=True)
3105+
expected = DataArray([3, 1, 2], coords=[('x', [2, 0, 1])])
3106+
assert_identical(expected, actual)
3107+
3108+
def test_roll_no_coords(self):
3109+
arr = DataArray([1, 2, 3], coords={'x': range(3)}, dims='x')
3110+
actual = arr.roll(x=1, roll_coords=False)
3111+
expected = DataArray([3, 1, 2], coords=[('x', [0, 1, 2])])
3112+
assert_identical(expected, actual)
3113+
3114+
def test_roll_coords_none(self):
3115+
arr = DataArray([1, 2, 3], coords={'x': range(3)}, dims='x')
3116+
3117+
with pytest.warns(FutureWarning):
3118+
actual = arr.roll(x=1, roll_coords=None)
3119+
31053120
expected = DataArray([3, 1, 2], coords=[('x', [2, 0, 1])])
31063121
assert_identical(expected, actual)
31073122

xarray/tests/test_dataset.py

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3862,18 +3862,42 @@ def test_shift(self):
38623862
with raises_regex(ValueError, 'dimensions'):
38633863
ds.shift(foo=123)
38643864

3865-
def test_roll(self):
3865+
def test_roll_coords(self):
38663866
coords = {'bar': ('x', list('abc')), 'x': [-4, 3, 2]}
38673867
attrs = {'meta': 'data'}
38683868
ds = Dataset({'foo': ('x', [1, 2, 3])}, coords, attrs)
3869-
actual = ds.roll(x=1)
3869+
actual = ds.roll(x=1, roll_coords=True)
38703870

38713871
ex_coords = {'bar': ('x', list('cab')), 'x': [2, -4, 3]}
38723872
expected = Dataset({'foo': ('x', [3, 1, 2])}, ex_coords, attrs)
38733873
assert_identical(expected, actual)
38743874

38753875
with raises_regex(ValueError, 'dimensions'):
3876-
ds.roll(foo=123)
3876+
ds.roll(foo=123, roll_coords=True)
3877+
3878+
def test_roll_no_coords(self):
3879+
coords = {'bar': ('x', list('abc')), 'x': [-4, 3, 2]}
3880+
attrs = {'meta': 'data'}
3881+
ds = Dataset({'foo': ('x', [1, 2, 3])}, coords, attrs)
3882+
actual = ds.roll(x=1, roll_coords=False)
3883+
3884+
expected = Dataset({'foo': ('x', [3, 1, 2])}, coords, attrs)
3885+
assert_identical(expected, actual)
3886+
3887+
with raises_regex(ValueError, 'dimensions'):
3888+
ds.roll(abc=321, roll_coords=False)
3889+
3890+
def test_roll_coords_none(self):
3891+
coords = {'bar': ('x', list('abc')), 'x': [-4, 3, 2]}
3892+
attrs = {'meta': 'data'}
3893+
ds = Dataset({'foo': ('x', [1, 2, 3])}, coords, attrs)
3894+
3895+
with pytest.warns(FutureWarning):
3896+
actual = ds.roll(x=1, roll_coords=None)
3897+
3898+
ex_coords = {'bar': ('x', list('cab')), 'x': [2, -4, 3]}
3899+
expected = Dataset({'foo': ('x', [3, 1, 2])}, ex_coords, attrs)
3900+
assert_identical(expected, actual)
38773901

38783902
def test_real_and_imag(self):
38793903
attrs = {'foo': 'bar'}

0 commit comments

Comments
 (0)