1

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?

smci
  • 32,567
  • 20
  • 113
  • 146
MJK
  • 13
  • 2
  • Are you looking to handle different masks or just the one you have? If you're just doing this one, you could straight up copy without the mask. And (mostly) unrelated, but it would help with the readability of your code if you followed the conventions of [PEP-8](https://www.python.org/dev/peps/pep-0008/) or at least used spaces between function parameters and before and after `=`. If you refactored into multiple functions, your code would also be a good bit easier to understand. – Daniel Underwood Nov 18 '16 at 03:32

1 Answers1

3

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)))
wwii
  • 23,232
  • 7
  • 37
  • 77
  • That is exactly what I was looking for. I tried a variety of things, but it was the extra dimension on mask that I was missing. – MJK Nov 18 '16 at 13:04