Add stft op (#1746)

FEATURE

Add tf.signal.stft op. 

One TODO is passing fft length parameter because rfft does not support fft length parameter. We can pass fft length parameter after rfft supports it.

See: tensorflow/tfjs#1362

Author: Lewuathe

src/ops/signal_ops.ts

@@ -18,8 +18,10 @@
 import {op} from '../ops/operation';
 import {Tensor, Tensor1D} from '../tensor';
 
+import {mul} from './binary_ops';
 import {concat} from './concat_split';
 import {slice} from './slice';
+import {rfft} from './spectral_ops';
 import {fill, tensor1d, tensor2d} from './tensor_ops';
 
 /**
@@ -94,6 +96,45 @@ function frame_(
   return concat(output).as2D(output.length, frameLength);
 }
 
+/**
+ * Computes the Short-time Fourier Transform of signals
+ * See: https://en.wikipedia.org/wiki/Short-time_Fourier_transform
+ *
+ * ```js
+ * const input = tf.tensor1d([1, 1, 1, 1, 1])
+ * tf.signal.stft(input, 3, 1).print();
+ * ```
+ * @param signal 1-dimensional real value tensor.
+ * @param frameLength The window length of samples.
+ * @param frameStep The number of samples to step.
+ * @param fftLength The size of the FFT to apply.
+ * @param windowFn A callable that takes a window length and returns 1-d tensor.
+ */
+/**
+ * @doc {heading: 'Operations', subheading: 'Signal', namespace: 'signal'}
+ */
+function stft_(
+    signal: Tensor1D, frameLength: number, frameStep: number,
+    fftLength?: number,
+    windowFn: (length: number) => Tensor1D = hannWindow): Tensor {
+  if (fftLength == null) {
+    fftLength = enclosingPowerOfTwo(frameLength);
+  }
+  const framedSignal = frame(signal, frameLength, frameStep);
+  const windowedSignal = mul(framedSignal, windowFn(frameLength));
+  const output: Tensor[] = [];
+  for (let i = 0; i < framedSignal.shape[0]; i++) {
+    output.push(rfft(windowedSignal.slice([i, 0], [1, frameLength]),
+      fftLength));
+  }
+  return concat(output);
+}
+
+function enclosingPowerOfTwo(value: number) {
+  // Return 2**N for integer N such that 2**N >= value.
+  return Math.floor(Math.pow(2, Math.ceil(Math.log(value) / Math.log(2.0))));
+}
+
 function cosineWindow(windowLength: number, a: number, b: number): Tensor1D {
   const even = 1 - windowLength % 2;
   const newValues = new Float32Array(windowLength);
@@ -107,3 +148,4 @@ function cosineWindow(windowLength: number, a: number, b: number): Tensor1D {
 export const hannWindow = op({hannWindow_});
 export const hammingWindow = op({hammingWindow_});
 export const frame = op({frame_});
+export const stft = op({stft_});

src/ops/signal_ops_test.ts

@@ -124,3 +124,166 @@ describeWithFlags('frame', ALL_ENVS, () => {
     expectArraysClose(await output.data(), [1, 2, 3, 4, 5, 100]);
   });
 });
+
+describeWithFlags('stft', ALL_ENVS, () => {
+  it('3 length with hann window', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 3;
+    const frameStep = 1;
+    const output = tf.signal.stft(input, frameLength, frameStep);
+    expect(output.shape).toEqual([3, 3]);
+    expectArraysClose(await output.data(), [
+      1.0, 0.0, 0.0, -1.0, -1.0, 0.0,
+      1.0, 0.0, 0.0, -1.0, -1.0, 0.0,
+      1.0, 0.0, 0.0, -1.0, -1.0, 0.0,
+    ]);
+  });
+
+  it('3 length with hann window (sequencial number)', async () => {
+    const input = tf.tensor1d([1, 2, 3, 4, 5]);
+    const frameLength = 3;
+    const frameStep = 1;
+    const output = tf.signal.stft(input, frameLength, frameStep);
+    expect(output.shape).toEqual([3, 3]);
+    expectArraysClose(await output.data(), [
+      2.0, 0.0, 0.0, -2.0, -2.0, 0.0,
+      3.0, 0.0, 0.0, -3.0, -3.0, 0.0,
+      4.0, 0.0, 0.0, -4.0, -4.0, 0.0
+    ]);
+  });
+
+  it('3 length, 2 step with hann window', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 3;
+    const frameStep = 2;
+    const output = tf.signal.stft(input, frameLength, frameStep);
+    expect(output.shape).toEqual([2, 3]);
+    expectArraysClose(await output.data(), [
+      1.0, 0.0, 0.0, -1.0, -1.0, 0.0,
+      1.0, 0.0, 0.0, -1.0, -1.0, 0.0
+    ]);
+  });
+
+  it('3 fftLength, 5 frameLength, 2 step', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1, 1]);
+    const frameLength = 5;
+    const frameStep = 1;
+    const fftLength = 3;
+    const output = tf.signal.stft(input, frameLength, frameStep, fftLength);
+    expect(output.shape[0]).toEqual(2);
+    expectArraysClose(await output.data(), [
+      1.5, 0.0, -0.749999, 0.433,
+      1.5, 0.0, -0.749999, 0.433
+    ]);
+  });
+
+  it('5 length with hann window', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 5;
+    const frameStep = 1;
+    const output = tf.signal.stft(input, frameLength, frameStep);
+    expect(output.shape).toEqual([1, 5]);
+    expectArraysClose(
+      await output.data(),
+      [2.0, 0.0, 0.0, -1.7071068, -1.0, 0.0, 0.0, 0.29289323, 0.0, 0.0]);
+  });
+
+  it('5 length with hann window (sequential)', async () => {
+    const input = tf.tensor1d([1, 2, 3, 4, 5]);
+    const frameLength = 5;
+    const frameStep = 1;
+    const output = tf.signal.stft(input, frameLength, frameStep);
+    expect(output.shape).toEqual([1, 5]);
+    expectArraysClose(
+      await output.data(),
+      [6.0, 0.0, -0.70710677, -5.1213202, -3.0, 1.0,
+        0.70710677, 0.87867975, 0.0, 0.0]);
+  });
+
+  it('3 length with hamming window', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 3;
+    const frameStep = 1;
+    const fftLength = 3;
+    const output = tf.signal.stft(input, frameLength, frameStep,
+      fftLength, (length) => tf.signal.hammingWindow(length));
+    expect(output.shape).toEqual([3, 2]);
+    expectArraysClose(await output.data(), [
+      1.16, 0.0, -0.46, -0.79674333,
+      1.16, 0.0, -0.46, -0.79674333,
+      1.16, 0.0, -0.46, -0.79674333
+    ]);
+  });
+
+  it('3 length, 2 step with hamming window', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 3;
+    const frameStep = 2;
+    const fftLength = 3;
+    const output = tf.signal.stft(input, frameLength, frameStep,
+      fftLength, (length) => tf.signal.hammingWindow(length));
+    expect(output.shape).toEqual([2, 2]);
+    expectArraysClose(await output.data(), [
+      1.16, 0.0, -0.46, -0.79674333,
+      1.16, 0.0, -0.46, -0.79674333
+    ]);
+  });
+
+  it('3 fftLength, 5 frameLength, 2 step with hamming window', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1, 1]);
+    const frameLength = 5;
+    const frameStep = 1;
+    const fftLength = 3;
+    const output = tf.signal.stft(input, frameLength, frameStep,
+      fftLength, (length) => tf.signal.hammingWindow(length));
+    expect(output.shape).toEqual([2, 2]);
+    expectArraysClose(await output.data(), [
+      1.619999, 0.0, -0.69, 0.39837,
+      1.619999, 0.0, -0.69, 0.39837
+    ]);
+  });
+
+  it('5 length with hann window (sequential)', async () => {
+    const input = tf.tensor1d([1, 2, 3, 4, 5]);
+    const frameLength = 5;
+    const frameStep = 1;
+    const fftLength = 5;
+    const output = tf.signal.stft(input, frameLength, frameStep,
+      fftLength, (length) => tf.signal.hammingWindow(length));
+    expect(output.shape).toEqual([1, 3]);
+    expectArraysClose(
+      await output.data(),
+      [6.72, 0.0, -3.6371822, -1.1404576, 0.4771822, 0.39919350]);
+  });
+
+  it('3 length without window function', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 3;
+    const frameStep = 1;
+    const fftLength = 3;
+    const ident = (length: number) => tf.ones([length]).as1D();
+    const output = tf.signal.stft(input, frameLength, frameStep,
+      fftLength, ident);
+    expect(output.shape).toEqual([3, 2]);
+    expectArraysClose(await output.data(), [
+      3.0, 0.0, 0.0, 0.0,
+      3.0, 0.0, 0.0, 0.0,
+      3.0, 0.0, 0.0, 0.0
+    ]);
+  });
+
+  it('3 length, 2 step without window function', async () => {
+    const input = tf.tensor1d([1, 1, 1, 1, 1]);
+    const frameLength = 3;
+    const frameStep = 2;
+    const fftLength = 3;
+    const ident = (length: number) => tf.ones([length]).as1D();
+    const output = tf.signal.stft(input, frameLength, frameStep,
+      fftLength, ident);
+    expect(output.shape).toEqual([2, 2]);
+    expectArraysClose(await output.data(), [
+      3.0, 0.0, 0.0, 0.0,
+      3.0, 0.0, 0.0, 0.0
+    ]);
+  });
+});