Search code examples
pythonnumpyvectorizationnested-loops

Vectorization on nested loop


I need to vectorize the following program :

y = np.empty((100, 100, 3))  
x = np.empty((300,)) 
for i in xrange(y.shape[0]):
    for j in xrange(y.shape[1]):
        y[i, j, 0] = x[y[i, j, 0]]

Of course, in my example, we suppose that y[:, :, :]<=299 Vectorization, as far as I know, can't simply work here as we are using the native python indexing on lists ...

I've heard of np.apply_along_axis, but it doesn't work on this special case, or may I missed something ?

Thank you very much for any help.


Solution

  • np.apply_along_axis could work, but it's overkill.

    First, there's a problem in your nested loop approach. np.empty, used to define y, returns an array of np.float values, which cannot be used to index an array. To take care of this, you have to cast the array as integers, e.g. y = np.empty((100, 100, 3)).astype(np.int).

    Once you do that, you can index using y, as follows:

    y = np.empty((100, 100, 3)).astype(np.uint8)
    x = np.empty((300,))
    y[:,:,0] = x[y[:,:,0]]
    

    Of course, y is all 0's, so it's not quite clear what this accomplishes.