Slow Fast Fourier Transform

polymul.js

// Cooley-Tukey FFT implementation in JS
// My DSP is very shaky so this probably has some bugs...

const N_TRIG = 4096;
const HALF_N_TRIG = N_TRIG >> 1;
const COS = new Float64Array(N_TRIG);
const SIN = new Float64Array(N_TRIG);

for (let i = 0; i < N_TRIG; ++i) {
  const rad = 2 * Math.PI * i / N_TRIG;
  COS[i] = Math.cos(rad);
  SIN[i] = Math.sin(rad);
}

const fft = (vals, fw) => {
  if (vals.length == 2) return vals.slice();
  const mem = new Float64Array(vals.length);
  const partlen = mem.length >> 1;
  for (let i = 0; i < partlen; i += 2) {
    mem[i] = vals[i << 1];
    mem[i + 1] = vals[(i << 1) + 1];
    mem[partlen + i] = vals[(i << 1) + 2];
    mem[partlen + i + 1] = vals[(i << 1) + 3];
  }
  const a = fft(mem.subarray(0, partlen), fw);
  const b = fft(mem.subarray(partlen), fw);
  for (let i = 0; i < partlen; i += 2) {
    let trigInd = Math.floor((i / vals.length) * N_TRIG);
    if (fw && trigInd) trigInd = N_TRIG - trigInd;
    const trigInd2 = (HALF_N_TRIG + trigInd) % N_TRIG;
    mem[i] = a[i] + COS[trigInd] * b[i] - SIN[trigInd] * b[i + 1];
    mem[i + 1] = a[i + 1] + SIN[trigInd] * b[i] + COS[trigInd] * b[i + 1];
    mem[partlen + i] = a[i] + COS[trigInd2] * b[i] - SIN[trigInd2] * b[i + 1];
    mem[partlen + i + 1] = a[i + 1] + SIN[trigInd2] * b[i] + COS[trigInd2] * b[i + 1];
  }
  return mem;
}

const cmInplace = (a, b) => {
  for (let i = 0; i < a.length; i += 2) {
    const re = a[i] * b[i] - a[i + 1] * b[i + 1];
    const im = a[i] * b[i + 1] + a[i + 1] * b[i];
    a[i] = re;
    a[i + 1] = im;
  }
}

class Polynomial {
  constructor(coeffs) {
    this.coeffs = coeffs;
  }
  static create(...coeffs) {
    return new Polynomial(coeffs.reverse());
  }
  add(other) {
    if (this.coeffs.length < other.coeffs.length) return other.add(this);
    const res = this.coeffs.slice();
    for (let i = 0; i < other.coeffs.length; ++i) res[i] += other.coeffs[i];
    return new Polynomial(res);
  }
  mul(other) {
    const finalLen = other.coeffs.length + this.coeffs.length - 1;
    let len = finalLen - 1;
    len |= len >> 1;
    len |= len >> 2;
    len |= len >> 4;
    len |= len >> 8;
    len = ((len | (len >> 16)) + 1) << 1;
    const a = new Float64Array(len);
    const b = new Float64Array(len);
    for (let i = 0; i < this.coeffs.length; ++i) {
      a[i << 1] = this.coeffs[i];
    }
    for (let i = 0; i < other.coeffs.length; ++i) {
      b[i << 1] = other.coeffs[i];
    }
    const r1 = fft(a, 1);
    const r2 = fft(b, 1);
    cmInplace(r1, r2);
    const orig = fft(r1, 0);
    const out = [];
    for (let i = 0; i < finalLen; ++i) {
      out.push(Math.round(1000 * orig[i << 1] / (len >> 1)) / 1000);
    }
    return new Polynomial(out);
  }
  eval(x) {
    let result = 0;
    for (let i = 0; i < this.coeffs.length; ++i) {
      result = x * result + this.coeffs[i];
    }
    return result;
  }
  [Symbol.toPrimitive](hint) {
    if (hint == 'string') {
      let res = '';
      for (let i = 0; i < this.coeffs.length; ++i) {
        if (this.coeffs[i] != 0) {
          const num = Math.abs(this.coeffs[i]);
          res = (this.coeffs[i] < 0 ? ' - ' : ' + ') + (num == 1 && i != 0 ? '' : num) + (i == 0 ? '' : i == 1 ? 'x' : 'x^' + i) + res;
        }
      }
      return res.slice(3);
    }
  }
  toString() {
    return this[Symbol.toPrimitive]('string');
  }
  [Symbol.for('nodejs.util.inspect.custom')]() {
    return this.toString();
  }
}

const randLen = (len, min, max) => Array.from({ length: len }, () => Math.floor(Math.random() * (max - min + 1)) + min);

const ex = Polynomial.create(...randLen(5000, -2, 4));
const ex2 = Polynomial.create(...randLen(500, -2, 4));

const ts = Date.now();
const res = ex.mul(ex2);
const te = Date.now() - ts;

console.log(`(${ex}) * (${ex2}) = ${res}`);
console.log(`done in ${te}ms`)