Roll first column by 1, second column by 2, etc-CodePudding

I have an array in numpy. I want to roll the first column by 1, second column by 2, etc.

Here is an example.

>>> x = np.reshape(np.arange(15), (5, 3))
>>> x
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

What I want to do:

>>> y = roll(x)
>>> y
array([[12, 10,  8],
       [ 0, 13, 11],
       [ 3,  1, 14],
       [ 6,  4,  2],
       [ 9,  7,  5]])

What is the best way to do it?

The real array will be very big. I'm using cupy, the GPU version of numpy. I will prefer solution fastest on GPU, but of course, any idea is welcomed.

CodePudding user response：

I implemented a naive solution (roll_for) and compares it to @Chrysophylaxs 's solution (roll_indexing).

Note that roll_indexing does not work if the data has more columns than rows, but I guess it could be adapted.

Implementations:

import numpy as np

def roll_for(x, shifts=None, axis=-1):
    if shifts is None:
        shifts = np.arange(1, x.shape[axis]   1)  # OP requirement
    xt = x.swapaxes(axis, 0)  # https://stackoverflow.com/a/31094758/13636407
    yt = np.empty_like(xt)
    for idx, shift in enumerate(shifts):
        yt[idx] = np.roll(xt[idx], shift=shift)
    return yt.swapaxes(0, axis)


def roll_indexing(x):
    h, w = x.shape
    rows, cols = np.arange(h), np.arange(w)
    offsets = cols   1
    shifted = np.subtract.outer(rows, offsets)
    return x[shifted, cols]

Tests:

M, N = 5, 3
x = np.arange(M * N).reshape(M, N)
expected = np.array([[12, 10, 8], [0, 13, 11], [3, 1, 14], [6, 4, 2], [9, 7, 5]])

assert np.array_equal(roll_for(x), expected)
assert np.array_equal(roll_indexing(x), expected)

M, N = 100, 200
# roll_indexing does not work when M < N

Benchmark:

M, N = 100, 100
x = np.arange(M * N).reshape(M, N)
assert np.array_equal(roll_for(x), roll_indexing(x))
%timeit roll_for(x)       # 833 µs ± 2.23 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit roll_indexing(x)  # 49.1 µs ± 111 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

M, N = 1_000, 1_000
x = np.arange(M * N).reshape(M, N)
assert np.array_equal(roll_for(x), roll_indexing(x))
%timeit roll_for(x)       # 12.4 ms ± 72.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit roll_indexing(x)  # 9.12 ms ± 25.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

M, N = 10_000, 10_000
x = np.arange(M * N).reshape(M, N)
assert np.array_equal(roll_for(x), roll_indexing(x))
%timeit roll_for(x)       # 1.42 s ± 94.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit roll_indexing(x)  # 1.24 s ± 17.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

CodePudding user response：

You could use advanced indexing:

import numpy as np

x = np.reshape(np.arange(15), (5, 3))

h, w = x.shape

rows, cols = np.arange(h), np.arange(w)
offsets = cols   1
shifted = np.subtract.outer(rows, offsets) % h

y = x[shifted, cols]

array([[12, 10,  8],
       [ 0, 13, 11],
       [ 3,  1, 14],
       [ 6,  4,  2],
       [ 9,  7,  5]])