diff --git a/src/expression/embeddedDocs/embeddedDocs.js b/src/expression/embeddedDocs/embeddedDocs.js index 50ee611fd..34ff8092d 100644 --- a/src/expression/embeddedDocs/embeddedDocs.js +++ b/src/expression/embeddedDocs/embeddedDocs.js @@ -142,6 +142,8 @@ import { subsetDocs } from './function/matrix/subset.js' import { traceDocs } from './function/matrix/trace.js' import { transposeDocs } from './function/matrix/transpose.js' import { zerosDocs } from './function/matrix/zeros.js' +import { fftDocs } from './function/matrix/fft.js' +import { ifftDocs } from './function/matrix/ifft.js' import { combinationsDocs } from './function/probability/combinations.js' import { combinationsWithRepDocs } from './function/probability/combinationsWithRep.js' import { factorialDocs } from './function/probability/factorial.js' @@ -461,6 +463,8 @@ export const embeddedDocs = { trace: traceDocs, transpose: transposeDocs, zeros: zerosDocs, + fft: fftDocs, + ifft: ifftDocs, // functions - probability combinations: combinationsDocs, diff --git a/src/expression/embeddedDocs/function/matrix/fft.js b/src/expression/embeddedDocs/function/matrix/fft.js new file mode 100644 index 000000000..d4d183520 --- /dev/null +++ b/src/expression/embeddedDocs/function/matrix/fft.js @@ -0,0 +1,14 @@ +export const fftDocs = { + name: 'fft', + category: 'Matrix', + syntax: [ + 'fft(x)' + ], + description: 'Calculate N-dimensional fourier transform', + examples: [ + 'fft([[1, 0], [1, 0]])' + ], + seealso: [ + 'ifft' + ] +} diff --git a/src/expression/embeddedDocs/function/matrix/ifft.js b/src/expression/embeddedDocs/function/matrix/ifft.js new file mode 100644 index 000000000..6584cc664 --- /dev/null +++ b/src/expression/embeddedDocs/function/matrix/ifft.js @@ -0,0 +1,14 @@ +export const ifftDocs = { + name: 'ifft', + category: 'Matrix', + syntax: [ + 'ifft(x)' + ], + description: 'Calculate N-dimensional inverse fourier transform', + examples: [ + 'ifft([[2, 2], [0, 0]])' + ], + seealso: [ + 'fft' + ] +} diff --git a/src/factoriesAny.js b/src/factoriesAny.js index 9a98646ee..6ff5739af 100644 --- a/src/factoriesAny.js +++ b/src/factoriesAny.js @@ -92,6 +92,8 @@ export { createSubset } from './function/matrix/subset.js' export { createTranspose } from './function/matrix/transpose.js' export { createCtranspose } from './function/matrix/ctranspose.js' export { createZeros } from './function/matrix/zeros.js' +export { createFft } from './function/matrix/fft.js' +export { createIfft } from './function/matrix/ifft.js' export { createErf } from './function/special/erf.js' export { createMode } from './function/statistics/mode.js' export { createProd } from './function/statistics/prod.js' diff --git a/src/function/matrix/fft.js b/src/function/matrix/fft.js new file mode 100644 index 000000000..d2d457938 --- /dev/null +++ b/src/function/matrix/fft.js @@ -0,0 +1,112 @@ +import { arraySize } from '../../utils/array.js' +import { factory } from '../../utils/factory.js' + +const name = 'fft' +const dependencies = [ + 'typed', + 'matrix', + 'addScalar', + 'multiplyScalar', + 'divideScalar', + 'exp', + 'tau', + 'i' +] + +export const createFft = /* #__PURE__ */ factory(name, dependencies, ({ + typed, + matrix, + addScalar, + multiplyScalar, + divideScalar, + exp, + tau, + i: I +}) => { + /** + * Calculate N-dimensional fourier transform + * + * Syntax: + * + * math.fft(arr) + * + * Examples: + * + * math.fft([[1, 0], [1, 0]]) // returns [[{re:2, im:0}, {re:2, im:0}], [{re:0, im:0}, {re:0, im:0}]] + * + * + * See Also: + * + * ifft + * + * @param {Array | Matrix} arr An array or matrix + * @return {Array | Matrix} N-dimensional fourier transformation of the array + */ + return typed(name, { + Array: _ndFft, + Matrix: function (matrix) { + return matrix.create(_ndFft(matrix.toArray())) + } + }) + + /** + * Perform an N-dimensional Fourier transform + * + * @param {Array} arr The array + * @return {Array} resulting array + */ + function _ndFft (arr) { + const size = arraySize(arr) + if (size.length === 1) return _fft(arr, size[0]) + // ndFft along dimension 1,...,N-1 then 1dFft along dimension 0 + return _1dFft(arr.map(slice => _ndFft(slice, size.slice(1))), 0) + } + + /** + * Perform an 1-dimensional Fourier transform + * + * @param {Array} arr The array + * @param {number} dim dimension of the array to perform on + * @return {Array} resulting array + */ + function _1dFft (arr, dim) { + const size = arraySize(arr) + if (dim !== 0) return new Array(size[0]).fill(0).map((_, i) => _1dFft(arr[i], dim - 1)) + if (size.length === 1) return _fft(arr) + function _transpose (arr) { // Swap first 2 dimensions + const size = arraySize(arr) + return new Array(size[1]).fill(0).map((_, j) => new Array(size[0]).fill(0).map((_, i) => arr[i][j])) + } + return _transpose(_1dFft(_transpose(arr), 1)) + } + + /** + * Perform an 1-dimensional Fourier transform + * + * @param {Array} arr The array + * @return {Array} resulting array + */ + function _fft (arr) { + const len = arr.length + if (len === 1) return [arr[0]] + if (len % 2 === 0) { + const ret = [ + ..._fft(arr.filter((_, i) => i % 2 === 0), len / 2), + ..._fft(arr.filter((_, i) => i % 2 === 1), len / 2) + ] + for (let k = 0; k < len / 2; k++) { + const p = ret[k] + const q = multiplyScalar( + ret[k + len / 2], + exp( + multiplyScalar(multiplyScalar(tau, I), divideScalar(-k, len)) + ) + ) + ret[k] = addScalar(p, q) + ret[k + len / 2] = addScalar(p, multiplyScalar(-1, q)) + } + return ret + } + throw new Error('Can only calculate FFT of power-of-two size') + } +}) diff --git a/src/function/matrix/ifft.js b/src/function/matrix/ifft.js new file mode 100644 index 000000000..60081a79e --- /dev/null +++ b/src/function/matrix/ifft.js @@ -0,0 +1,43 @@ +import { arraySize } from '../../utils/array.js' +import { factory } from '../../utils/factory.js' +import { isMatrix } from '../../utils/is.js' + +const name = 'ifft' +const dependencies = [ + 'typed', + 'fft', + 'dotDivide', + 'conj' +] + +export const createIfft = /* #__PURE__ */ factory(name, dependencies, ({ + typed, + fft, + dotDivide, + conj +}) => { + /** + * Calculate N-dimensional inverse fourier transform + * + * Syntax: + * + * math.ifft(arr) + * + * Examples: + * + * math.ifft([[2, 2], [0, 0]]) // returns [[{re:1, im:0}, {re:0, im:0}], [{re:1, im:0}, {re:0, im:0}]] + * + * See Also: + * + * fft + * + * @param {Array | Matrix} arr An array or matrix + * @return {Array | Matrix} N-dimensional fourier transformation of the array + */ + return typed(name, { + 'Array | Matrix': function (arr) { + const size = isMatrix(arr) ? arr.size() : arraySize(arr) + return dotDivide(conj(fft(conj(arr))), size.reduce((acc, curr) => acc * curr, 1)) + } + }) +}) diff --git a/test/unit-tests/function/matrix/fft.test.js b/test/unit-tests/function/matrix/fft.test.js new file mode 100644 index 000000000..8b7b86f3a --- /dev/null +++ b/test/unit-tests/function/matrix/fft.test.js @@ -0,0 +1,41 @@ +// test fft + +import approx from '../../../../tools/approx.js' +import math from '../../../../src/defaultInstance.js' +const fft = math.fft + +describe('fft', function () { + it('should calculate 1-dimensional fourier transformation', function () { + const in1 = [1, math.complex(2, -1), math.complex(0, -1), math.complex(-1, 2)] + const out1 = [2, math.complex(-2, -2), math.complex(0, -2), math.complex(4, 4)] + approx.deepEqual(fft(in1.valueOf()), out1.valueOf()) + approx.deepEqual(fft(math.matrix(in1)), math.matrix(out1)) + }) + + it('should calculate multidimensional fourier transformation', function () { + const in1 = [ + [1, 0], + [1, 0] + ] + const out1 = [ + [2, 2], + [0, 0] + ] + approx.deepEqual(fft(in1.valueOf()), out1.valueOf()) + approx.deepEqual(fft(math.matrix(in1)), math.matrix(out1)) + const in2 = [ + [0, 0, 1, 1], + [0, 0, 1, 1], + [1, 1, 0, 0], + [1, 1, 0, 0] + ] + const out2 = [ + [8, 0, 0, 0], + [0, math.complex(0, 4), 0, -4], + [0, 0, 0, 0], + [0, -4, 0, math.complex(0, -4)] + ] + approx.deepEqual(fft(in2.valueOf()), out2.valueOf()) + approx.deepEqual(fft(math.matrix(in2)), math.matrix(out2)) + }) +}) diff --git a/test/unit-tests/function/matrix/ifft.test.js b/test/unit-tests/function/matrix/ifft.test.js new file mode 100644 index 000000000..fa0c10ab0 --- /dev/null +++ b/test/unit-tests/function/matrix/ifft.test.js @@ -0,0 +1,20 @@ +// test ifft + +import approx from '../../../../tools/approx.js' +import math from '../../../../src/defaultInstance.js' +const ifft = math.ifft + +describe('ifft', function () { + it('should calculate 1-dimensional inverse fourier transformation', function () { + approx.deepEqual(ifft([2, math.complex(-2, -2), math.complex(0, -2), math.complex(4, 4)]), [1, math.complex(2, -1), math.complex(0, -1), math.complex(-1, 2)]) + }) + + it('should calculate multidimensional inverse fourier transformation', function () { + const in1 = [ + [1, 0], + [1, 0] + ] + approx.deepEqual(math.fft(ifft(in1.valueOf())), in1.valueOf()) + approx.deepEqual(math.fft(ifft(math.matrix(in1))), math.matrix(in1)) + }) +}) diff --git a/types/index.d.ts b/types/index.d.ts index 4d94313a5..75ada9a61 100644 --- a/types/index.d.ts +++ b/types/index.d.ts @@ -1997,6 +1997,20 @@ declare namespace math { */ zeros(m: number, n: number, format?: string): MathCollection + /** + * Calculate N-dimensional fourier transform + * @param {Array | Matrix} arr An array or matrix + * @return {Array | Matrix} N-dimensional fourier transformation of the array + */ + fft(arr: T): T + + /** + * Calculate N-dimensional inverse fourier transform + * @param {Array | Matrix} arr An array or matrix + * @return {Array | Matrix} N-dimensional fourier transformation of the array + */ + ifft(arr: T): T + /************************************************************************* * Probability functions ************************************************************************/ diff --git a/types/index.ts b/types/index.ts index ba54f7e39..5046bb43b 100644 --- a/types/index.ts +++ b/types/index.ts @@ -1023,6 +1023,62 @@ Matrices examples { assert.strictEqual(math.matrix([1, 2, 3]) instanceof math.Matrix, true) } + + // Fourier transform and inverse + { + assert.ok( + math.deepEqual( + math.fft([ + [1, 0], + [1, 0], + ]), + [ + [math.complex(2, 0), math.complex(2, 0)], + [math.complex(0, 0), math.complex(0, 0)], + ] + ) + ) + assert.ok( + math.deepEqual( + math.fft( + math.matrix([ + [1, 0], + [1, 0], + ]) + ), + math.matrix([ + [math.complex(2, 0), math.complex(2, 0)], + [math.complex(0, 0), math.complex(0, 0)], + ]) + ) + ) + assert.ok( + math.deepEqual( + math.ifft([ + [2, 2], + [0, 0], + ]), + [ + [math.complex(1, 0), math.complex(0, 0)], + [math.complex(1, 0), math.complex(0, 0)], + ] + ) + ) + assert.ok( + math.deepEqual( + math.ifft( + math.matrix([ + [2, 2], + [0, 0], + ]) + ), + math.matrix([ + [math.complex(1, 0), math.complex(0, 0)], + [math.complex(1, 0), math.complex(0, 0)], + ]) + ) + ) + } } /*