mlab.py 35 KB

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