Generalized version of np.roll

Question

I have a 2D array

a = np.array([[0,1,2,3],[4,5,6,7]])

that is a 2x4 array. I need to shift the elements of each of the two arrays in axis 0 in but with different steps, say 1 for the first and 2 for the second, so that the output will be

np.array([[1,2,3,0],[6,7,4,5]])

With np.roll it doesn't seem possible to do it, at least looking at the documentation, I don't see any useful hint. There exists another function doing this?

Does this answer your question? [Roll rows of a matrix independently](https://stackoverflow.com/questions/20360675/roll-rows-of-a-matrix-independently) — Michael Szczesny, Aug 29 '22 at 23:22
@Domenico I think there may be a mistake in your question. Did you mean *axis 0*? — Michael Sohnen, Aug 29 '22 at 23:44
`roll` basically does a couple of sliced copies. `result[:n] = arr[-n:]` etc, generalized by axis. Nothing exotic or specially compiled. — hpaulj, Aug 30 '22 at 00:01
Besides the mentioned link by [michael-szczesny comment](https://stackoverflow.com/users/14277722/michael-szczesny), such command can be rewrite using loops and be jitted by numba to get better efficiency if it could (see https://stackoverflow.com/a/72917320/13394817). — Ali_Sh, Aug 30 '22 at 00:12
@MichaelSzczesny partially yes, I didn't specify to avoid for loops as the script would be slow with big arryas. — Domenico, Aug 30 '22 at 05:42
@Domenico - Then the answer is: No, there isn't. `np.roll` is implemented with [python loops](https://github.com/numpy/numpy/blob/v1.23.0/numpy/core/numeric.py#L1223). — Michael Szczesny, Aug 30 '22 at 05:48

score 0 · Answer 1 · answered Aug 30 '22 at 01:33

This is an attempt at a generalized version of numpy.roll.

import numpy as np

a = np.array([[0,1,2,3],[4,5,6,7]])

def roll(a, shifts, axis):
    assert a.shape[axis] == len(shifts)
    return np.stack([
        np.roll(np.take(a, i, axis), shifts[i]) for i in range(len(shifts))
    ], axis)

print(a)
print(roll(a, [-1, -2], 0))
print(roll(a, [1, 2, 1, 0], 1))

prints

[[0 1 2 3]
 [4 5 6 7]]

[[1 2 3 0]
 [6 7 4 5]]

[[4 1 6 3]
 [0 5 2 7]]

Here, the parameter a is a numpy.array, shifts is an Iterable containing the shift amounts per element and axis is the axis along which to shift. Note that was only tested on two-dimensional arrays however.

Generalized version of np.roll

1 Answers1