numpy set element in 3d array to nodata if nodata occurs anywhere

Question

I have an image, expressed as a numpy array. Its dimensions are 3000,3000 and there are 21 layers/bands. Thus my array is expressed as (21, 3000, 3000). I am using -9999 as a nodata value. However, where I have -9999 in one cell (pixel) in one band, I may have valid data in another. I would like to set all values in the stack to -9999 if a single -9999 occurs. For example, given my data is:

a = np.array([[[1,-9999,3,4],[1,2,3,4]],
              [[1,2,3,4],[1,2,-9999,4]],
              [[-9999,2,3,4],[1,2,3,4]]])

I would like to end up at:

a = np.array([[[-9999,-9999,3,4],[1,2,-9999,4]],
              [[-9999,-9999,3,4],[1,2,-9999,4]],
              [[-9999,-9999,3,4],[1,2,-9999,4]]])

where the occurrence of a -9999 in one band, would cause -9999 to be placed throughout. What its the most pythonic way of doing this?

Answer 1

You can find the columns that contains a -9999 using np.any and then set them:

a[:, np.any(a == -9999, axis=0)] = -9999