Feature: Implementation of fourier transform (#2540)

* Fix #46

Draft Implementations

* Add docs

* Fixup

* Add type declaration and test

* Fixup tests

* Fixup test

* Format

* Add examples in docs

* Update fft.js

Edit example in docs (`math.fft` returns complex matrix).

* Update ifft.js

Edit example in docs (`math.ffti` returns complex matrix).

* Update fft.js

Edit docs examples, representation of complex number from `a+bi` to `{re:a, im:b}`

* Update ifft.js

Edit docs examples, representation of complex number from `a+bi` to `{re:a, im:b}`

* Update index.ts

Edit test.
Add test for `math.ifft`
`math.fft` returns complex matrix.

* Update index.ts

Use `approx.deepEqual` instead off `assert.deepStrictEqual`.

* Update index.ts

Format code

* Update index.ts

Use `assert.ok(math.deepEqual(...))` instead of `approx.deepEqual`.

* Update index.ts

Format

* Update index.ts

Typo: replace `approx` with `assert`.

Co-authored-by: Jos de Jong <wjosdejong@gmail.com>

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,

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'
+  ]
+}

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'
+  ]
+}

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'

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')
+  }
+})

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))
+    }
+  })
+})

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))
+  })
+})

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))
+  })
+})

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<T extends MathCollection>(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<T extends MathCollection>(arr: T): T
+
     /*************************************************************************
      * Probability functions
      ************************************************************************/

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)],
+        ])
+      )
+    )
+  }
 }
 
 /*