CODEWITHSUNDEEP
pythonpython-2.7numpytheano
csuo
csuo

Reputation: 820

slice matrix in Numpy (or Theano)

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

Upvotes: 3

Views: 345

Answers (2)

Divakar
Divakar

Reputation: 221574

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]])

Upvotes: 1

Andrew Che
Andrew Che

Reputation: 968

Not sure if it's optimal, but at least it doesn't throw any errors :)

Python 2.7.12 (default, Nov 19 2016, 06:48:10) 
[GCC 5.4.0 20160609] on linux2
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy
>>> a = numpy.array([[1,  1,  2,  3,  4, 5, 6],
...                  [1,  11, 12, 13, 14, 15, 16],
...                  [3,  22, 23, 24, 25, 26, 27]])
>>> stride = 2
>>> numpy.array(map(lambda row: row[row[0]:row[0] + stride + 1], a))
array([[ 1,  2,  3],
       [11, 12, 13],
       [24, 25, 26]])

Upvotes: 0

Related Questions