Python - slice array at different position on every row

Question

I have a 2D python array that I want to slice in an odd way - I want a constant width slice starting on a different position on every row. I would like to do this in a vectorised way if possible.

e.g. I have the array A=np.array([range(5), range(5)]) which looks like

array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])

I would like to slice this as follows: 2 elements from each row, starting at positions 0 and 3. The starting posiitons are stored in b=np.array([0,3]). Desired output is thus: np.array([[0,1],[3,4]]) i.e.

array([[0, 1],
       [3, 4]])

The obvious thing I tried to get this result was A[:,b:b+2] but that doesn't work, and I can't find anything that will.

Speed is important as this will operate on a largish array in a loop, and I don't want to bottleneck other parts of my code.

Answer 1

You can use np.take():

In [21]: slices = np.dstack([b, b+1])

In [22]: np.take(arr, slices)
Out[22]: 
array([[[0, 1],
        [3, 4]]])

Answer 2

Approach #1 : Here's one approach with broadcasting to get all indices and then using advanced-indexing to extract those -

def take_per_row(A, indx, num_elem=2):
    all_indx = indx[:,None] + np.arange(num_elem)
    return A[np.arange(all_indx.shape[0])[:,None], all_indx]

Sample run -

In [340]: A
Out[340]: 
array([[0, 5, 2, 6, 3, 7, 0, 0],
       [3, 2, 3, 1, 3, 1, 3, 7],
       [1, 7, 4, 0, 5, 1, 5, 4],
       [0, 8, 8, 6, 8, 6, 3, 1],
       [2, 5, 2, 5, 6, 7, 4, 3]])

In [341]: indx = np.array([0,3,1,5,2])

In [342]: take_per_row(A, indx)
Out[342]: 
array([[0, 5],
       [1, 3],
       [7, 4],
       [6, 3],
       [2, 5]])

Approach #2 : Using np.lib.stride_tricks.as_strided -

from numpy.lib.stride_tricks import as_strided

def take_per_row_strided(A, indx, num_elem=2):
    m,n = A.shape
    A.shape = (-1)
    s0 = A.strides[0]
    l_indx = indx + n*np.arange(len(indx))
    out = as_strided(A, (len(A)-num_elem+1, num_elem), (s0,s0))[l_indx]
    A.shape = m,n
    return out

Runtime test for taking 200 per row from a 2000x4000 matrix

In [447]: A = np.random.randint(0,9,(2000,4000))

In [448]: indx = np.random.randint(0,4000-200,(2000))

In [449]: out1 = take_per_row(A, indx, 200)

In [450]: out2 = take_per_row_strided(A, indx, 200)

In [451]: np.allclose(out1, out2)
Out[451]: True

In [452]: %timeit take_per_row(A, indx, 200)
100 loops, best of 3: 2.14 ms per loop

In [453]: %timeit take_per_row_strided(A, indx, 200)
1000 loops, best of 3: 435 µs per loop