Fast Alternative To Conditionally Set An Array Elements

I have two given 3d arrays x_dist and y_dist each of shape (36,50,50). Elements in x_dist and y_dist are of type np.float32 and can be positive, negative or zero. I need to create

Solution 1:

Numba JIT can be used to do that efficiently. Here is an implementation:

@njit
def fastImpl(x_dist, y_dist):
    res_array = np.empty(x_dist.shape)
    for z in range(res_array.shape[0]):
        for y in range(res_array.shape[1]):
            for x in range(res_array.shape[2]):
                xDist = x_dist[z,y,x]
                yDist = y_dist[z,y,x]
                if xDist > 0.0 and yDist <= (1.0 + xDist):
                    res_array[z,y,x] = (1.0 - yDist) * xDist
    return res_array

Here are performance results on random input matrices:

Original implementation: 494 µs ± 6.23 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)
New implementation: 37.8 µs ± 236 ns per loop (mean ± std. dev. of 7 runs, 500 loops each)

The new implementation is about 13 times faster (without taking into account the compilation/warm-up time).