1"""
2Numerical python functions written for compatibility with MATLAB
3commands with the same names. Most numerical python functions can be found in
4the `numpy` and `scipy` libraries. What remains here is code for performing
5spectral computations.
6
7Spectral functions
8------------------
9
10`cohere`
11    Coherence (normalized cross spectral density)
12
13`csd`
14    Cross spectral density using Welch's average periodogram
15
16`detrend`
17    Remove the mean or best fit line from an array
18
19`psd`
20    Power spectral density using Welch's average periodogram
21
22`specgram`
23    Spectrogram (spectrum over segments of time)
24
25`complex_spectrum`
26    Return the complex-valued frequency spectrum of a signal
27
28`magnitude_spectrum`
29    Return the magnitude of the frequency spectrum of a signal
30
31`angle_spectrum`
32    Return the angle (wrapped phase) of the frequency spectrum of a signal
33
34`phase_spectrum`
35    Return the phase (unwrapped angle) of the frequency spectrum of a signal
36
37`detrend_mean`
38    Remove the mean from a line.
39
40`detrend_linear`
41    Remove the best fit line from a line.
42
43`detrend_none`
44    Return the original line.
45
46`stride_windows`
47    Get all windows in an array in a memory-efficient manner
48"""
49
50import functools
51from numbers import Number
52
53import numpy as np
54
55from matplotlib import _api
56import matplotlib.cbook as cbook
57from matplotlib import docstring
58
59
60def window_hanning(x):
61    """
62    Return x times the hanning window of len(x).
63
64    See Also
65    --------
66    window_none : Another window algorithm.
67    """
68    return np.hanning(len(x))*x
69
70
71def window_none(x):
72    """
73    No window function; simply return x.
74
75    See Also
76    --------
77    window_hanning : Another window algorithm.
78    """
79    return x
80
81
82def detrend(x, key=None, axis=None):
83    """
84    Return x with its trend removed.
85
86    Parameters
87    ----------
88    x : array or sequence
89        Array or sequence containing the data.
90
91    key : {'default', 'constant', 'mean', 'linear', 'none'} or function
92        The detrending algorithm to use. 'default', 'mean', and 'constant' are
93        the same as `detrend_mean`. 'linear' is the same as `detrend_linear`.
94        'none' is the same as `detrend_none`. The default is 'mean'. See the
95        corresponding functions for more details regarding the algorithms. Can
96        also be a function that carries out the detrend operation.
97
98    axis : int
99        The axis along which to do the detrending.
100
101    See Also
102    --------
103    detrend_mean : Implementation of the 'mean' algorithm.
104    detrend_linear : Implementation of the 'linear' algorithm.
105    detrend_none : Implementation of the 'none' algorithm.
106    """
107    if key is None or key in ['constant', 'mean', 'default']:
108        return detrend(x, key=detrend_mean, axis=axis)
109    elif key == 'linear':
110        return detrend(x, key=detrend_linear, axis=axis)
111    elif key == 'none':
112        return detrend(x, key=detrend_none, axis=axis)
113    elif callable(key):
114        x = np.asarray(x)
115        if axis is not None and axis + 1 > x.ndim:
116            raise ValueError(f'axis(={axis}) out of bounds')
117        if (axis is None and x.ndim == 0) or (not axis and x.ndim == 1):
118            return key(x)
119        # try to use the 'axis' argument if the function supports it,
120        # otherwise use apply_along_axis to do it
121        try:
122            return key(x, axis=axis)
123        except TypeError:
124            return np.apply_along_axis(key, axis=axis, arr=x)
125    else:
126        raise ValueError(
127            f"Unknown value for key: {key!r}, must be one of: 'default', "
128            f"'constant', 'mean', 'linear', or a function")
129
130
131def detrend_mean(x, axis=None):
132    """
133    Return x minus the mean(x).
134
135    Parameters
136    ----------
137    x : array or sequence
138        Array or sequence containing the data
139        Can have any dimensionality
140
141    axis : int
142        The axis along which to take the mean.  See numpy.mean for a
143        description of this argument.
144
145    See Also
146    --------
147    detrend_linear : Another detrend algorithm.
148    detrend_none : Another detrend algorithm.
149    detrend : A wrapper around all the detrend algorithms.
150    """
151    x = np.asarray(x)
152
153    if axis is not None and axis+1 > x.ndim:
154        raise ValueError('axis(=%s) out of bounds' % axis)
155
156    return x - x.mean(axis, keepdims=True)
157
158
159def detrend_none(x, axis=None):
160    """
161    Return x: no detrending.
162
163    Parameters
164    ----------
165    x : any object
166        An object containing the data
167
168    axis : int
169        This parameter is ignored.
170        It is included for compatibility with detrend_mean
171
172    See Also
173    --------
174    detrend_mean : Another detrend algorithm.
175    detrend_linear : Another detrend algorithm.
176    detrend : A wrapper around all the detrend algorithms.
177    """
178    return x
179
180
181def detrend_linear(y):
182    """
183    Return x minus best fit line; 'linear' detrending.
184
185    Parameters
186    ----------
187    y : 0-D or 1-D array or sequence
188        Array or sequence containing the data
189
190    axis : int
191        The axis along which to take the mean.  See numpy.mean for a
192        description of this argument.
193
194    See Also
195    --------
196    detrend_mean : Another detrend algorithm.
197    detrend_none : Another detrend algorithm.
198    detrend : A wrapper around all the detrend algorithms.
199    """
200    # This is faster than an algorithm based on linalg.lstsq.
201    y = np.asarray(y)
202
203    if y.ndim > 1:
204        raise ValueError('y cannot have ndim > 1')
205
206    # short-circuit 0-D array.
207    if not y.ndim:
208        return np.array(0., dtype=y.dtype)
209
210    x = np.arange(y.size, dtype=float)
211
212    C = np.cov(x, y, bias=1)
213    b = C[0, 1]/C[0, 0]
214
215    a = y.mean() - b*x.mean()
216    return y - (b*x + a)
217
218
219def stride_windows(x, n, noverlap=None, axis=0):
220    """
221    Get all windows of x with length n as a single array,
222    using strides to avoid data duplication.
223
224    .. warning::
225
226        It is not safe to write to the output array.  Multiple
227        elements may point to the same piece of memory,
228        so modifying one value may change others.
229
230    Parameters
231    ----------
232    x : 1D array or sequence
233        Array or sequence containing the data.
234    n : int
235        The number of data points in each window.
236    noverlap : int, default: 0 (no overlap)
237        The overlap between adjacent windows.
238    axis : int
239        The axis along which the windows will run.
240
241    References
242    ----------
243    `stackoverflow: Rolling window for 1D arrays in Numpy?
244    <http://stackoverflow.com/a/6811241>`_
245    `stackoverflow: Using strides for an efficient moving average filter
246    <http://stackoverflow.com/a/4947453>`_
247    """
248    if noverlap is None:
249        noverlap = 0
250
251    if noverlap >= n:
252        raise ValueError('noverlap must be less than n')
253    if n < 1:
254        raise ValueError('n cannot be less than 1')
255
256    x = np.asarray(x)
257
258    if x.ndim != 1:
259        raise ValueError('only 1-dimensional arrays can be used')
260    if n == 1 and noverlap == 0:
261        if axis == 0:
262            return x[np.newaxis]
263        else:
264            return x[np.newaxis].transpose()
265    if n > x.size:
266        raise ValueError('n cannot be greater than the length of x')
267
268    # np.lib.stride_tricks.as_strided easily leads to memory corruption for
269    # non integer shape and strides, i.e. noverlap or n. See #3845.
270    noverlap = int(noverlap)
271    n = int(n)
272
273    step = n - noverlap
274    if axis == 0:
275        shape = (n, (x.shape[-1]-noverlap)//step)
276        strides = (x.strides[0], step*x.strides[0])
277    else:
278        shape = ((x.shape[-1]-noverlap)//step, n)
279        strides = (step*x.strides[0], x.strides[0])
280    return np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides)
281
282
283def _spectral_helper(x, y=None, NFFT=None, Fs=None, detrend_func=None,
284                     window=None, noverlap=None, pad_to=None,
285                     sides=None, scale_by_freq=None, mode=None):
286    """
287    Private helper implementing the common parts between the psd, csd,
288    spectrogram and complex, magnitude, angle, and phase spectrums.
289    """
290    if y is None:
291        # if y is None use x for y
292        same_data = True
293    else:
294        # The checks for if y is x are so that we can use the same function to
295        # implement the core of psd(), csd(), and spectrogram() without doing
296        # extra calculations.  We return the unaveraged Pxy, freqs, and t.
297        same_data = y is x
298
299    if Fs is None:
300        Fs = 2
301    if noverlap is None:
302        noverlap = 0
303    if detrend_func is None:
304        detrend_func = detrend_none
305    if window is None:
306        window = window_hanning
307
308    # if NFFT is set to None use the whole signal
309    if NFFT is None:
310        NFFT = 256
311
312    if mode is None or mode == 'default':
313        mode = 'psd'
314    _api.check_in_list(
315        ['default', 'psd', 'complex', 'magnitude', 'angle', 'phase'],
316        mode=mode)
317
318    if not same_data and mode != 'psd':
319        raise ValueError("x and y must be equal if mode is not 'psd'")
320
321    # Make sure we're dealing with a numpy array. If y and x were the same
322    # object to start with, keep them that way
323    x = np.asarray(x)
324    if not same_data:
325        y = np.asarray(y)
326
327    if sides is None or sides == 'default':
328        if np.iscomplexobj(x):
329            sides = 'twosided'
330        else:
331            sides = 'onesided'
332    _api.check_in_list(['default', 'onesided', 'twosided'], sides=sides)
333
334    # zero pad x and y up to NFFT if they are shorter than NFFT
335    if len(x) < NFFT:
336        n = len(x)
337        x = np.resize(x, NFFT)
338        x[n:] = 0
339
340    if not same_data and len(y) < NFFT:
341        n = len(y)
342        y = np.resize(y, NFFT)
343        y[n:] = 0
344
345    if pad_to is None:
346        pad_to = NFFT
347
348    if mode != 'psd':
349        scale_by_freq = False
350    elif scale_by_freq is None:
351        scale_by_freq = True
352
353    # For real x, ignore the negative frequencies unless told otherwise
354    if sides == 'twosided':
355        numFreqs = pad_to
356        if pad_to % 2:
357            freqcenter = (pad_to - 1)//2 + 1
358        else:
359            freqcenter = pad_to//2
360        scaling_factor = 1.
361    elif sides == 'onesided':
362        if pad_to % 2:
363            numFreqs = (pad_to + 1)//2
364        else:
365            numFreqs = pad_to//2 + 1
366        scaling_factor = 2.
367
368    if not np.iterable(window):
369        window = window(np.ones(NFFT, x.dtype))
370    if len(window) != NFFT:
371        raise ValueError(
372            "The window length must match the data's first dimension")
373
374    result = stride_windows(x, NFFT, noverlap, axis=0)
375    result = detrend(result, detrend_func, axis=0)
376    result = result * window.reshape((-1, 1))
377    result = np.fft.fft(result, n=pad_to, axis=0)[:numFreqs, :]
378    freqs = np.fft.fftfreq(pad_to, 1/Fs)[:numFreqs]
379
380    if not same_data:
381        # if same_data is False, mode must be 'psd'
382        resultY = stride_windows(y, NFFT, noverlap)
383        resultY = detrend(resultY, detrend_func, axis=0)
384        resultY = resultY * window.reshape((-1, 1))
385        resultY = np.fft.fft(resultY, n=pad_to, axis=0)[:numFreqs, :]
386        result = np.conj(result) * resultY
387    elif mode == 'psd':
388        result = np.conj(result) * result
389    elif mode == 'magnitude':
390        result = np.abs(result) / np.abs(window).sum()
391    elif mode == 'angle' or mode == 'phase':
392        # we unwrap the phase later to handle the onesided vs. twosided case
393        result = np.angle(result)
394    elif mode == 'complex':
395        result /= np.abs(window).sum()
396
397    if mode == 'psd':
398
399        # Also include scaling factors for one-sided densities and dividing by
400        # the sampling frequency, if desired. Scale everything, except the DC
401        # component and the NFFT/2 component:
402
403        # if we have a even number of frequencies, don't scale NFFT/2
404        if not NFFT % 2:
405            slc = slice(1, -1, None)
406        # if we have an odd number, just don't scale DC
407        else:
408            slc = slice(1, None, None)
409
410        result[slc] *= scaling_factor
411
412        # MATLAB divides by the sampling frequency so that density function
413        # has units of dB/Hz and can be integrated by the plotted frequency
414        # values. Perform the same scaling here.
415        if scale_by_freq:
416            result /= Fs
417            # Scale the spectrum by the norm of the window to compensate for
418            # windowing loss; see Bendat & Piersol Sec 11.5.2.
419            result /= (np.abs(window)**2).sum()
420        else:
421            # In this case, preserve power in the segment, not amplitude
422            result /= np.abs(window).sum()**2
423
424    t = np.arange(NFFT/2, len(x) - NFFT/2 + 1, NFFT - noverlap)/Fs
425
426    if sides == 'twosided':
427        # center the frequency range at zero
428        freqs = np.roll(freqs, -freqcenter, axis=0)
429        result = np.roll(result, -freqcenter, axis=0)
430    elif not pad_to % 2:
431        # get the last value correctly, it is negative otherwise
432        freqs[-1] *= -1
433
434    # we unwrap the phase here to handle the onesided vs. twosided case
435    if mode == 'phase':
436        result = np.unwrap(result, axis=0)
437
438    return result, freqs, t
439
440
441def _single_spectrum_helper(
442        mode, x, Fs=None, window=None, pad_to=None, sides=None):
443    """
444    Private helper implementing the commonality between the complex, magnitude,
445    angle, and phase spectrums.
446    """
447    _api.check_in_list(['complex', 'magnitude', 'angle', 'phase'], mode=mode)
448
449    if pad_to is None:
450        pad_to = len(x)
451
452    spec, freqs, _ = _spectral_helper(x=x, y=None, NFFT=len(x), Fs=Fs,
453                                      detrend_func=detrend_none, window=window,
454                                      noverlap=0, pad_to=pad_to,
455                                      sides=sides,
456                                      scale_by_freq=False,
457                                      mode=mode)
458    if mode != 'complex':
459        spec = spec.real
460
461    if spec.ndim == 2 and spec.shape[1] == 1:
462        spec = spec[:, 0]
463
464    return spec, freqs
465
466
467# Split out these keyword docs so that they can be used elsewhere
468docstring.interpd.update(
469    Spectral="""\
470Fs : float, default: 2
471    The sampling frequency (samples per time unit).  It is used to calculate
472    the Fourier frequencies, *freqs*, in cycles per time unit.
473
474window : callable or ndarray, default: `.window_hanning`
475    A function or a vector of length *NFFT*.  To create window vectors see
476    `.window_hanning`, `.window_none`, `numpy.blackman`, `numpy.hamming`,
477    `numpy.bartlett`, `scipy.signal`, `scipy.signal.get_window`, etc.  If a
478    function is passed as the argument, it must take a data segment as an
479    argument and return the windowed version of the segment.
480
481sides : {'default', 'onesided', 'twosided'}, optional
482    Which sides of the spectrum to return. 'default' is one-sided for real
483    data and two-sided for complex data. 'onesided' forces the return of a
484    one-sided spectrum, while 'twosided' forces two-sided.""",
485
486    Single_Spectrum="""\
487pad_to : int, optional
488    The number of points to which the data segment is padded when performing
489    the FFT.  While not increasing the actual resolution of the spectrum (the
490    minimum distance between resolvable peaks), this can give more points in
491    the plot, allowing for more detail. This corresponds to the *n* parameter
492    in the call to fft().  The default is None, which sets *pad_to* equal to
493    the length of the input signal (i.e. no padding).""",
494
495    PSD="""\
496pad_to : int, optional
497    The number of points to which the data segment is padded when performing
498    the FFT.  This can be different from *NFFT*, which specifies the number
499    of data points used.  While not increasing the actual resolution of the
500    spectrum (the minimum distance between resolvable peaks), this can give
501    more points in the plot, allowing for more detail. This corresponds to
502    the *n* parameter in the call to fft(). The default is None, which sets
503    *pad_to* equal to *NFFT*
504
505NFFT : int, default: 256
506    The number of data points used in each block for the FFT.  A power 2 is
507    most efficient.  This should *NOT* be used to get zero padding, or the
508    scaling of the result will be incorrect; use *pad_to* for this instead.
509
510detrend : {'none', 'mean', 'linear'} or callable, default: 'none'
511    The function applied to each segment before fft-ing, designed to remove
512    the mean or linear trend.  Unlike in MATLAB, where the *detrend* parameter
513    is a vector, in Matplotlib is it a function.  The :mod:`~matplotlib.mlab`
514    module defines `.detrend_none`, `.detrend_mean`, and `.detrend_linear`,
515    but you can use a custom function as well.  You can also use a string to
516    choose one of the functions: 'none' calls `.detrend_none`. 'mean' calls
517    `.detrend_mean`. 'linear' calls `.detrend_linear`.
518
519scale_by_freq : bool, default: True
520    Whether the resulting density values should be scaled by the scaling
521    frequency, which gives density in units of Hz^-1.  This allows for
522    integration over the returned frequency values.  The default is True for
523    MATLAB compatibility.""")
524
525
526@docstring.dedent_interpd
527def psd(x, NFFT=None, Fs=None, detrend=None, window=None,
528        noverlap=None, pad_to=None, sides=None, scale_by_freq=None):
529    r"""
530    Compute the power spectral density.
531
532    The power spectral density :math:`P_{xx}` by Welch's average
533    periodogram method.  The vector *x* is divided into *NFFT* length
534    segments.  Each segment is detrended by function *detrend* and
535    windowed by function *window*.  *noverlap* gives the length of
536    the overlap between segments.  The :math:`|\mathrm{fft}(i)|^2`
537    of each segment :math:`i` are averaged to compute :math:`P_{xx}`.
538
539    If len(*x*) < *NFFT*, it will be zero padded to *NFFT*.
540
541    Parameters
542    ----------
543    x : 1-D array or sequence
544        Array or sequence containing the data
545
546    %(Spectral)s
547
548    %(PSD)s
549
550    noverlap : int, default: 0 (no overlap)
551        The number of points of overlap between segments.
552
553    Returns
554    -------
555    Pxx : 1-D array
556        The values for the power spectrum :math:`P_{xx}` (real valued)
557
558    freqs : 1-D array
559        The frequencies corresponding to the elements in *Pxx*
560
561    References
562    ----------
563    Bendat & Piersol -- Random Data: Analysis and Measurement Procedures, John
564    Wiley & Sons (1986)
565
566    See Also
567    --------
568    specgram
569        `specgram` differs in the default overlap; in not returning the mean of
570        the segment periodograms; and in returning the times of the segments.
571
572    magnitude_spectrum : returns the magnitude spectrum.
573
574    csd : returns the spectral density between two signals.
575    """
576    Pxx, freqs = csd(x=x, y=None, NFFT=NFFT, Fs=Fs, detrend=detrend,
577                     window=window, noverlap=noverlap, pad_to=pad_to,
578                     sides=sides, scale_by_freq=scale_by_freq)
579    return Pxx.real, freqs
580
581
582@docstring.dedent_interpd
583def csd(x, y, NFFT=None, Fs=None, detrend=None, window=None,
584        noverlap=None, pad_to=None, sides=None, scale_by_freq=None):
585    """
586    Compute the cross-spectral density.
587
588    The cross spectral density :math:`P_{xy}` by Welch's average
589    periodogram method.  The vectors *x* and *y* are divided into
590    *NFFT* length segments.  Each segment is detrended by function
591    *detrend* and windowed by function *window*.  *noverlap* gives
592    the length of the overlap between segments.  The product of
593    the direct FFTs of *x* and *y* are averaged over each segment
594    to compute :math:`P_{xy}`, with a scaling to correct for power
595    loss due to windowing.
596
597    If len(*x*) < *NFFT* or len(*y*) < *NFFT*, they will be zero
598    padded to *NFFT*.
599
600    Parameters
601    ----------
602    x, y : 1-D arrays or sequences
603        Arrays or sequences containing the data
604
605    %(Spectral)s
606
607    %(PSD)s
608
609    noverlap : int, default: 0 (no overlap)
610        The number of points of overlap between segments.
611
612    Returns
613    -------
614    Pxy : 1-D array
615        The values for the cross spectrum :math:`P_{xy}` before scaling (real
616        valued)
617
618    freqs : 1-D array
619        The frequencies corresponding to the elements in *Pxy*
620
621    References
622    ----------
623    Bendat & Piersol -- Random Data: Analysis and Measurement Procedures, John
624    Wiley & Sons (1986)
625
626    See Also
627    --------
628    psd : equivalent to setting ``y = x``.
629    """
630    if NFFT is None:
631        NFFT = 256
632    Pxy, freqs, _ = _spectral_helper(x=x, y=y, NFFT=NFFT, Fs=Fs,
633                                     detrend_func=detrend, window=window,
634                                     noverlap=noverlap, pad_to=pad_to,
635                                     sides=sides, scale_by_freq=scale_by_freq,
636                                     mode='psd')
637
638    if Pxy.ndim == 2:
639        if Pxy.shape[1] > 1:
640            Pxy = Pxy.mean(axis=1)
641        else:
642            Pxy = Pxy[:, 0]
643    return Pxy, freqs
644
645
646_single_spectrum_docs = """\
647Compute the {quantity} of *x*.
648Data is padded to a length of *pad_to* and the windowing function *window* is
649applied to the signal.
650
651Parameters
652----------
653x : 1-D array or sequence
654    Array or sequence containing the data
655
656{Spectral}
657
658{Single_Spectrum}
659
660Returns
661-------
662spectrum : 1-D array
663    The {quantity}.
664freqs : 1-D array
665    The frequencies corresponding to the elements in *spectrum*.
666
667See Also
668--------
669psd
670    Returns the power spectral density.
671complex_spectrum
672    Returns the complex-valued frequency spectrum.
673magnitude_spectrum
674    Returns the absolute value of the `complex_spectrum`.
675angle_spectrum
676    Returns the angle of the `complex_spectrum`.
677phase_spectrum
678    Returns the phase (unwrapped angle) of the `complex_spectrum`.
679specgram
680    Can return the complex spectrum of segments within the signal.
681"""
682
683
684complex_spectrum = functools.partial(_single_spectrum_helper, "complex")
685complex_spectrum.__doc__ = _single_spectrum_docs.format(
686    quantity="complex-valued frequency spectrum",
687    **docstring.interpd.params)
688magnitude_spectrum = functools.partial(_single_spectrum_helper, "magnitude")
689magnitude_spectrum.__doc__ = _single_spectrum_docs.format(
690    quantity="magnitude (absolute value) of the frequency spectrum",
691    **docstring.interpd.params)
692angle_spectrum = functools.partial(_single_spectrum_helper, "angle")
693angle_spectrum.__doc__ = _single_spectrum_docs.format(
694    quantity="angle of the frequency spectrum (wrapped phase spectrum)",
695    **docstring.interpd.params)
696phase_spectrum = functools.partial(_single_spectrum_helper, "phase")
697phase_spectrum.__doc__ = _single_spectrum_docs.format(
698    quantity="phase of the frequency spectrum (unwrapped phase spectrum)",
699    **docstring.interpd.params)
700
701
702@docstring.dedent_interpd
703def specgram(x, NFFT=None, Fs=None, detrend=None, window=None,
704             noverlap=None, pad_to=None, sides=None, scale_by_freq=None,
705             mode=None):
706    """
707    Compute a spectrogram.
708
709    Compute and plot a spectrogram of data in x.  Data are split into
710    NFFT length segments and the spectrum of each section is
711    computed.  The windowing function window is applied to each
712    segment, and the amount of overlap of each segment is
713    specified with noverlap.
714
715    Parameters
716    ----------
717    x : array-like
718        1-D array or sequence.
719
720    %(Spectral)s
721
722    %(PSD)s
723
724    noverlap : int, default: 128
725        The number of points of overlap between blocks.
726    mode : str, default: 'psd'
727        What sort of spectrum to use:
728            'psd'
729                Returns the power spectral density.
730            'complex'
731                Returns the complex-valued frequency spectrum.
732            'magnitude'
733                Returns the magnitude spectrum.
734            'angle'
735                Returns the phase spectrum without unwrapping.
736            'phase'
737                Returns the phase spectrum with unwrapping.
738
739    Returns
740    -------
741    spectrum : array-like
742        2D array, columns are the periodograms of successive segments.
743
744    freqs : array-like
745        1-D array, frequencies corresponding to the rows in *spectrum*.
746
747    t : array-like
748        1-D array, the times corresponding to midpoints of segments
749        (i.e the columns in *spectrum*).
750
751    See Also
752    --------
753    psd : differs in the overlap and in the return values.
754    complex_spectrum : similar, but with complex valued frequencies.
755    magnitude_spectrum : similar single segment when mode is 'magnitude'.
756    angle_spectrum : similar to single segment when mode is 'angle'.
757    phase_spectrum : similar to single segment when mode is 'phase'.
758
759    Notes
760    -----
761    detrend and scale_by_freq only apply when *mode* is set to 'psd'.
762
763    """
764    if noverlap is None:
765        noverlap = 128  # default in _spectral_helper() is noverlap = 0
766    if NFFT is None:
767        NFFT = 256  # same default as in _spectral_helper()
768    if len(x) <= NFFT:
769        _api.warn_external("Only one segment is calculated since parameter "
770                           f"NFFT (={NFFT}) >= signal length (={len(x)}).")
771
772    spec, freqs, t = _spectral_helper(x=x, y=None, NFFT=NFFT, Fs=Fs,
773                                      detrend_func=detrend, window=window,
774                                      noverlap=noverlap, pad_to=pad_to,
775                                      sides=sides,
776                                      scale_by_freq=scale_by_freq,
777                                      mode=mode)
778
779    if mode != 'complex':
780        spec = spec.real  # Needed since helper implements generically
781
782    return spec, freqs, t
783
784
785@docstring.dedent_interpd
786def cohere(x, y, NFFT=256, Fs=2, detrend=detrend_none, window=window_hanning,
787           noverlap=0, pad_to=None, sides='default', scale_by_freq=None):
788    r"""
789    The coherence between *x* and *y*.  Coherence is the normalized
790    cross spectral density:
791
792    .. math::
793
794        C_{xy} = \frac{|P_{xy}|^2}{P_{xx}P_{yy}}
795
796    Parameters
797    ----------
798    x, y
799        Array or sequence containing the data
800
801    %(Spectral)s
802
803    %(PSD)s
804
805    noverlap : int, default: 0 (no overlap)
806        The number of points of overlap between segments.
807
808    Returns
809    -------
810    Cxy : 1-D array
811        The coherence vector.
812    freqs : 1-D array
813            The frequencies for the elements in *Cxy*.
814
815    See Also
816    --------
817    :func:`psd`, :func:`csd` :
818        For information about the methods used to compute :math:`P_{xy}`,
819        :math:`P_{xx}` and :math:`P_{yy}`.
820    """
821    if len(x) < 2 * NFFT:
822        raise ValueError(
823            "Coherence is calculated by averaging over *NFFT* length "
824            "segments.  Your signal is too short for your choice of *NFFT*.")
825    Pxx, f = psd(x, NFFT, Fs, detrend, window, noverlap, pad_to, sides,
826                 scale_by_freq)
827    Pyy, f = psd(y, NFFT, Fs, detrend, window, noverlap, pad_to, sides,
828                 scale_by_freq)
829    Pxy, f = csd(x, y, NFFT, Fs, detrend, window, noverlap, pad_to, sides,
830                 scale_by_freq)
831    Cxy = np.abs(Pxy) ** 2 / (Pxx * Pyy)
832    return Cxy, f
833
834
835class GaussianKDE:
836    """
837    Representation of a kernel-density estimate using Gaussian kernels.
838
839    Parameters
840    ----------
841    dataset : array-like
842        Datapoints to estimate from. In case of univariate data this is a 1-D
843        array, otherwise a 2D array with shape (# of dims, # of data).
844
845    bw_method : str, scalar or callable, optional
846        The method used to calculate the estimator bandwidth.  This can be
847        'scott', 'silverman', a scalar constant or a callable.  If a
848        scalar, this will be used directly as `kde.factor`.  If a
849        callable, it should take a `GaussianKDE` instance as only
850        parameter and return a scalar. If None (default), 'scott' is used.
851
852    Attributes
853    ----------
854    dataset : ndarray
855        The dataset with which `gaussian_kde` was initialized.
856
857    dim : int
858        Number of dimensions.
859
860    num_dp : int
861        Number of datapoints.
862
863    factor : float
864        The bandwidth factor, obtained from `kde.covariance_factor`, with which
865        the covariance matrix is multiplied.
866
867    covariance : ndarray
868        The covariance matrix of *dataset*, scaled by the calculated bandwidth
869        (`kde.factor`).
870
871    inv_cov : ndarray
872        The inverse of *covariance*.
873
874    Methods
875    -------
876    kde.evaluate(points) : ndarray
877        Evaluate the estimated pdf on a provided set of points.
878
879    kde(points) : ndarray
880        Same as kde.evaluate(points)
881
882    """
883
884    # This implementation with minor modification was too good to pass up.
885    # from scipy: https://github.com/scipy/scipy/blob/master/scipy/stats/kde.py
886
887    def __init__(self, dataset, bw_method=None):
888        self.dataset = np.atleast_2d(dataset)
889        if not np.array(self.dataset).size > 1:
890            raise ValueError("`dataset` input should have multiple elements.")
891
892        self.dim, self.num_dp = np.array(self.dataset).shape
893
894        if bw_method is None:
895            pass
896        elif cbook._str_equal(bw_method, 'scott'):
897            self.covariance_factor = self.scotts_factor
898        elif cbook._str_equal(bw_method, 'silverman'):
899            self.covariance_factor = self.silverman_factor
900        elif isinstance(bw_method, Number):
901            self._bw_method = 'use constant'
902            self.covariance_factor = lambda: bw_method
903        elif callable(bw_method):
904            self._bw_method = bw_method
905            self.covariance_factor = lambda: self._bw_method(self)
906        else:
907            raise ValueError("`bw_method` should be 'scott', 'silverman', a "
908                             "scalar or a callable")
909
910        # Computes the covariance matrix for each Gaussian kernel using
911        # covariance_factor().
912
913        self.factor = self.covariance_factor()
914        # Cache covariance and inverse covariance of the data
915        if not hasattr(self, '_data_inv_cov'):
916            self.data_covariance = np.atleast_2d(
917                np.cov(
918                    self.dataset,
919                    rowvar=1,
920                    bias=False))
921            self.data_inv_cov = np.linalg.inv(self.data_covariance)
922
923        self.covariance = self.data_covariance * self.factor ** 2
924        self.inv_cov = self.data_inv_cov / self.factor ** 2
925        self.norm_factor = (np.sqrt(np.linalg.det(2 * np.pi * self.covariance))
926                            * self.num_dp)
927
928    def scotts_factor(self):
929        return np.power(self.num_dp, -1. / (self.dim + 4))
930
931    def silverman_factor(self):
932        return np.power(
933            self.num_dp * (self.dim + 2.0) / 4.0, -1. / (self.dim + 4))
934
935    #  Default method to calculate bandwidth, can be overwritten by subclass
936    covariance_factor = scotts_factor
937
938    def evaluate(self, points):
939        """
940        Evaluate the estimated pdf on a set of points.
941
942        Parameters
943        ----------
944        points : (# of dimensions, # of points)-array
945            Alternatively, a (# of dimensions,) vector can be passed in and
946            treated as a single point.
947
948        Returns
949        -------
950        (# of points,)-array
951            The values at each point.
952
953        Raises
954        ------
955        ValueError : if the dimensionality of the input points is different
956                     than the dimensionality of the KDE.
957
958        """
959        points = np.atleast_2d(points)
960
961        dim, num_m = np.array(points).shape
962        if dim != self.dim:
963            raise ValueError("points have dimension {}, dataset has dimension "
964                             "{}".format(dim, self.dim))
965
966        result = np.zeros(num_m)
967
968        if num_m >= self.num_dp:
969            # there are more points than data, so loop over data
970            for i in range(self.num_dp):
971                diff = self.dataset[:, i, np.newaxis] - points
972                tdiff = np.dot(self.inv_cov, diff)
973                energy = np.sum(diff * tdiff, axis=0) / 2.0
974                result = result + np.exp(-energy)
975        else:
976            # loop over points
977            for i in range(num_m):
978                diff = self.dataset - points[:, i, np.newaxis]
979                tdiff = np.dot(self.inv_cov, diff)
980                energy = np.sum(diff * tdiff, axis=0) / 2.0
981                result[i] = np.sum(np.exp(-energy), axis=0)
982
983        result = result / self.norm_factor
984
985        return result
986
987    __call__ = evaluate
988