[jupyter][JavaScriot][Japanese] Understanding FFT

dft.mjs

// complex number arithmetic for DFT/FFT
function expi(theta) {return [Math.cos(theta), Math.sin(theta)];}
function iadd([ax, ay], [bx, by]) {return [ax + bx, ay + by];}
function imul([ax, ay], [bx, by]) {
  return [ax * bx - ay * by, ax * by + ay * bx];
}
function idivr([x, y], r) {return [x / r, y / r];}
function isum(cs) {return cs.reduce((s, c) => iadd(s, c), [0, 0]);}

// DFT
function dftc(c, T) {
  return [...Array(c.length).keys()].map(i => isum(
    c.map((cn, n) => imul(cn, expi(T * n * i)))
  ));
}
export function dft(f) {
  const N = f.length, T = -2 * Math.PI / N;
  return dftc(f, T);
}
export function idft(F) {
  const N = F.length, T = 2 * Math.PI / N;
  return dftc(F, T).map(z => idivr(z, N));
}

fft-rec.js

// complex number arithmetic for DFT/FFT
function expi(theta) {return [Math.cos(theta), Math.sin(theta)];}
function iadd([ax, ay], [bx, by]) {return [ax + bx, ay + by];}
function isub([ax, ay], [bx, by]) {return [ax - bx, ay - by];}
function imul([ax, ay], [bx, by]) {
  return [ax * bx - ay * by, ax * by + ay * bx];
}
function idivr([x, y], r) {return [x / r, y / r];}

// FFT: Cooley-Tukey FFT
function fftrec(c, T, N = c.length) {
  if (N === 1) return c;
  const l = fftrec(c.filter((_, i) => i % 2 === 0), T, N / 2);
  const r = fftrec(c.filter((_, i) => i % 2 === 1), T, N / 2);
  const re = r.map((ri, i) => imul(ri, expi(T / N * i)));
  const laddr = l.map((li, i) => iadd(li, re[i]));
  const lsubr = l.map((li, i) => isub(li, re[i]));
  return laddr.concat(lsubr);
}
function fft0(f) {return fftrec(f, -2 * Math.PI);}
function ifft0(F) {return fftrec(F, 2 * Math.PI).map(z => idivr(z, F.length));}

// ordering
{
  function reorder(c, N = c.length) {
    if (N === 1) return c;
    const l = reorder(c.filter((_, i) => i % 2 === 0), N / 2);
    const r = reorder(c.filter((_, i) => i % 2 === 1), N / 2);
    return l.concat(r);
  }
  function toBin(k) {
    return n => n.toString(2).padStart(k, "0");
  }
  console.log([0,1,2,3].map(toBin(2)));
  console.log(reorder([0,1,2,3]).map(toBin(2)));
  
  console.log([0,1,2,3,4,5,6,7].map(toBin(3)));
  console.log(reorder([0,1,2,3,4,5,6,7]).map(toBin(3)));
  
  console.log([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15].map(toBin(4)));
  console.log(reorder([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]).map(toBin(4)));
}

// example
{
  const fr0 = [1,3,4,2, 5,6,2,4, 0,1,3,4, 5,62,2,3];
  const f0 = fr0.map(r => [r, 0]);
  
  const F = fft0(f0);
  const f1 = ifft0(F);
  const fr1 = f1.map(([r]) => r);
  
  console.log("fr0:", fr0);
  console.log("F:", F);
  console.log("f1:", f1);
  console.log("fr1:", fr1.map(Math.round));
}

fft.mjs

// complex number arithmetic for DFT/FFT
function conj([x, y]) {return [x, -y];}
function expi(theta) {return [Math.cos(theta), Math.sin(theta)];}
function iadd([ax, ay], [bx, by]) {return [ax + bx, ay + by];}
function isub([ax, ay], [bx, by]) {return [ax - bx, ay - by];}
function imul([ax, ay], [bx, by]) {
  return [ax * bx - ay * by, ax * by + ay * bx];
}
function idivr([x, y], r) {return [x / r, y / r];}

