nonshape.py 文件源码

def f_psd(x, Fs, method,
        Hzmed=0, welch_params={'window':'hanning','nperseg':1000,'noverlap':None}):
    '''
    Calculate the power spectrum of a signal

    Parameters
    ----------
    x : array
        temporal signal
    Fs : integer
        sampling rate
    method : str in ['fftmed','welch']
        Method for calculating PSD
    Hzmed : float
        relevant if method == 'fftmed'
        Frequency width of the median filter
    welch_params : dict
        relevant if method == 'welch'
        Parameters to sp.signal.welch

    Returns
    -------
    f : array
        frequencies corresponding to the PSD output
    psd : array
        power spectrum
    '''

    if method == 'fftmed':
        # Calculate frequencies
        N = len(x)
        f = np.arange(0,Fs/2,Fs/N)

        # Calculate PSD
        rawfft = np.fft.fft(x)
        psd = np.abs(rawfft[:len(f)])**2

        # Median filter
        if Hzmed > 0:
            sampmed = np.argmin(np.abs(f-Hzmed/2.0))
            psd = signal.medfilt(psd,sampmed*2+1)

    elif method == 'welch':
        f, psd = sp.signal.welch(x, fs=Fs, **welch_params)

    else:
        raise ValueError('input for PSD method not recognized')

    return f, psd