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