How to replace a double for-loop with mask and indexing?

Question

I have a couple of nested for-loops that do the right thing (masked copy of array). However the performance is too slow, and I feel there must be a better Pythonic way to do this. The goal is use the mask to determine when to copy data from the source, using coord as an index into source. The looping code that works is below:

import numpy as np
dest = np.zeros((4,4,2))
source = range(32)
source = np.reshape(source,(4,4,2))
mask = np.ones((4,4),bool)
mask[1,0] = 0
coord = np.ones((4,4,2),int)

for y in range (0,dest.shape[0]):
    for x in range (0, dest.shape[1]):
        if np.all(mask[y,x]):
            dest[y,x] = source[coord[y,x,0], coord[y,x,1]]

print dest

After running, dest looks like this:

[[[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[  0.   0.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]]

source[1,1] is copied to all of dest, except for dest[1,0] because mask[1,0] is set to False. The rest of mask is True. Can anyone please show me how replace the loops with something more efficient?

Answer 1

Use numpy.where. You have to add an extra dimension to mask so it will broadcast.

dest = np.where(mask[:,:,None], source[coord[:,:,0], coord[:,:,1]], dest)

Or if appropriate:

dest = np.where(mask[:,:,None], source[coord[:,:,0], coord[:,:,1]], np.zeros((4,4,2)))