implemented LArray.roll

Author: gdementen

doc/source/api.rst

@@ -444,6 +444,7 @@ Miscellaneous
    LArray.divnot0
    LArray.clip
    LArray.shift
+   LArray.roll
    LArray.diff
    LArray.to_clipboard
 

doc/source/changes/version_0_30.rst.inc

@@ -132,8 +132,13 @@ New features
 * implemented :py:obj:`wrap_elementwise_array_func()` function to make a function defined in another library work with
   LArray arguments instead of with numpy arrays.
 
+* implemented :py:obj:`LArray.roll()` to roll the cells of an array n-times to the right along an axis. This is similar
+  to :py:obj:`LArray.shift()`, except that cells which are pushed "outside of the axis" are reintroduced on the opposite
+  side of the axis instead of being dropped.
+
 * implemented :py:obj:`AxisCollection.rename()` to rename axes of an AxisCollection, independently of any array.
 
+
 Miscellaneous improvements
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 

larray/core/array.py

@@ -6853,6 +6853,11 @@ def shift(self, axis, n=1):
         -------
         LArray
 
+        See Also
+        --------
+        LArray.roll : cells which are pushed "outside of the axis" are reintroduced on the opposite side of the axis
+                      instead of being dropped.
+
         Examples
         --------
         >>> arr = ndtest('sex=M,F;year=2019..2021')
@@ -6877,6 +6882,63 @@ def shift(self, axis, n=1):
         else:
             return self[:]
 
+    def roll(self, axis=None, n=1):
+        r"""Rolls the cells of the array n-times to the right along axis. Cells which would be pushed "outside of the
+        axis" are reintroduced on the opposite side of the axis.
+
+        Parameters
+        ----------
+        axis : int, str or Axis, optional
+            Axis along which to roll. Defaults to None (all axes).
+        n : int or LArray, optional
+            Number of positions to roll. Defaults to 1. Use a negative integers to roll left.
+            If n is an LArray the number of positions rolled can vary along the axes of n.
+
+        Returns
+        -------
+        LArray
+
+        See Also
+        --------
+        LArray.shift : cells which are pushed "outside of the axis" are dropped instead of being reintroduced on the
+                       opposite side of the axis.
+
+        Examples
+        --------
+        >>> arr = ndtest('sex=M,F;year=2019..2021')
+        >>> arr
+        sex\year  2019  2020  2021
+               M     0     1     2
+               F     3     4     5
+        >>> arr.roll('year')
+        sex\year  2019  2020  2021
+               M     2     0     1
+               F     5     3     4
+
+        One can also roll by a different amount depending on another axis
+
+        >>> # let us roll by 1 for men and by 2 for women
+        >>> n = sequence(arr.sex, initial=1)
+        >>> n
+        sex  M  F
+             1  2
+        >>> arr.roll('year', n)
+        sex\year  2019  2020  2021
+               M     2     0     1
+               F     4     5     3
+        """
+        if isinstance(n, (int, np.integer)):
+            axis_idx = None if axis is None else self.axes.index(axis)
+            return LArray(np.roll(self.data, n, axis=axis_idx), self.axes)
+        else:
+            if not isinstance(n, LArray):
+                raise TypeError("n should either be an integer or an LArray")
+            if axis is None:
+                raise TypeError("axis may not be None if n is an LArray")
+            axis = self.axes[axis]
+            seq = sequence(axis)
+            return self[axis.i[(seq - n) % len(axis)]]
+
     # TODO: add support for groups as axis (like aggregates)
     # eg a.diff(x.year[2018:]) instead of a[2018:].diff(x.year)
     def diff(self, axis=-1, d=1, n=1, label='upper'):