feat: add FFTConvolution class

This will allow to reuse arrays for better performance.

Author: targos

benchmark/convolution.js

@@ -1,6 +1,6 @@
 'use strict';
 
-const { fftConvolution, directConvolution } = require('..');
+const { FFTConvolution, directConvolution } = require('..');
 
 const tests = {
   data: [128, 512, 2048, 4096, 16384, 65536, 262144, 1048576],
@@ -11,6 +11,11 @@ function test(dataLength, kernelLength) {
   const data = Array.from({ length: dataLength }, Math.random);
   const kernel = Array.from({ length: kernelLength }, Math.random);
 
+  const fft = new FFTConvolution(dataLength, kernel);
+  const fftConvolution = (data) => {
+    return fft.convolute(data);
+  };
+
   const fftResult = measure(data, kernel, fftConvolution);
   const directResult = measure(data, kernel, directConvolution);
 

ml-convolution.d.ts

@@ -14,4 +14,11 @@ declare module 'ml-convolution' {
     kernel: ArrayLike<number>,
     borderType?: BorderType
   ): number[];
+  export class FFTConvolution {
+    public constructor(size: number, kernel: ArrayLike<number>);
+    public convolute(
+      input: ArrayLike<number>,
+      borderType?: BorderType
+    ): number[];
+  }
 }

src/fftConvolution.js

@@ -3,53 +3,66 @@ import nextPOT from 'next-power-of-two';
 
 import { checkKernel } from './utils';
 
-export default function fftConvolution(input, kernel, borderType = 'CONSTANT') {
-  checkKernel(kernel);
-  switch (borderType) {
-    case 'CONSTANT': {
-      return fftConvolutionImpl(input, kernel, false);
-    }
-    case 'CUT': {
-      return fftConvolutionImpl(input, kernel, true);
-    }
-    default: {
-      throw new Error(`unexpected border type: ${borderType}`);
+export class FFTConvolution {
+  constructor(size, kernel) {
+    if (!Number.isInteger(size) || size < 1) {
+      throw new TypeError('size must be a positive integer');
     }
+    checkKernel(kernel);
+    this.kernelOffset = (kernel.length - 1) / 2;
+    this.doubleOffset = 2 * this.kernelOffset;
+    const resultLength = size + this.doubleOffset;
+    this.fftLength = nextPOT(Math.max(resultLength, 2));
+    this.fft = new FFT(this.fftLength);
+    kernel = kernel.slice().reverse();
+    const { output: fftKernel, input: result } = createPaddedFFt(
+      kernel,
+      this.fft,
+      this.fftLength
+    );
+    this.fftKernel = fftKernel;
+    this.ifftOutput = this.fft.createComplexArray();
+    this.result = result;
   }
-}
 
-function fftConvolutionImpl(input, kernel, cutBorder) {
-  const kernelOffset = (kernel.length - 1) / 2;
-  const doubleOffset = 2 * kernelOffset;
-  const resultLength = input.length + doubleOffset;
-  const fftLength = nextPOT(resultLength);
+  convolute(input, borderType = 'CONSTANT') {
+    // if (input.length) // TODO CHECK SIZE
+    const { output: fftInput } = createPaddedFFt(
+      input,
+      this.fft,
+      this.fftLength
+    );
 
-  const fft = new FFT(fftLength);
-
-  kernel = kernel.slice().reverse();
-  const { output: fftKernel, input: result } = createPaddedFFt(
-    kernel,
-    fft,
-    fftLength
-  );
-  const { output: fftInput } = createPaddedFFt(input, fft, fftLength);
+    for (var i = 0; i < fftInput.length; i += 2) {
+      const tmp =
+        fftInput[i] * this.fftKernel[i] -
+        fftInput[i + 1] * this.fftKernel[i + 1];
+      fftInput[i + 1] =
+        fftInput[i] * this.fftKernel[i + 1] +
+        fftInput[i + 1] * this.fftKernel[i];
+      fftInput[i] = tmp;
+    }
 
-  for (var i = 0; i < fftInput.length; i += 2) {
-    const tmp = fftInput[i] * fftKernel[i] - fftInput[i + 1] * fftKernel[i + 1];
-    fftInput[i + 1] =
-      fftInput[i] * fftKernel[i + 1] + fftInput[i + 1] * fftKernel[i];
-    fftInput[i] = tmp;
-  }
-  const inverse = fftKernel;
-  fft.inverseTransform(inverse, fftInput);
-  const r = fft.fromComplexArray(inverse, result);
-  if (cutBorder) {
-    return r.slice(doubleOffset, input.length);
-  } else {
-    return r.slice(kernelOffset, kernelOffset + input.length);
+    this.fft.inverseTransform(this.ifftOutput, fftInput);
+    const r = this.fft.fromComplexArray(this.ifftOutput, this.result);
+    switch (borderType) {
+      case 'CONSTANT': {
+        return r.slice(this.kernelOffset, this.kernelOffset + input.length);
+      }
+      case 'CUT': {
+        return r.slice(this.doubleOffset, input.length);
+      }
+      default: {
+        throw new Error(`unexpected border type: ${borderType}`);
+      }
+    }
   }
 }
 
+export function fftConvolution(input, kernel, borderType = 'CONSTANT') {
+  return new FFTConvolution(input.length, kernel).convolute(input, borderType);
+}
+
 function createPaddedFFt(data, fft, length) {
   const input = [];
   let i = 0;

src/index.js

@@ -1,7 +1,6 @@
-import directConvolution from './directConvolution';
-import fftConvolution from './fftConvolution';
+export { default as directConvolution } from './directConvolution';
+export * from './fftConvolution';
 
-export { directConvolution, fftConvolution };
 export const BorderType = {
   CONSTANT: 'CONSTANT',
   CUT: 'CUT'