Vectorized sliding / rolling numpy nanmean

user_12039011 0
Viewed 118

I currently use this code to do a moving window average:

n=500

x_copy=np.hstack((np.full(n,np.nan),copy.deepcopy(x),np.full(n,np.nan)))

x_values=[]
for i in range(n,len(x)+n):
    x_values.append(np.nanmean(x[i-n:i+n+1]))

plt.plot(x_values)

with x the array I'm working on and n is half the length of the window. However, I need to do this quickly, as I have to do roughly 4400*10 of this operation, with arrays around 60000 elements in length. After searching for a while, I found that np.convolve should work, so I have this code instead:

plt.plot(np.convolve(x, np.ones(((2*n),))/(2*n), mode='valid'),zorder=2)

While this is really fast, it's not doing exactly what I need it to do, as it seems to stop 500*2 units before the end of the array. For reference, here is the image of the plots of both of them: the blue is my own code, the orange is the convolution. I want to use convolve to speed up my moving window average, though I don't know how.

Reference

convolution moving-average numpy arrays python
1 Answers

Set those NaNs to 0s and then use np.convolve on those masked versions -

# Window-size
W = 2*n+1

# Non-nans mask
m = ~np.isnan(x)

# "Masked" input array
x0 = np.where(m,x,0)

# Setup conv kernel and perform conv on x0 and mask m for the counts to divide
K = np.ones(W)
out = (np.convolve(x0,K)/np.convolve(m,K))[W-1:]
user_32938810
Related