slice matrix in Numpy (or Theano)

Question

Is there an optimal way to slice each row of a matrix in Numpy (or Theano) by a stride of N, given start index for each column?

For example, in the matrix A below, starting slice index for each row is given in the first column, and for a row i, I want to have A[i, A[0]:A[0]+stride]

A = [[1,  1,  2,  3,  4, 5, 6],
     [1,  11, 12, 13, 14, 15, 16],
     [3,  22, 23, 24, 25, 26, 27]]
stride = 2
Desired output:
[[  1.   2.   3.]
 [ 11.  12.  13.]
 [ 24.  25.  26.]]

I tried the code below:

b = [range(A.shape[0]), A[:, 0]]
c = [range(A.shape[0]), A[:, 0] + stride]
A[b:c]

but I got the following error:

IndexError: failed to coerce slice entry of type list to integer

Answer 1

Here's a vectorized approach making use of broadcasting to get those indices for indexing into columns across each row and then using NumPy's advanced-indexing to extract out those elements along each row in a vectorized manner -

idx = A[:,0,None] + np.arange(stride+1)
out = A[np.arange(idx.shape[0])[:,None], idx]

Sample run -

In [273]: A
Out[273]: 
array([[ 1,  1,  2,  3,  4,  5,  6],
       [ 1, 11, 12, 13, 14, 15, 16],
       [ 3, 22, 23, 24, 25, 26, 27]])

In [274]: idx = A[:,0,None] + np.arange(stride+1)

In [275]: idx
Out[275]: 
array([[1, 2, 3],
       [1, 2, 3],
       [3, 4, 5]])

In [276]: A[np.arange(idx.shape[0])[:,None], idx]
Out[276]: 
array([[ 1,  2,  3],
       [11, 12, 13],
       [24, 25, 26]])