Replace all nan in a 2D array by the nanmean() value of adjacent neighboring cells

I have a 2D array witch include some nan and I would like do the following :

• replace each nan in the array by the nanmean() value of adjacent neighboring cells.
• if the neighboring cells are all nan, then the value is nan
• if there are some nan in the neighboring cells then the value is the nanmean() value of adjacent neighboring cells. Let's take some examples :

a = np.random.randint(0,100,60).astype(float)
a = a.reshape(6,10) 
a[1,1]= np.nan
a[2:5,4:7]=np.nan
a
array([[48., 84., 80., 59., 43., 60., 31., 37.,  4., 75.],
       [83., nan, 52., 34., 95., 15., 69.,  7.,  7., 16.],
       [10.,  6.,  6., 44., nan, nan, nan, 95., 28.,  4.],
       [12.,  1., 62., 96., nan, nan, nan, 66., 21., 80.],
       [41., 18.,  1., 49., nan, nan, nan, 27., 64., nan],
       [13., 33., 98., 85., 77., 20., 73., 57., 15., 28.]])

# In this array a[1,1] is a nan and should be replaced by 46.125
np.nanmean(np.array([48.,84.,80.,52.,6.,6.,10.,83.]))
46.125

# Another example a[2,4] is a nan and should be replaced by 56.8
np.nanmean(np.array([34.,95.,15.,44.,96.]))
56.8

etc.

So far I have tried the following techniques :

• Using some loop : very easy but very slow if the array is big (1000 x 1000)
• Scipy NearestNDInterpolator, interesting but replace all nan (nans surrounded by nans must remain nans!)

How to do this in the fastest way, without looping ?

Use a 2D convolution (scipy.signal.convolve2d):

from scipy.signal import convolve2d

m = np.isnan(a)
kernel = np.ones((3, 3))

# sum of values
S1 = convolve2d(np.nan_to_num(a), kernel, mode='same')
# count of non-NaNs
S2 = convolve2d(~m, kernel, mode='same')

# optional: update the mask to ensure excluding cells
# fully surrounded with NaNs:
m &= S2!=0

# replace NaNs with mean (sum/count)
a[m] = S1[m]/S2[m]

Updated a:

array([[48.  , 84.  , 80.  , 59.  , 43.  , 60.  , 31.  , 37.  ,  4.  , 75.  ],
       [83.  , 46.12, 52.  , 34.  , 95.  , 15.  , 69.  ,  7.  ,  7.  , 16.  ],
       [10.  ,  6.  ,  6.  , 44.  , 56.8 , 59.67, 50.4 , 95.  , 28.  ,  4.  ],
       [12.  ,  1.  , 62.  , 96.  , 63.  ,   nan, 62.67, 66.  , 21.  , 80.  ],
       [41.  , 18.  ,  1.  , 49.  , 65.4 , 56.67, 48.6 , 27.  , 64.  , 41.6 ],
       [13.  , 33.  , 98.  , 85.  , 77.  , 20.  , 73.  , 57.  , 15.  , 28.  ]])

Intermediates:

# m
array([[False, False, False, False, False, False, False, False, False, False],
       [False,  True, False, False, False, False, False, False, False, False],
       [False, False, False, False,  True,  True,  True, False, False, False],
       [False, False, False, False,  True,  True,  True, False, False, False],
       [False, False, False, False,  True,  True,  True, False, False,  True],
       [False, False, False, False, False, False, False, False, False, False]])

# S1
array([[215., 347., 309., 363., 306., 313., 219., 155., 146., 102.],
       [231., 369., 365., 413., 350., 313., 314., 278., 273., 134.],
       [112., 232., 301., 389., 284., 179., 252., 293., 324., 156.],
       [ 88., 157., 283., 258., 189.,   0., 188., 301., 385., 197.],
       [118., 279., 443., 468., 327., 170., 243., 323., 358., 208.],
       [105., 204., 284., 310., 231., 170., 177., 236., 191., 107.]])

# S2
array([[3., 5., 5., 6., 6., 6., 6., 6., 6., 4.],
       [5., 8., 8., 8., 7., 6., 7., 8., 9., 6.],
       [5., 8., 8., 7., 5., 3., 5., 7., 9., 6.],
       [6., 9., 9., 6., 3., 0., 3., 6., 8., 5.],
       [6., 9., 9., 7., 5., 3., 5., 7., 8., 5.],
       [4., 6., 6., 5., 4., 3., 4., 5., 5., 3.]])