Filter a numpy ndarray using another numpy ndarray

I have two numpy ndarrays of the same shape (15081, 56724, 3, 3). What I want to do is as follows:

Say we have a cross section of the first array, array1[1, 1, :, :], looks like this:

[[120, 110, 220],
 [ 85,  99,  72],
 [197,  80,  75]]

I want to convert it to a boolean in a way that the max of each row is True and the rest is False. In the whole array, this corresponds to axis=3. So the array looks like this after conversion:

[[False, False,  True],
 [False,  True, False],
 [ True, False, False]]

Now I want to filter the other array, array2, using this boolean array to have something that looks like below. I only want to keeping those values of array2 that correspond to True in array1 and set the rest to zero.

[[  0,   0,  65],
 [  0, 179,   0],
 [125,   0,   0]]

I can do this using a loop but it takes an age (even more). I expect something like numpy.where(array1.is_max(axis=3), True, False), but there is no function like is_max in python, besides this way the axis 3 is collapsed and I cannot filter array2 using array1.

Answer

In numpy, is_max is approximately argmax:

indices = array1.argmax(axis=-1)

Instead of a mask, this will give you a linear index, so you can do

array2b = np.zeros_like(array2)
index = np.indices(array2.shape[:-1], sparse=True) + (indices,)
array2b[index] = array1[index]

You can do something similar with np.take_along_axis and np.put_along_axis:

indices = np.expand_dims(array1.argmax(-1), array1.ndim - 1)
array2b = np.zeros_like(array2)
np.put_along_axis(array2b, indices, np.take_along_axis(array2, indices, -1), -1)

If you want a mask-based approach, you can create the mask like this:

mask = array1 != array1.max(-1, keepdims=True)

Now you can set all elements to zero directly:

array2[mask] = 0

Alternatively you can do something like

mask = array1 == array1.max(-1, keepdims=True)
array2 *= mask

Update

From your description in the comments, you are looking for a different operation entirely. You can start by thresholding array1 (which I assume represents the difference between the blurred and original image):

mask = array1 >= 100  # find a threshold that works
array2 *= mask

OR

mask = array1 < 100
array2[mask] = 0

You may also be looking for local minima in an image. You can get those by finding pixels that are bigger than their surroundings. To do that, run a non-linear filter on the image, like scipy.ndimage.maximum_filter:

 mask = array1 == scipy.ndimage.maximum_filter(array1, 5)
 array2 *= mask