mlab.py 30 KB

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