[Numpy-discussion] use index array of len n to select columns of n x m array

[Martin Spacek]

On 2010-08-06 13:11, Martin Spacek wrote:
> Josef, I'd forgotten you could use None to increase the dimensionality of an
> array. Neat. And, somehow, it's almost twice as fast as the Cython version!:
>
>   >>>  timeit a[np.arange(a.shape[0])[:, None], i]
> 100000 loops, best of 3: 5.76 us per loop

I just realized why the Cython version was slower - the two assertion lines. 
commenting those out:

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True) # might be necessary to release the GIL?
def rowtake_cy(np.ndarray[np.int32_t, ndim=2] a,
                np.ndarray[np.int32_t, ndim=2] i):
     """For each row in a, return values according to column indices in the
     corresponding row in i. Returned shape == i.shape"""

     cdef Py_ssize_t nrows, ncols, rowi, coli
     cdef np.ndarray[np.int32_t, ndim=2] out

     nrows = i.shape[0]
     ncols = i.shape[1] # num cols to take for each row
     #assert a.shape[0] == nrows
     #assert i.max() < a.shape[1]
     out = np.empty((nrows, ncols), dtype=np.int32)

     for rowi in range(nrows):
         for coli in range(ncols):
             out[rowi, coli] = a[rowi, i[rowi, coli]]

     return out


gives me:


 >>> timeit rowtake_cy(a, i)
1000000 loops, best of 3: 1.44 us per loop

which is 4X faster than the a[np.arange(a.shape[0])[:, None], i] broadcasting 
method.

Martin