You can use np.where to conditionally pick 0 or the result of your computation. We can also suppress NumPy warnings that are triggered by the divide by zero (we discard those values anyways, so the warning is safe to ignore).
An alternative is to mask all zero elements in a masked array, and a third way to index all valid elements first and do the computation on those elements only:
A = np.random.rand(300, 400)
A *= np.random.randint(2, size=A.shape)
def using_where(A):
with np.errstate(divide='ignore', invalid='ignore'):
return np.where(
A == 0,
0,
1 - (1/(2 * A)) * (1 - np.exp(-2 * A))
)
def using_ma(A):
mA = np.ma.masked_equal(A, 0)
mA = 1 - (1/(2 * mA)) * (1 - np.exp(-2 * mA))
return mA.filled(0)
def using_mask(A):
A = A.copy() # we are modifying A inplace. To make this function pure we need to work on a copy of A
mask = A != 0
A[mask] = 1 - (1/(2 * A[mask])) * (1 - np.exp(-2 * A[mask]))
return A
%timeit using_where(A)
# 5.51 ms ± 329 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit using_ma(A)
# 20.3 ms ± 508 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit using_mask(A)
# 6.61 ms ± 301 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
As you can see, the np.where approach is the fastest.