Extend numpy mask by n cells to the right for each bad value, efficiently

Question

Let\'s say I have a length 30 array with 4 bad values in it. I want to create a mask for those bad values, but since I will be using rolling window functions, I\'d also like a f

Answer 1

Something like this?

maskleft = np.where(np.isnan(a))[0]
maskright = maskleft + n
mask = np.zeros(len(a),dtype=bool)
for l,r in itertools.izip(maskleft,maskright): 
   mask[l:r] = True

Or, since n is small, it might be better to loop over it instead:

maskleft = np.where(np.isnan(a))[0]
mask = np.zeros(len(a),dtype=bool)
for i in range(0,n):
  thismask = maskleft+i
  mask[thismask] = True

Except for the loop over n, the above is fully vectorized. But the loop is fully parallelizable, so you could be able to get a factor-n speedup using e.g. multiprocessing or Cython, if you're willing to go to the trouble.

Edit: per @askewchan solution 2 can potentially cause out of range errors. It also has indexing problems in the range(0,n). Possible correction:

maskleft = np.where(np.isnan(a))[0]
ghost_mask = np.zeros(len(a)+n,dtype=bool)
for i in range(0, n+1):
    thismask = maskleft + i
    ghost_mask[thismask] = True
mask = ghost_mask[:len(ghost_mask)-n]