def medianPSD(d0,win=np.hanning,dx=1.,axis=0,nans=False):
"""Return the 1D PSD "medianed" over a surface.
Axis indicates the axis over which to FFT
If nans is True, each slice will be stripped,
internally interpolated, and then the power spectra
interpolated to common frequency grid"""
d = stripnans(d0)
if win is not 1:
win = win(np.shape(d)[axis])/\
np.sqrt(np.mean(win(np.shape(d)[axis])**2))
win = np.repeat(win,np.shape(d)[axis-1])
win = np.reshape(win,(np.shape(d)[axis],np.shape(d)[axis-1]))
if axis is 1:
win = np.transpose(win)
c = np.abs(np.fft.fft(d*win,axis=axis)/np.shape(d)[axis])**2
c = np.median(c,axis=axis-1)
f = np.fft.fftfreq(np.size(c),d=dx)
f = f[:np.size(c)/2]
c = c[:np.size(c)/2]
c[1:] = 2*c[1:]
return f,c
def realPSD(d0,win=np.hanning,dx=1.,axis=None,nans=False,minpx=10):
"""This function returns the PSD of a real function
Gets rid of zero frequency and puts all power in positive frequencies
Returns only positive frequencies
"""
if nans is True:
d = stripnans(d0)
else:
d = d0
if len(d) < minpx:
return np.nan
#Get Fourier components
c = components(d,win=win)
#Handle collapsing to 1D PSD if axis keyword is set
if axis==0:
c = c[:,0]
elif axis==1:
c = c[0,:]
#Reform into PSD
if np.size(np.shape(c)) is 2:
f = [np.fft.fftfreq(np.shape(c)[0],d=dx)[:np.shape(c)[0]/2],\
np.fft.fftfreq(np.shape(c)[1],d=dx)[:np.shape(c)[1]/2]]
c = c[:np.shape(c)[0]/2,:np.shape(c)[1]/2]
c[0,0] = 0.
#Handle normalization
c = 2*c
c[0,:] = c[0,:]/np.sqrt(2.)
c[:,0] = c[:,0]/np.sqrt(2.)
elif np.size(np.shape(c)) is 1:
f = np.fft.fftfreq(np.size(c),d=dx)
f = f[:np.size(c)/2]
c = c[:np.size(c)/2]
c[0] = 0.
c = c*np.sqrt(2.)
return f[1:],np.abs(c[1:])**2
def meanPSD(d0,win=np.hanning,dx=1.,axis=0,irregular=False,returnInd=False,minpx=10):
"""Return the 1D PSD averaged over a surface.
Axis indicates the axis over which to FFT
If irregular is True, each slice will be stripped
and then the power spectra
interpolated to common frequency grid
Presume image has already been interpolated internally
If returnInd is true, return array of power spectra
Ignores slices with less than minpx non-nans
"""
#Handle which axis is transformed
if axis==0:
d0 = np.transpose(d0)
#Create list of slices
if irregular is True:
d0 = [stripnans(di) for di in d0]
else:
d0 = [di for di in d0]
#Create power spectra from each slice
pows = [realPSD(s,win=win,dx=dx,minpx=minpx) for s in d0 \
if np.sum(~np.isnan(s)) >= minpx]
#Interpolate onto common frequency grid of shortest slice
if irregular is True:
#Determine smallest frequency grid
ln = [len(s[0]) for s in pows]
freq = pows[np.argmin(ln)][0]
#Interpolate
pp = [griddata(p[0],p[1],freq) for p in pows]
else:
pp = [p[1] for p in pows]
freq = pows[0][0]
#Average
pa = np.mean(pp,axis=0)
if returnInd is True:
return freq,pp
return freq,pa