// bit operations for FFT
function revBit(k, n0) {
  if (k === 1) return n0;
  const s1 = ((n0 & 0xaaaaaaaa) >>> 1) | ((n0 & 0x55555555) << 1);
  if (k === 2) return s1;
  const s2 = ((s1 & 0xcccccccc) >>> 2) | ((s1 & 0x33333333) << 2);
  if (k <= 4) return s2 >>> (4 - k);
  const s3 = ((s2 & 0xf0f0f0f0) >>> 4) | ((s2 & 0x0f0f0f0f) << 4);
  if (k <= 8) return s3 >>> (8 - k);
  const s4 = ((s3 & 0xff00ff00) >>> 8) | ((s3 & 0x00ff00ff) << 8);
  if (k <= 16) return s3 >>> (16 - k);
  const s5 = ((s4 & 0xffff0000) >>> 16) | ((s4 & 0x0000ffff) << 16);
  return s5 >>> (32 - k);
}
function isPowerOf2(n) {return (n & (n - 1)) === 0;}
function nextPowerOf2(n) {
  return isPowerOf2(n) ? n : 1 << (32 - Math.clz32(n));
}

// FFT: Cooley-Tukey FFT
function fftin1(c, T, N = c.length) {
  const k = Math.log2(N);
  const rec = c.map((_, i) => c[revBit(k, i)]);
  for (let Nh = 1; Nh < N; Nh *= 2) {
    T /= 2;
    for (let s = 0; s < N; s += Nh * 2) {
      for (let i = 0, li = s, ri = s + Nh; i < Nh; i++, li++, ri++) { 
        const l = rec[li], re = imul(rec[ri], expi(T * i));
        [rec[li], rec[ri]] = [iadd(l, re), isub(l, re)];
      }
    }
  }
  return rec;
}
function fft1(f) {return fftin1(f, -2 * Math.PI);}
function ifft1(F) {return fftin1(F, 2 * Math.PI).map(z => idivr(z, F.length));}

// FFT: Blustein's algorithm
function fftin2(c, T, N = c.length) {
  const Nd = N * 2;
  const bc = c.map((_, n) => expi(T * n * n / Nd));
  const b = bc.map(conj);
  const a = c.map((cn, n) => imul(cn, bc[n]));
  
  const N2 = nextPowerOf2(Nd - 1);
  const a2 = a.concat(Array(N2 - N).fill([0, 0]));
  const b2 = b.concat(Array(N2 - Nd + 1).fill([0, 0]), b.slice(1).reverse());
  const A2 = fft1(a2);
  const B2 = fft1(b2);
  const AB2 = A2.map((A2n, n) => imul(A2n, B2[n]));
  const ab2 = ifft1(AB2);
  return bc.map((bcn, n) => imul(bcn, ab2[n]));
}
export function fft(f) {
  const N = f.length, T = 2 * Math.PI;
  const fftin = isPowerOf2(N) ? fftin1 : fftin2;
  return fftin(f, -2 * Math.PI);
}
export function ifft(F) {
  const N = F.length, T = 2 * Math.PI;
  const fftin = isPowerOf2(N) ? fftin1 : fftin2;
  return fftin(F, T).map(z => idivr(z, N));
}

main-large.mjs

// $ node --experimental-modules main.mjs
import {fft, ifft} from "./fft.mjs";

{
  // 15 elements example
  const fr0 = [...Array(1024 * 1024).keys()].map(i => Math.sin(i) * i);
  const f0 = fr0.map(r => [r, 0]);
  
  console.time("fft-ifft");
  const F = fft(f0);
  const f1 = ifft(F);
  console.timeEnd("fft-ifft");
  const fr1 = f1.map(([r]) => r);
}

main.mjs

// $ node --experimental-modules main.mjs
import {dft, idft} from "./dft.mjs";
import {fft, ifft} from "./fft.mjs";

{
  // 15 elements example
  const fr0 = [1,3,4,2, 5,6,2,4, 0,1,3,4, 5,62,2];
  const f0 = fr0.map(r => [r, 0]);
  
  const F = dft(f0);
  const f1 = idft(F);
  const fr1 = f1.map(([r]) => r);
  
  console.log("fr0:", fr0);
  console.log("F:", F);
  console.log("f1:", f1);
  console.log("fr1:", fr1.map(Math.round));
}

{
  // 15 elements example
  const fr0 = [1,3,4,2, 5,6,2,4, 0,1,3,4, 5,62,2];
  const f0 = fr0.map(r => [r, 0]);
  
  const F = fft(f0);
  const f1 = ifft(F);
  const fr1 = f1.map(([r]) => r);
  
  console.log("fr0:", fr0);
  console.log("F:", F);
  console.log("f1:", f1);
  console.log("fr1:", fr1.map(Math.round));
}

understanding-fft.ipynb