Estimate Power spectral density using Continuous wavelet transform form pycwt

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

It is hard to answer this question without knowing what you mean by "correct".

As far as I understand, your code allows to:

• Select a given wavelet
• Compute the cwt of input data for selected wavelet
• Compute the "PSD" of the CWT-transformed data for the entire recording.

I was able to run your code against an electrogram data example from the scipy library and it ran as expected:

from scipy.misc import electrocardiogram
ecg = electrocardiogram()

Since these data are sampled at 360Hz, I use dt=1/360:

db_x, freqs, PSD, scales = trace_PSD_wavelet(ecg, 1/360, 1/24, 'morlet')

Plotting the output db_x:

fig = plt.imshow(np.abs(db_x), extent=[db_x.shape[1],db_x.shape[0],scales[-1],scales[]], aspect='auto')
plt.xlabel('time')
plt.ylabel('scale')

Plotting the corresponding "PSD":

What you call "PSD" measures the energy contained in the CWT-transformed data at each scale, averaged over the entire recording, 5 minutes of data in this example. I am not sure how you plan to use this information but be careful that this is not the PSD of the original input time domain data.

Finally, concerning the Python implementation, you may simplify the way you call the default wavelet. Simply add the Morlet wavelet to your dictionary:

mother_wave_dict = {
    'gaussian': wavelet.DOG(),
    'paul': wavelet.Paul(),
    'mexican_hat': wavelet.MexicanHat()
    'morlet': wavelet.Morlet()
}

Then you can avoid the if statement in your function and simply call:

mother_morlet = mother_wave_dict[mother_wave].