CODEWITHSUNDEEP
pythonnumpy
BayerSe
BayerSe

Reputation: 1171

Speeding up Numpy array slicing

I need to slice a moderately sized 2d Numpy array along two dimensions. As example,

import numpy as np
X = np.random.normal(loc=0, scale=1, size=(3000, 100))

From this array, I need to select a large number of rows and a rather small number of columns, e.g.

row_idx = np.random.random_integers(0, 2999, 2500)
col_idx = np.random.random_integers(0, 99, 10)

Right now, I do this by the following command:

X.take(col_idx, axis=1).take(row_idx, axis=0)

This takes approximately 115µs on my computer. The problem is that I need to do this step several million times per run.

Do you see any chance to speed this up?

Edit (additional information): I have a matrix X which is nxk. The n rows contain 1xk vectors. There are three sets: an active set (V), a left set (L) and a right set (R). Moreover, there are coefficients v0 and v. I need to compute this quantity: http://goo.gl/KNoSl3 (sorry, I can't post images). The formula from the question selects all rows of X which are in the left (right) set and all columns that are in the active set.

Edit 2

I found another small improvement.

X.take(col_idx, axis=1, mode='clip').take(row_idx, axis=0, mode='clip')

is a little faster (roughly 25% on my machine).

Upvotes: 2

Views: 781

Answers (2)

Daniel
Daniel

Reputation: 19547

Lets do something where we make a 1D array of indices that satisfy our conditions for a n-dimensional grid.

def make_multi_index(arr, *inds):
    tmp = np.meshgrid(*inds, indexing='ij')
    idx = np.vstack([x.ravel() for x in tmp])
    return np.ravel_multi_index(idx, X.shape)

Using your test arrays and the original case for reference:

%timeit X.take(col_idx, axis=1).take(row_idx, axis=0)
10000 loops, best of 3: 95.4 µs per loop

Lets use this functional to build indices, hold them, and then use take to return your desired output:

inds = make_multi_index(X, row_idx, col_idx)
tmp = np.take(X,inds).reshape(row_idx.shape[0], col_idx.shape[0])

np.allclose(tmp, X.take(col_idx, axis=1).take(row_idx, axis=0))
Out[128]: True

So building our indices and keeping them appears to work, now for timings:

%timeit make_multi_index(X, row_idx, col_idx)
1000 loops, best of 3: 356 µs per loop

%timeit np.take(X,inds).reshape(row_idx.shape[0], col_idx.shape[0])
10000 loops, best of 3: 59.9 µs per loop

So as it happens not terribly great- this probably gets better with more dimension that you would like to take from. Anyhow if you keep these indices for more then 10-15 iterations it could help some or if you add an additional dimension and take all inactive datasets simultaneously.

Upvotes: 1

perimosocordiae
perimosocordiae

Reputation: 17847

You could use two-dimensional fancy indexing:

X[row_idx,col_idx[:,None]]

But that takes ~1ms on my machine, vs ~300us using your method.

It seems like your method is the best you can do, unless you have additional information about the values in row_idx and col_idx.

Upvotes: 0

Related Questions