fft.js 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. import { arraySize } from '../../utils/array.js';
  2. import { factory } from '../../utils/factory.js';
  3. var name = 'fft';
  4. var dependencies = ['typed', 'matrix', 'addScalar', 'multiplyScalar', 'divideScalar', 'exp', 'tau', 'i'];
  5. export var createFft = /* #__PURE__ */factory(name, dependencies, _ref => {
  6. var {
  7. typed,
  8. matrix,
  9. addScalar,
  10. multiplyScalar,
  11. divideScalar,
  12. exp,
  13. tau,
  14. i: I
  15. } = _ref;
  16. /**
  17. * Calculate N-dimensional fourier transform
  18. *
  19. * Syntax:
  20. *
  21. * math.fft(arr)
  22. *
  23. * Examples:
  24. *
  25. * math.fft([[1, 0], [1, 0]]) // returns [[{re:2, im:0}, {re:2, im:0}], [{re:0, im:0}, {re:0, im:0}]]
  26. *
  27. *
  28. * See Also:
  29. *
  30. * ifft
  31. *
  32. * @param {Array | Matrix} arr An array or matrix
  33. * @return {Array | Matrix} N-dimensional fourier transformation of the array
  34. */
  35. return typed(name, {
  36. Array: _ndFft,
  37. Matrix: function Matrix(matrix) {
  38. return matrix.create(_ndFft(matrix.toArray()));
  39. }
  40. });
  41. /**
  42. * Perform an N-dimensional Fourier transform
  43. *
  44. * @param {Array} arr The array
  45. * @return {Array} resulting array
  46. */
  47. function _ndFft(arr) {
  48. var size = arraySize(arr);
  49. if (size.length === 1) return _fft(arr, size[0]);
  50. // ndFft along dimension 1,...,N-1 then 1dFft along dimension 0
  51. return _1dFft(arr.map(slice => _ndFft(slice, size.slice(1))), 0);
  52. }
  53. /**
  54. * Perform an 1-dimensional Fourier transform
  55. *
  56. * @param {Array} arr The array
  57. * @param {number} dim dimension of the array to perform on
  58. * @return {Array} resulting array
  59. */
  60. function _1dFft(arr, dim) {
  61. var size = arraySize(arr);
  62. if (dim !== 0) return new Array(size[0]).fill(0).map((_, i) => _1dFft(arr[i], dim - 1));
  63. if (size.length === 1) return _fft(arr);
  64. function _transpose(arr) {
  65. // Swap first 2 dimensions
  66. var size = arraySize(arr);
  67. return new Array(size[1]).fill(0).map((_, j) => new Array(size[0]).fill(0).map((_, i) => arr[i][j]));
  68. }
  69. return _transpose(_1dFft(_transpose(arr), 1));
  70. }
  71. /**
  72. * Perform an 1-dimensional Fourier transform
  73. *
  74. * @param {Array} arr The array
  75. * @return {Array} resulting array
  76. */
  77. function _fft(arr) {
  78. var len = arr.length;
  79. if (len === 1) return [arr[0]];
  80. if (len % 2 === 0) {
  81. var ret = [..._fft(arr.filter((_, i) => i % 2 === 0), len / 2), ..._fft(arr.filter((_, i) => i % 2 === 1), len / 2)];
  82. for (var k = 0; k < len / 2; k++) {
  83. var p = ret[k];
  84. var q = multiplyScalar(ret[k + len / 2], exp(multiplyScalar(multiplyScalar(tau, I), divideScalar(-k, len))));
  85. ret[k] = addScalar(p, q);
  86. ret[k + len / 2] = addScalar(p, multiplyScalar(-1, q));
  87. }
  88. return ret;
  89. }
  90. throw new Error('Can only calculate FFT of power-of-two size');
  91. }
  92. });