how to avoid division by zero in 2d numpy array when taking average?

Question

Let's say I have three arrays

A = np.array([[2,2,2],[1,0,0],[1,2,1]])
B = np.array([[2,0,2],[0,1,0],[1,2,1]])
C = np.array([[2,0,1],[0,1,0],[1,1,2]])
A,B,C
(array([[2, 2, 2],
        [1, 0, 0],
        [1, 2, 1]]),
 array([[2, 0, 2],
        [0, 1, 0],
        [1, 2, 1]]),
 array([[2, 0, 1],
        [0, 1, 0],
        [1, 1, 2]]))

when i take average of C/ (A+B), i get nan/inf value with RunTimeWarning.. the resultant array looks like the following.

np.average(C/(A+B), axis = 1)
array([0.25      ,        nan, 0.58333333])

I would like to change any inf/nan value to 0.

What I tried so far was

#doesn't work. ( maybe im doing this wrong..)
mask = A+B >0
np.average(C[mask]/(A[mask]+B[mask]), axis = 1)


#does not work and not an ideal solution.
avg = np.average(C/(A+B), axis = 1)
avg[avg == np.nan] =0

any help would be appreciated!

Answer 1

Your tried approaches are both a valid way of dealing with it, but you need to change them slightly.

1. Avoiding the division upfront, by only calculating the result where it's valid (eg non-zero):

The use of the boolean mask you defined makes the resulting arrays (after indexing) to become 1D. So using this would mean you have to allocate the resulting array upfront, and assign it using that same mask.

mask = A+B > 0
result = np.zeros_like(A, dtype=np.float32)
result[mask] = C[mask]/(A[mask]+B[mask])

It does require the averaging over the second dimension to be done separate, and also masking the incorrect result to zero for elements where the division could not be done due to the zeros.

result = result.mean(axis=1)
result[(~mask).any(axis=1)] = 0

To me the main benefit would be avoiding the warning from Numpy, and perhaps in the case of a large amount of zeros (in A+B) you could gain a little performance by avoiding that calculation all together. But overall it seems a lot of effort to me.

2. Masking invalid values afterwards:

The main takeaway here is that you should never ever compare against np.nan directly since it will always be False. You can check this yourself by looking at the result from np.nan == np.nan. The way to handle this is use the dedicated np.isnan function. Or alternatively negate the np.isfinite function if you also want to catch +/- np.inf values at the same time.

avg = np.average(C/(A+B), axis = 1)
avg[np.isnan(avg)] = 0

# or to include inf
avg[~np.isfinite(avg)] = 0