text: # Authors : Alexandre Gramfort, alexandre.gramfort@telecom-paristech.fr (2011)
  #           Denis A. Engemann <denis.engemann@gmail.com>
  # License : BSD 3-clause
  
  import numpy as np
  
  from ..parallel import parallel_func
  from ..io.pick import _picks_to_idx
  from ..utils import logger, verbose, _time_mask
  from .multitaper import psd_array_multitaper
  
  
  def _psd_func(epoch, noverlap, n_per_seg, nfft, fs, freq_mask, func):
      """Aux function."""
      return func(epoch, fs=fs, nperseg=n_per_seg, noverlap=noverlap,
                  nfft=nfft, window='hamming')[2][..., freq_mask, :]
  
  
  def _check_nfft(n, n_fft, n_per_seg, n_overlap):
      """Ensure n_fft, n_per_seg and n_overlap make sense."""
      if n_per_seg is None and n_fft > n:
          raise ValueError(('If n_per_seg is None n_fft is not allowed to be > '
                            'n_times. If you want zero-padding, you have to set '
                            'n_per_seg to relevant length. Got n_fft of %d while'
                            ' signal length is %d.') % (n_fft, n))
      n_per_seg = n_fft if n_per_seg is None or n_per_seg > n_fft else n_per_seg
      n_per_seg = n if n_per_seg > n else n_per_seg
      if n_overlap >= n_per_seg:
          raise ValueError(('n_overlap cannot be greater than n_per_seg (or '
                            'n_fft). Got n_overlap of %d while n_per_seg is '
                            '%d.') % (n_overlap, n_per_seg))
      return n_fft, n_per_seg, n_overlap
  
  
  def _check_psd_data(inst, tmin, tmax, picks, proj, reject_by_annotation=False):
      """Check PSD data / pull arrays from inst."""
      from ..io.base import BaseRaw
      from ..epochs import BaseEpochs
      from ..evoked import Evoked
      if not isinstance(inst, (BaseEpochs, BaseRaw, Evoked)):
          raise ValueError('epochs must be an instance of Epochs, Raw, or'
                           'Evoked. Got type {}'.format(type(inst)))
  
      time_mask = _time_mask(inst.times, tmin, tmax, sfreq=inst.info['sfreq'])
      picks = _picks_to_idx(inst.info, picks, 'data', with_ref_meg=False)
      if proj:
          # Copy first so it's not modified
          inst = inst.copy().apply_proj()
  
      sfreq = inst.info['sfreq']
      if isinstance(inst, BaseRaw):
          start, stop = np.where(time_mask)[0][[0, -1]]
          rba = 'NaN' if reject_by_annotation else None
          data = inst.get_data(picks, start, stop + 1, reject_by_annotation=rba)
      elif isinstance(inst, BaseEpochs):
          data = inst.get_data()[:, picks][:, :, time_mask]
      else:  # Evoked
          data = inst.data[picks][:, time_mask]
  
      return data, sfreq
  
  
  @verbose
  def psd_array_welch(x, sfreq, fmin=0, fmax=np.inf, n_fft=256, n_overlap=0,
                      n_per_seg=None, n_jobs=1, verbose=None):
      """Compute power spectral density (PSD) using Welch's method.
  
      Parameters
      ----------
      x : array, shape=(..., n_times)
          The data to compute PSD from.
      sfreq : float
          The sampling frequency.
      fmin : float
          The lower frequency of interest.
      fmax : float
          The upper frequency of interest.
      n_fft : int
          The length of FFT used, must be ``>= n_per_seg`` (default: 256).
          The segments will be zero-padded if ``n_fft > n_per_seg``.
      n_overlap : int
          The number of points of overlap between segments. Will be adjusted
          to be <= n_per_seg. The default value is 0.
      n_per_seg : int | None
          Length of each Welch segment (windowed with a Hamming window). Defaults
          to None, which sets n_per_seg equal to n_fft.
      n_jobs : int
          Number of CPUs to use in the computation.
      %(verbose)s
  
      Returns
      -------
      psds : ndarray, shape (..., n_freqs) or
          The power spectral densities. All dimensions up to the last will
          be the same as input.
      freqs : ndarray, shape (n_freqs,)
          The frequencies.
  
      Notes
      -----
      .. versionadded:: 0.14.0
      """
      from scipy.signal import spectrogram
      dshape = x.shape[:-1]
      n_times = x.shape[-1]
      x = x.reshape(-1, n_times)
  
      # Prep the PSD
      n_fft, n_per_seg, n_overlap = _check_nfft(n_times, n_fft, n_per_seg,
                                                n_overlap)
      win_size = n_fft / float(sfreq)
      logger.info("Effective window size : %0.3f (s)" % win_size)
      freqs = np.arange(n_fft // 2 + 1, dtype=float) * (sfreq / n_fft)
      freq_mask = (freqs >= fmin) & (freqs <= fmax)
      freqs = freqs[freq_mask]
  
      # Parallelize across first N-1 dimensions
      parallel, my_psd_func, n_jobs = parallel_func(_psd_func, n_jobs=n_jobs)
      x_splits = np.array_split(x, n_jobs)
      f_spectrogram = parallel(my_psd_func(d, noverlap=n_overlap, nfft=n_fft,
                                           fs=sfreq, freq_mask=freq_mask,
                                           func=spectrogram, n_per_seg=n_per_seg)
                               for d in x_splits)
  
      # Combining, reducing windows and reshaping to original data shape
      psds = np.concatenate([np.nanmean(f_s, axis=-1)
                             for f_s in f_spectrogram], axis=0)
      psds.shape = dshape + (-1,)
      return psds, freqs
  
  
  @verbose
  def psd_welch(inst, fmin=0, fmax=np.inf, tmin=None, tmax=None, n_fft=256,
                n_overlap=0, n_per_seg=None, picks=None, proj=False, n_jobs=1,
                reject_by_annotation=True, verbose=None):
      """Compute the power spectral density (PSD) using Welch's method.
  
      Calculates periodograms for a sliding window over the time dimension, then
      averages them together for each channel/epoch.
  
      Parameters
      ----------
      inst : instance of Epochs or Raw or Evoked
          The data for PSD calculation
      fmin : float
          Min frequency of interest
      fmax : float
          Max frequency of interest
      tmin : float | None
          Min time of interest
      tmax : float | None
          Max time of interest
      n_fft : int
          The length of FFT used, must be ``>= n_per_seg`` (default: 256).
          The segments will be zero-padded if ``n_fft > n_per_seg``.
          If n_per_seg is None, n_fft must be >= number of time points
          in the data.
      n_overlap : int
          The number of points of overlap between segments. Will be adjusted
          to be <= n_per_seg. The default value is 0.
      n_per_seg : int | None
          Length of each Welch segment (windowed with a Hamming window). Defaults
          to None, which sets n_per_seg equal to n_fft.
      %(picks_good_data_noref)s
      proj : bool
          Apply SSP projection vectors. If inst is ndarray this is not used.
      n_jobs : int
          Number of CPUs to use in the computation.
      reject_by_annotation : bool
          Whether to omit bad segments from the data while computing the
          PSD. If True, annotated segments with a description that starts
          with 'bad' are omitted. Has no effect if ``inst`` is an Epochs or
          Evoked object. Defaults to True.
  
          .. versionadded:: 0.15.0
      %(verbose)s
  
      Returns
      -------
      psds : ndarray, shape (..., n_freqs)
          The power spectral densities. If input is of type Raw,
          then psds will be shape (n_channels, n_freqs), if input is type Epochs
          then psds will be shape (n_epochs, n_channels, n_freqs).
      freqs : ndarray, shape (n_freqs,)
          The frequencies.
  
      See Also
      --------
      mne.io.Raw.plot_psd
      mne.Epochs.plot_psd
      psd_multitaper
      psd_array_welch
  
      Notes
      -----
      .. versionadded:: 0.12.0
      """
      # Prep data
      data, sfreq = _check_psd_data(inst, tmin, tmax, picks, proj,
                                    reject_by_annotation=reject_by_annotation)
      return psd_array_welch(data, sfreq, fmin=fmin, fmax=fmax, n_fft=n_fft,
                             n_overlap=n_overlap, n_per_seg=n_per_seg,
                             n_jobs=n_jobs, verbose=verbose)
  
  
  @verbose
  def psd_multitaper(inst, fmin=0, fmax=np.inf, tmin=None, tmax=None,
                     bandwidth=None, adaptive=False, low_bias=True,
                     normalization='length', picks=None, proj=False,
                     n_jobs=1, verbose=None):
      """Compute the power spectral density (PSD) using multitapers.
  
      Calculates spectral density for orthogonal tapers, then averages them
      together for each channel/epoch. See [1] for a description of the tapers
      and [2] for the general method.
  
      Parameters
      ----------
      inst : instance of Epochs or Raw or Evoked
          The data for PSD calculation.
      fmin : float
          Min frequency of interest
      fmax : float
          Max frequency of interest
      tmin : float | None
          Min time of interest
      tmax : float | None
          Max time of interest
      bandwidth : float
          The bandwidth of the multi taper windowing function in Hz. The default
          value is a window half-bandwidth of 4.
      adaptive : bool
          Use adaptive weights to combine the tapered spectra into PSD
          (slow, use n_jobs >> 1 to speed up computation).
      low_bias : bool
          Only use tapers with more than 90% spectral concentration within
          bandwidth.
      normalization : str
          Either "full" or "length" (default). If "full", the PSD will
          be normalized by the sampling rate as well as the length of
          the signal (as in nitime).
      %(picks_good_data_noref)s
      proj : bool
          Apply SSP projection vectors. If inst is ndarray this is not used.
      n_jobs : int
          Number of CPUs to use in the computation.
      %(verbose)s
  
      Returns
      -------
      psds : ndarray, shape (..., n_freqs)
          The power spectral densities. If input is of type Raw,
          then psds will be shape (n_channels, n_freqs), if input is type Epochs
          then psds will be shape (n_epochs, n_channels, n_freqs).
      freqs : ndarray, shape (n_freqs,)
          The frequencies.
  
      References
      ----------
      .. [1] Slepian, D. "Prolate spheroidal wave functions, Fourier analysis,
             and uncertainty V: The discrete case." Bell System Technical
             Journal, vol. 57, 1978.
  
      .. [2] Percival D.B. and Walden A.T. "Spectral Analysis for Physical
             Applications: Multitaper and Conventional Univariate Techniques."
             Cambridge University Press, 1993.
  
      See Also
      --------
      mne.io.Raw.plot_psd
      mne.Epochs.plot_psd
      psd_array_multitaper
      psd_welch
      csd_multitaper
  
      Notes
      -----
      .. versionadded:: 0.12.0
      """
      # Prep data
      data, sfreq = _check_psd_data(inst, tmin, tmax, picks, proj)
      return psd_array_multitaper(data, sfreq, fmin=fmin, fmax=fmax,
                                  bandwidth=bandwidth, adaptive=adaptive,
                                  low_bias=low_bias, normalization=normalization,
                                  n_jobs=n_jobs, verbose=verbose)
  