I want to estimate the Power spectral density using Continuous wavelet transform and a Morlet Wavelet. Bellow you can find the function I am using. Any comments or suggestions on whether or not the following code is correct?
import pycwt as wavelet
mother_wave_dict = {
'gaussian': wavelet.DOG(),
'paul': wavelet.Paul(),
'mexican_hat': wavelet.MexicanHat()
}
def trace_PSD_wavelet(x, dt, dj, mother_wave='morlet'):
"""
Method to calculate the power spectral density using wavelet method.
Parameters
----------
x : array-like
the components of the field to apply wavelet tranform
dt: int
the sampling time of the timeseries
dj: determines how many scales are used to estimate wavelet coeff
(e.g., for dj=1 -> 2**numb_scales
mother_wave: str
Returns
-------
db_x,db_y,db_zz: array-like
component coeficients of th wavelet tranform
freq : list
Frequency of the corresponding psd points.
psd : list
Power Spectral Density of the signal.
psd : list
The scales at which wavelet was estimated
"""
if mother_wave in mother_wave_dict.keys():
mother_morlet = mother_wave_dict[mother_wave]
else:
mother_morlet = wavelet.Morlet()
N = len(x)
db_x, _, freqs, _, _, _ = wavelet.cwt(x, dt, dj, wavelet=mother_morlet)
# Estimate trace powerspectral density
PSD = (np.nanmean(np.abs(db_x)**2, axis=1))*( 2*dt)
# Also estimate the scales to use later
scales = ((1/freqs)/dt)#.astype(int)
return db_x, freqs, PSD, scales

