Python 2D array replace value by NaN when other array is NaN

Question

I have 2 arrays:

import numpy as np
A = np.array([[np.nan,np.nan,0],[4,5,6],[7,8,9]])
B = np.zeros((len(A),len(A[0])))

For every NaN in A, I would like to replace the zero in B at the corresponding indices by np.nan. How can I do that ? I tried with masks but couldn't make it work.

B = np.ma.masked_where(np.isnan(A) == True, np.nan) 

In reality I am working with bigger arrays (582x533) so I don't want to do it by hand.

Answer 1

You can use np.isnan(A) directly:

B = np.zeros_like(A)
B[np.isnan(A)] = np.nan

np.where is useful for when you want to construct B directly:

B = np.where(np.isnan(A), np.nan, 0)

Answer 2

You can create np.zeros with A.shape then use np.where like below:

(this approach construct B two times)

>>> import numpy as np
>>> A = np.array([[np.nan,np.nan,0],[4,5,6],[7,8,9]])
>>> B = np.zeros(A.shape)
>>> B = np.where(np.isnan(A) , np.nan, B)
>>> B
array([[nan, nan,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.]])