I’d like to reduce a multidimensional array along one axis by finding the first value in the other dims that is truthy/matches some condition, or a default value if no matching elements can be found. So far I’m trying to do this for a 3D array with some simple looping as follows:

def first(arr, condition, default):
    out = default.copy()

    for u in range(arr.shape[1]):
        for e in range(arr.shape[2]):
            (idx,) = np.nonzero(condition[:, u, e])
            if len(idx):
                out[u, e] = arr[idx[0], u, e]

    return out

This is a bit ugly, not efficient, not vectorized, and not easily generalized to arrays of arbitrary dims though. Does numpy have a nicer, built-in way to achieve this?