Finding min values in a 3D array across one axis and replacing the non-corresponding values in another 3D array with 0 without loops in python

Question

Let's say we have two 3D arrays, A(x,y,z) and B(x,y,z) that x,y,z are dimensions. I want to identify all the minimum values across the z-axis in the A array and then based on those values and their indices choose the corresponding values in the B, keep them and replace other values with zero.

Answer 1

You can think of it a little differently. Finding the locations of the minima in A is straightforward:

ind = np.expand_dims(np.argmin(A, axis=2), axis=2)

You can do one of the following things:

1. Simplest: create a replacement of B and populate the relevant elements:

    C = np.zeros_like(B)
    np.put_along_axis(C, ind, np.take_along_axis(B, ind, 2), 2)
   

2. Same thing, but in-place:

    values = np.take_along_axis(B, ind, 2)
    B[:] = 0
    np.put_along_axis(B, ind, values, 2)
   

3. Convert the index to a mask:

    mask = np.ones(B.shape, dtype=bool)
    np.put_along_axis(mask, ind, False, 2)
    B[mask] = 0
   

You can replace the calls to take_along_axis and put_along_axis with suitable indexing expressions. In particular:

indx, indy = np.indices(A.shape[:-1])
indz = np.argmin(A, axis=-1)

The examples above then transform into

1. New array:

    C = np.zeros_like(B)
    C[indx, indy, indz] = B[indx, indy, indz]
   

2. In-place:

    values = B[indx, indy, indz]
    B[:] = 0
    B[indx, indy, indz] = values
   

3. Masked:

    mask = np.ones(B.shape, dtype=bool)
    mask[indx, indy, indz] = False
    B[mask] = 0