multitaper.py 文件源码

def _mt_spectra(x, dpss, sfreq, n_fft=None):
    """ Compute tapered spectra

    Parameters
    ----------
    x : array, shape=(n_signals, n_times)
        Input signal
    dpss : array, shape=(n_tapers, n_times)
        The tapers
    sfreq : float
        The sampling frequency
    n_fft : int | None
        Length of the FFT. If None, the number of samples in the input signal
        will be used.

    Returns
    -------
    x_mt : array, shape=(n_signals, n_tapers, n_times)
        The tapered spectra
    freqs : array
        The frequency points in Hz of the spectra
    """

    if n_fft is None:
        n_fft = x.shape[1]

    # remove mean (do not use in-place subtraction as it may modify input x)
    x = x - np.mean(x, axis=-1)[:, np.newaxis]

    # only keep positive frequencies
    freqs = fftpack.fftfreq(n_fft, 1. / sfreq)
    freq_mask = (freqs >= 0)
    freqs = freqs[freq_mask]

    # The following is equivalent to this, but uses less memory:
    # x_mt = fftpack.fft(x[:, np.newaxis, :] * dpss, n=n_fft)
    n_tapers = dpss.shape[0] if dpss.ndim > 1 else 1
    x_mt = np.zeros((len(x), n_tapers, freq_mask.sum()), dtype=np.complex128)
    for idx, sig in enumerate(x):
        x_mt[idx] = fftpack.fft(sig[np.newaxis, :] * dpss,
                                n=n_fft)[:, freq_mask]
    return x_mt, freqs