fft.js 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
  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', 'dotDivide', 'conj', 'pow', 'ceil', 'log2'];
  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. dotDivide,
  16. conj,
  17. pow,
  18. ceil,
  19. log2
  20. } = _ref;
  21. /**
  22. * Calculate N-dimensional fourier transform
  23. *
  24. * Syntax:
  25. *
  26. * math.fft(arr)
  27. *
  28. * Examples:
  29. *
  30. * math.fft([[1, 0], [1, 0]]) // returns [[{re:2, im:0}, {re:2, im:0}], [{re:0, im:0}, {re:0, im:0}]]
  31. *
  32. *
  33. * See Also:
  34. *
  35. * ifft
  36. *
  37. * @param {Array | Matrix} arr An array or matrix
  38. * @return {Array | Matrix} N-dimensional fourier transformation of the array
  39. */
  40. return typed(name, {
  41. Array: _ndFft,
  42. Matrix: function Matrix(matrix) {
  43. return matrix.create(_ndFft(matrix.toArray()));
  44. }
  45. });
  46. /**
  47. * Perform an N-dimensional Fourier transform
  48. *
  49. * @param {Array} arr The array
  50. * @return {Array} resulting array
  51. */
  52. function _ndFft(arr) {
  53. var size = arraySize(arr);
  54. if (size.length === 1) return _fft(arr, size[0]);
  55. // ndFft along dimension 1,...,N-1 then 1dFft along dimension 0
  56. return _1dFft(arr.map(slice => _ndFft(slice, size.slice(1))), 0);
  57. }
  58. /**
  59. * Perform an 1-dimensional Fourier transform
  60. *
  61. * @param {Array} arr The array
  62. * @param {number} dim dimension of the array to perform on
  63. * @return {Array} resulting array
  64. */
  65. function _1dFft(arr, dim) {
  66. var size = arraySize(arr);
  67. if (dim !== 0) return new Array(size[0]).fill(0).map((_, i) => _1dFft(arr[i], dim - 1));
  68. if (size.length === 1) return _fft(arr);
  69. function _transpose(arr) {
  70. // Swap first 2 dimensions
  71. var size = arraySize(arr);
  72. return new Array(size[1]).fill(0).map((_, j) => new Array(size[0]).fill(0).map((_, i) => arr[i][j]));
  73. }
  74. return _transpose(_1dFft(_transpose(arr), 1));
  75. }
  76. /**
  77. * Perform an 1-dimensional non-power-of-2 Fourier transform using Chirp-Z Transform
  78. *
  79. * @param {Array} arr The array
  80. * @return {Array} resulting array
  81. */
  82. function _czt(arr) {
  83. var n = arr.length;
  84. var w = exp(divideScalar(multiplyScalar(-1, multiplyScalar(I, tau)), n));
  85. var chirp = [];
  86. for (var i = 1 - n; i < n; i++) {
  87. chirp.push(pow(w, divideScalar(pow(i, 2), 2)));
  88. }
  89. var N2 = pow(2, ceil(log2(n + n - 1)));
  90. var xp = [...new Array(n).fill(0).map((_, i) => multiplyScalar(arr[i], chirp[n - 1 + i])), ...new Array(N2 - n).fill(0)];
  91. var ichirp = [...new Array(n + n - 1).fill(0).map((_, i) => divideScalar(1, chirp[i])), ...new Array(N2 - (n + n - 1)).fill(0)];
  92. var fftXp = _fft(xp);
  93. var fftIchirp = _fft(ichirp);
  94. var fftProduct = new Array(N2).fill(0).map((_, i) => multiplyScalar(fftXp[i], fftIchirp[i]));
  95. var ifftProduct = dotDivide(conj(_ndFft(conj(fftProduct))), N2);
  96. var ret = [];
  97. for (var _i = n - 1; _i < n + n - 1; _i++) {
  98. ret.push(multiplyScalar(ifftProduct[_i], chirp[_i]));
  99. }
  100. return ret;
  101. }
  102. /**
  103. * Perform an 1-dimensional Fourier transform
  104. *
  105. * @param {Array} arr The array
  106. * @return {Array} resulting array
  107. */
  108. function _fft(arr) {
  109. var len = arr.length;
  110. if (len === 1) return [arr[0]];
  111. if (len % 2 === 0) {
  112. var ret = [..._fft(arr.filter((_, i) => i % 2 === 0), len / 2), ..._fft(arr.filter((_, i) => i % 2 === 1), len / 2)];
  113. for (var k = 0; k < len / 2; k++) {
  114. var p = ret[k];
  115. var q = multiplyScalar(ret[k + len / 2], exp(multiplyScalar(multiplyScalar(tau, I), divideScalar(-k, len))));
  116. ret[k] = addScalar(p, q);
  117. ret[k + len / 2] = addScalar(p, multiplyScalar(-1, q));
  118. }
  119. return ret;
  120. } else {
  121. // use chirp-z transform for non-power-of-2 FFT
  122. return _czt(arr);
  123. }
  124. // throw new Error('Can only calculate FFT of power-of-two size')
  125. }
  126. });