Roll rows of a matrix independently

Question

I have a matrix (2d numpy ndarray, to be precise):

A = np.array([[4, 0, 0],
              [1, 2, 3],
              [0, 0, 5]])

And I want to roll each row of A independently, according to roll values in another array:

r = np.array([2, 0, -1])

That is, I want to do this:

print np.array([np.roll(row, x) for row,x in zip(A, r)])

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

Is there a way to do this efficiently? Perhaps using fancy indexing tricks?

Answer 1

You can do it using advanced indexing. Whether or not it is the fastest way likely depends on the array size. For instance, for large rows, this may be slower than other methods.

rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]

# Always use a negative shift, so that column_indices are valid.
# Alternative: r %= A.shape[1]
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:, np.newaxis]

result = A[rows, column_indices]

Answer 2

By using a fast fourrier transform we can apply a transformation in the frequency domain and then use the inverse fast fourrier transform to obtain the row shift.

So this is a pure numpy solution that take only one line:

import numpy as np
from numpy.fft import fft, ifft

# The row shift function using the fast fourrier transform
#   rshift(A,r) where A is a 2D array, r the row shift vector
def rshift(A,r):
     return np.real(ifft(fft(A,axis=1)*np.exp(2*1j*np.pi/A.shape[1]*r[:,None]*np.r_[0:A.shape[1]][None,:]),axis=1).round())

This will apply a left shift, but we can simply negate the exponential exponant to turn the function into a right shift function:

ifft(fft(...)*np.exp(-2*1j...)

It can be used like that:

# Example:

A = np.array([[1,2,3,4],
              [1,2,3,4],
              [1,2,3,4]])

r = np.array([1,-1,3])

print(rshift(A,r))

Answer 3

I implement a pure numpy.lib.stride_tricks.as_strided solution as follows

from numpy.lib.stride_tricks import as_strided

def custom_roll(arr, r_tup):
    m = np.asarray(r_tup)
    arr_roll = arr[:, [*range(arr.shape[1]),*range(arr.shape[1]-1)]].copy() #need `copy`
    strd_0, strd_1 = arr_roll.strides
    n = arr.shape[1]
    result = as_strided(arr_roll, (*arr.shape, n), (strd_0 ,strd_1, strd_1))

    return result[np.arange(arr.shape[0]), (n-m)%n]

A = np.array([[4, 0, 0],
              [1, 2, 3],
              [0, 0, 5]])

r = np.array([2, 0, -1])

out = custom_roll(A, r)

Out[789]:
array([[0, 0, 4],
       [1, 2, 3],
       [0, 5, 0]])

Answer 4

Building on divakar's excellent answer, you can apply this logic to 3D array easily (which was the problematic that brought me here in the first place). Here's an example - basically flatten your data, roll it & reshape it after::

def applyroll_30(cube, threshold=25, offset=500):
    flattened_cube = cube.copy().reshape(cube.shape[0]*cube.shape[1], cube.shape[2])

    roll_matrix = calc_roll_matrix_flattened(flattened_cube, threshold, offset)

    rolled_cube = strided_indexing_roll(flattened_cube, roll_matrix, cube_shape=cube.shape)

    rolled_cube = triggered_cube.reshape(cube.shape[0], cube.shape[1], cube.shape[2])
    return rolled_cube


def calc_roll_matrix_flattened(cube_flattened, threshold, offset):
    """ Calculates the number of position along time axis we need to shift
        elements in order to trig the data.
        We return a 1D numpy array of shape (X*Y, time) elements
    """

    # armax(...) finds the position in the cube (3d) where we are above threshold
    roll_matrix = np.argmax(cube_flattened > threshold, axis=1) + offset
    # ensure we don't have index out of bound
    roll_matrix[roll_matrix>cube_flattened.shape[1]] = cube_flattened.shape[1]          
    return roll_matrix



def strided_indexing_roll(cube_flattened, roll_matrix_flattened, cube_shape):
    # Concatenate with sliced to cover all rolls
    # otherwise we shift in the wrong direction for my application
    roll_matrix_flattened = -1 * roll_matrix_flattened

    a_ext = np.concatenate((cube_flattened, cube_flattened[:, :-1]), axis=1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = cube_flattened.shape[1]
    result = viewW(a_ext,(1,n))[np.arange(len(roll_matrix_flattened)), (n - roll_matrix_flattened) % n, 0]
    result = result.reshape(cube_shape)
    return result

Divakar's answer doesn't do justice to how much more efficient this is on large cube of data. I've timed it on a 400x400x2000 data formatted as int8. An equivalent for-loop does ~5.5seconds, Seberg's answer ~3.0seconds and strided_indexing.... ~0.5second.

Answer 5

In case you want more general solution (dealing with any shape and with any axis), I modified @seberg's solution:

def indep_roll(arr, shifts, axis=1):
    """Apply an independent roll for each dimensions of a single axis.

    Parameters
    ----------
    arr : np.ndarray
        Array of any shape.

    shifts : np.ndarray
        How many shifting to use for each dimension. Shape: `(arr.shape[axis],)`.

    axis : int
        Axis along which elements are shifted. 
    """
    arr = np.swapaxes(arr,axis,-1)
    all_idcs = np.ogrid[[slice(0,n) for n in arr.shape]]

    # Convert to a positive shift
    shifts[shifts < 0] += arr.shape[-1] 
    all_idcs[-1] = all_idcs[-1] - shifts[:, np.newaxis]

    result = arr[tuple(all_idcs)]
    arr = np.swapaxes(result,-1,axis)
    return arr

Answer 6

numpy.lib.stride_tricks.as_strided stricks (abbrev pun intended) again!

Speaking of fancy indexing tricks, there's the infamous - np.lib.stride_tricks.as_strided. The idea/trick would be to get a sliced portion starting from the first column until the second last one and concatenate at the end. This ensures that we can stride in the forward direction as needed to leverage np.lib.stride_tricks.as_strided and thus avoid the need of actually rolling back. That's the whole idea!

Now, in terms of actual implementation we would use scikit-image's view_as_windows to elegantly use np.lib.stride_tricks.as_strided under the hoods. Thus, the final implementation would be -

from skimage.util.shape import view_as_windows as viewW

def strided_indexing_roll(a, r):
    # Concatenate with sliced to cover all rolls
    a_ext = np.concatenate((a,a[:,:-1]),axis=1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[1]
    return viewW(a_ext,(1,n))[np.arange(len(r)), (n-r)%n,0]

Here's a sample run -

In [327]: A = np.array([[4, 0, 0],
     ...:               [1, 2, 3],
     ...:               [0, 0, 5]])

In [328]: r = np.array([2, 0, -1])

In [329]: strided_indexing_roll(A, r)
Out[329]: 
array([[0, 0, 4],
       [1, 2, 3],
       [0, 5, 0]])

Benchmarking

# @seberg's solution
def advindexing_roll(A, r):
    rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]    
    r[r < 0] += A.shape[1]
    column_indices = column_indices - r[:,np.newaxis]
    return A[rows, column_indices]

Let's do some benchmarking on an array with large number of rows and columns -

In [324]: np.random.seed(0)
     ...: a = np.random.rand(10000,1000)
     ...: r = np.random.randint(-1000,1000,(10000))

# @seberg's solution
In [325]: %timeit advindexing_roll(a, r)
10 loops, best of 3: 71.3 ms per loop

#  Solution from this post
In [326]: %timeit strided_indexing_roll(a, r)
10 loops, best of 3: 44 ms per loop