[FIXED] 3d numpy array calculate mean of columns with nans

Issue

I have this 3d array:

import numpy as np
import numpy.ma as ma

a = np.array([[[1,2,3],[4,np.nan,6], [7,8,9]],
             [[11,12,13],[14,np.nan,16], [17,18,19]]])

a.shape
(2, 3, 3)

array([[[ 1.,  2.,  3.],
        [ 4., nan,  6.],
        [ 7.,  8.,  9.]],

       [[11., 12., 13.],
        [14., nan, 16.],
        [17., 18., 19.]]])

So, I have 2 sets of 2d data. I want to calculate the column mean of each set and replace the nans with that value.

So, I want the result:

array([[[ 1.,  2.,  3.],
        [ 4.,  5,   6.],
        [ 7.,  8.,  9.]],

       [[11., 12., 13.],
        [14., 15,  16.],
        [17., 18., 19.]]])

( 2 + 8 = 10 / 2 = 5)

( 12 + 18 = 30 / 2 = 15)

I tried:

a = np.where(np.isnan(a), 
             ma.array(a, mask=np.isnan(a)).mean(axis=1), 
             a)

with mean along axis 1 but it gives:

operands could not be broadcast together with shapes (2,3,3) (2,3) (2,3,3) 

Solution

Can't you just calculate the mean along the second axis (axis=1) while ignoring NaNs using np.nanmean then broadcast it back to the original array shape:

import numpy as np

a = np.array([[[1, 2, 3], [4, np.nan, 6], [7, 8, 9]],
             [[11, 12, 13], [14, np.nan, 16], [17, 18, 19]]])

means = np.nanmean(a, axis=1)
means_reshaped = means[:, np.newaxis, :]
a = np.where(np.isnan(a), means_reshaped, a)

print(a)

Output:

array([[[ 1.,  2.,  3.],
        [ 4.,  5.,  6.],
        [ 7.,  8.,  9.]],

       [[11., 12., 13.],
        [14., 15., 16.],
        [17., 18., 19.]]])

Answered By - Sash Sinha