Speed up loops over numpy arrays that modify the array

Question

I have numpy arrays with the shape (X, Y, 4), where X and Y are large. I want to do the following operation on axis 2 of the array (i.e., A[x, y] for each x and y): for each 1-dimensional vector A[x, y], if each of the first 3 elements is within some distance d to the corresponding element in constant vector C, then set the 4th element to 0. This is how I do it with Python loops:

for i in range(X):
    for j in range(Y):
        if np.all(np.abs(A[i, j, 0:3] - C) <= d):
            A[i, j, 3] = 0

These loops are very slow, obviously. I suspect there's a faster, ufunc-based way of doing this operation, but for the life of me can't figure out how to do it. Does anyone have any suggestions?

(Note on application: this is to add alpha channel to an image so that colors close to C become transparent.)

Answer 1

Try:

mask = (np.abs(A[:,:,:3]-C)<=d).all(-1)
A[mask,3] = 0