diff --git a/benchmark/reimFFT.ts b/benchmark/reimFFT.ts new file mode 100644 index 00000000..1ac7d00f --- /dev/null +++ b/benchmark/reimFFT.ts @@ -0,0 +1,99 @@ +/* eslint-disable no-console */ +import { reimFFT } from '../src/reim/reimFFT.ts'; +import { reimArrayFFT } from '../src/reimArray/reimArrayFFT.ts'; + +const size = 2 ** 16; +const count = 10; // number of spectra in the array benchmark + +// Build input data +const spectra = Array.from({ length: count }, () => { + const re = new Float64Array(size); + const im = new Float64Array(size); + for (let i = 0; i < size; i++) { + re[i] = Math.random(); + im[i] = Math.random(); + } + return { re, im }; +}); + +// Warmup +for (const s of spectra) reimFFT(s); +for (const s of spectra) reimFFT(s, { inPlace: true }); +reimArrayFFT(spectra); +reimArrayFFT(spectra, { inPlace: true }); + +const targetMs = 5000; + +// --- reimFFT (loop over each spectrum individually) --- +{ + let iterations = 0; + const start = performance.now(); + console.time('reimFFT (loop)'); + while (performance.now() - start < targetMs) { + for (const s of spectra) reimFFT(s); + iterations++; + } + const elapsed = performance.now() - start; + console.timeEnd('reimFFT (loop)'); + console.log( + ` ${iterations * count} total FFTs, ${count} spectra × ${iterations} rounds`, + ); + console.log(` ${(elapsed / (iterations * count)).toFixed(3)} ms per FFT`); +} + +console.log(''); + +// --- reimFFT inPlace (loop over each spectrum individually) --- +{ + let iterations = 0; + const start = performance.now(); + console.time('reimFFT inPlace (loop)'); + while (performance.now() - start < targetMs) { + for (const s of spectra) reimFFT(s, { inPlace: true }); + iterations++; + } + const elapsed = performance.now() - start; + console.timeEnd('reimFFT inPlace (loop)'); + console.log( + ` ${iterations * count} total FFTs, ${count} spectra × ${iterations} rounds`, + ); + console.log(` ${(elapsed / (iterations * count)).toFixed(3)} ms per FFT`); +} + +console.log(''); + +// --- reimArrayFFT (single call for the whole array) --- +{ + let iterations = 0; + const start = performance.now(); + console.time('reimArrayFFT'); + while (performance.now() - start < targetMs) { + reimArrayFFT(spectra); + iterations++; + } + const elapsed = performance.now() - start; + console.timeEnd('reimArrayFFT'); + console.log( + ` ${iterations * count} total FFTs, ${count} spectra × ${iterations} rounds`, + ); + console.log(` ${(elapsed / (iterations * count)).toFixed(3)} ms per FFT`); +} + +console.log(''); + +// --- reimArrayFFT inPlace (single call for the whole array) --- +{ + let iterations = 0; + const start = performance.now(); + console.time('reimArrayFFT inPlace'); + while (performance.now() - start < targetMs) { + reimArrayFFT(spectra, { inPlace: true }); + iterations++; + } + const elapsed = performance.now() - start; + console.timeEnd('reimArrayFFT inPlace'); + console.log( + ` ${iterations * count} total FFTs, ${count} spectra × ${iterations} rounds`, + ); + console.log(` ${(elapsed / (iterations * count)).toFixed(3)} ms per FFT`); +} diff --git a/package.json b/package.json index 128a9c40..93990a7b 100644 --- a/package.json +++ b/package.json @@ -47,19 +47,19 @@ "ml-xsadd": "^3.0.1" }, "devDependencies": { - "@types/node": "^25.0.3", - "@vitest/coverage-v8": "^4.0.16", + "@types/node": "^25.3.0", + "@vitest/coverage-v8": "^4.0.18", "@zakodium/tsconfig": "^1.0.2", "cheminfo-build": "^1.3.2", "eslint": "^9.39.2", - "eslint-config-cheminfo-typescript": "^21.0.1", + "eslint-config-cheminfo-typescript": "^21.1.0", "jest-matcher-deep-close-to": "^3.0.2", "ml-spectra-fitting": "^5.0.1", - "prettier": "^3.7.4", - "rimraf": "^6.1.2", + "prettier": "^3.8.1", + "rimraf": "^6.1.3", "spectrum-generator": "^8.1.1", "typescript": "^5.9.3", - "vitest": "^4.0.16" + "vitest": "^4.0.18" }, "repository": { "type": "git", diff --git a/src/__tests__/__snapshots__/index.test.ts.snap b/src/__tests__/__snapshots__/index.test.ts.snap index 2dacb647..7330ba0b 100644 --- a/src/__tests__/__snapshots__/index.test.ts.snap +++ b/src/__tests__/__snapshots__/index.test.ts.snap @@ -7,6 +7,7 @@ exports[`existence of exported functions 1`] = ` "reimFFT", "reimPhaseCorrection", "reimZeroFilling", + "reimArrayFFT", "getOutputArray", "xAbsolute", "xAbsoluteMedian", diff --git a/src/index.ts b/src/index.ts index 9134d30d..abb3c006 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,4 +1,5 @@ export * from './reim/index.ts'; +export * from './reimArray/index.ts'; export * from './x/index.ts'; diff --git a/src/reim/__tests__/reimFFT.test.ts b/src/reim/__tests__/reimFFT.test.ts index 0af28fd4..7ab8ae4c 100644 --- a/src/reim/__tests__/reimFFT.test.ts +++ b/src/reim/__tests__/reimFFT.test.ts @@ -13,3 +13,29 @@ test('reimFFT', () => { expect(inverse.re).toStrictEqual(re); }); + +test('reimFFT inPlace: returns same object and mutates input arrays', () => { + const re = Float64Array.from([0, 3, 6, 5]); + const im = Float64Array.from([0, 4, 8, 3]); + const data = { re, im }; + + const result = reimFFT(data, { inPlace: true }); + + expect(result).toBe(data); + expect(result.re).toBe(re); + expect(result.im).toBe(im); +}); + +test('reimFFT inPlace: round-trip restores original values', () => { + const re = Float64Array.from([0, 3, 6, 5]); + const im = Float64Array.from([0, 4, 8, 3]); + const reOrig = Float64Array.from(re); + const imOrig = Float64Array.from(im); + const data = { re, im }; + + reimFFT(data, { inPlace: true, applyZeroShift: true }); + reimFFT(data, { inPlace: true, inverse: true, applyZeroShift: true }); + + expect(data.re).toStrictEqual(reOrig); + expect(data.im).toStrictEqual(imOrig); +}); diff --git a/src/reim/reimFFT.ts b/src/reim/reimFFT.ts index f1cd6c7b..8c2380ba 100644 --- a/src/reim/reimFFT.ts +++ b/src/reim/reimFFT.ts @@ -1,11 +1,17 @@ import FFT from 'fft.js'; import type { DataReIm } from '../types/index.ts'; -import { xRotate } from '../x/index.ts'; + +import { zeroShift } from './zeroShift.ts'; export interface ReimFFTOptions { inverse?: boolean; applyZeroShift?: boolean; + /** + * Write the result back into the input arrays instead of allocating new ones. + * @default false + */ + inPlace?: boolean; } /** @@ -18,7 +24,7 @@ export function reimFFT( data: DataReIm, options: ReimFFTOptions = {}, ): DataReIm { - const { inverse = false, applyZeroShift = false } = options; + const { inverse = false, applyZeroShift = false, inPlace = false } = options; const { re, im } = data; const size = re.length; @@ -40,6 +46,14 @@ export function reimFFT( if (applyZeroShift) output = zeroShift(output); } + if (inPlace) { + for (let i = 0; i < csize; i += 2) { + re[i >>> 1] = output[i]; + im[i >>> 1] = output[i + 1]; + } + return data as DataReIm; + } + const newRe = new Float64Array(size); const newIm = new Float64Array(size); for (let i = 0; i < csize; i += 2) { @@ -49,13 +63,3 @@ export function reimFFT( return { re: newRe, im: newIm }; } - -function zeroShift( - data: Float64Array, - inverse?: boolean, -): Float64Array { - const middle = inverse - ? Math.ceil(data.length / 2) - : Math.floor(data.length / 2); - return xRotate(data, middle); -} diff --git a/src/reim/zeroShift.ts b/src/reim/zeroShift.ts new file mode 100644 index 00000000..70aa838f --- /dev/null +++ b/src/reim/zeroShift.ts @@ -0,0 +1,12 @@ +import { xRotate } from '../x/index.ts'; + +export function zeroShift( + data: Float64Array, + inverse?: boolean, +): Float64Array { + const middle = inverse + ? Math.ceil(data.length / 2) + : Math.floor(data.length / 2); + + return xRotate(data, middle); +} diff --git a/src/reimArray/__tests__/reimArrayFFT.test.ts b/src/reimArray/__tests__/reimArrayFFT.test.ts new file mode 100644 index 00000000..778d48d1 --- /dev/null +++ b/src/reimArray/__tests__/reimArrayFFT.test.ts @@ -0,0 +1,169 @@ +import { expect, test } from 'vitest'; + +import { reimArrayFFT } from '../reimArrayFFT.ts'; + +test('reimArrayFFT: round-trip on a single spectrum', () => { + const re = Float64Array.from([0, 3, 6, 5]); + const im = Float64Array.from([0, 4, 8, 3]); + + const [transformed] = reimArrayFFT([{ re, im }], { applyZeroShift: true }); + const [inverse] = reimArrayFFT([transformed], { + inverse: true, + applyZeroShift: true, + }); + + expect(inverse.re).toStrictEqual(re); + expect(inverse.im).toStrictEqual(im); +}); + +test('reimArrayFFT: round-trip on multiple spectra, single FFT instance', () => { + const spectra = [ + { + re: Float64Array.from([1, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 1, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 3, 6, 5]), + im: Float64Array.from([0, 4, 8, 3]), + }, + ]; + + const transformed = reimArrayFFT(spectra); + const restored = reimArrayFFT(transformed, { inverse: true }); + + for (let i = 0; i < spectra.length; i++) { + expect(restored[i].re).toStrictEqual(spectra[i].re); + expect(restored[i].im).toStrictEqual(spectra[i].im); + } +}); + +test('reimArrayFFT: applyZeroShift round-trip on multiple spectra', () => { + const spectra = [ + { + re: Float64Array.from([0, 3, 6, 5]), + im: Float64Array.from([0, 4, 8, 3]), + }, + { + re: Float64Array.from([1, 2, 3, 4]), + im: Float64Array.from([0, 1, 0, 1]), + }, + ]; + + const transformed = reimArrayFFT(spectra, { applyZeroShift: true }); + const restored = reimArrayFFT(transformed, { + inverse: true, + applyZeroShift: true, + }); + + for (let i = 0; i < spectra.length; i++) { + expect(restored[i].re).toStrictEqual(spectra[i].re); + expect(restored[i].im).toStrictEqual(spectra[i].im); + } +}); + +test('reimArrayFFT: output arrays are independent (no shared buffers between results)', () => { + const spectra = [ + { + re: Float64Array.from([1, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 1, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + ]; + + const results = reimArrayFFT(spectra); + + expect(results[0].re).not.toBe(results[1].re); + expect(results[0].im).not.toBe(results[1].im); + expect(results[0].re).not.toBe(spectra[0].re); + expect(results[0].im).not.toBe(spectra[0].im); +}); + +test('reimArrayFFT: returns empty array for empty input', () => { + expect(reimArrayFFT([])).toStrictEqual([]); +}); + +test('reimArrayFFT: throws RangeError when elements have different lengths', () => { + const spectra = [ + { + re: Float64Array.from([1, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 1, 0, 0, 0, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0, 0, 0, 0, 0]), + }, + ]; + + expect(() => reimArrayFFT(spectra)).toThrowError(RangeError); +}); + +test('reimArrayFFT: throws RangeError indicating which element has the wrong length', () => { + const spectra = [ + { + re: Float64Array.from([1, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 1, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 0, 1, 0, 0, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0, 0, 0, 0, 0]), + }, + ]; + + expect(() => reimArrayFFT(spectra)).toThrowError(/element 2/); +}); + +test('reimArrayFFT inPlace: result shares re/im references with input', () => { + const spectra = [ + { + re: Float64Array.from([1, 0, 0, 0]), + im: Float64Array.from([0, 0, 0, 0]), + }, + { + re: Float64Array.from([0, 3, 6, 5]), + im: Float64Array.from([0, 4, 8, 3]), + }, + ]; + + const results = reimArrayFFT(spectra, { inPlace: true }); + + expect(results[0].re).toBe(spectra[0].re); + expect(results[0].im).toBe(spectra[0].im); + expect(results[1].re).toBe(spectra[1].re); + expect(results[1].im).toBe(spectra[1].im); +}); + +test('reimArrayFFT inPlace: round-trip restores original values', () => { + const spectra = [ + { + re: Float64Array.from([0, 3, 6, 5]), + im: Float64Array.from([0, 4, 8, 3]), + }, + { + re: Float64Array.from([1, 2, 3, 4]), + im: Float64Array.from([0, 1, 0, 1]), + }, + ]; + const originals = spectra.map((s) => ({ + re: Float64Array.from(s.re), + im: Float64Array.from(s.im), + })); + + reimArrayFFT(spectra, { inPlace: true, applyZeroShift: true }); + reimArrayFFT(spectra, { inPlace: true, inverse: true, applyZeroShift: true }); + + for (let i = 0; i < spectra.length; i++) { + expect(spectra[i].re).toStrictEqual(originals[i].re); + expect(spectra[i].im).toStrictEqual(originals[i].im); + } +}); diff --git a/src/reimArray/index.ts b/src/reimArray/index.ts new file mode 100644 index 00000000..32f99d38 --- /dev/null +++ b/src/reimArray/index.ts @@ -0,0 +1 @@ +export * from './reimArrayFFT.ts'; diff --git a/src/reimArray/reimArrayFFT.ts b/src/reimArray/reimArrayFFT.ts new file mode 100644 index 00000000..3645376d --- /dev/null +++ b/src/reimArray/reimArrayFFT.ts @@ -0,0 +1,83 @@ +import FFT from 'fft.js'; + +import { zeroShift } from '../reim/zeroShift.ts'; +import type { DataReIm } from '../types/index.ts'; + +export interface ReimArrayFFTOptions { + inverse?: boolean; + applyZeroShift?: boolean; + /** + * Write the result back into the input arrays instead of allocating new ones. + * @default false + */ + inPlace?: boolean; +} + +/** + * Apply FFT to an array of complex spectra, reusing a single FFT instance for + * all elements to avoid repeated instantiation overhead. + * All elements must have the same length. + * @param data - array of complex spectra + * @param options - options + * @returns array of transformed complex spectra + */ +export function reimArrayFFT( + data: DataReIm[], + options: ReimArrayFFTOptions = {}, +): Array> { + if (data.length === 0) return []; + + const { inverse = false, applyZeroShift = false, inPlace = false } = options; + + const size = data[0].re.length; + const csize = size << 1; + + for (let j = 1; j < data.length; j++) { + if (data[j].re.length !== size || data[j].im.length !== size) { + throw new RangeError( + `All elements must have the same length. Expected ${size} but element ${j} has length ${data[j].re.length}.`, + ); + } + } + + // Single FFT instance and working buffers reused across all spectra + const fft = new FFT(size); + const complexArray = new Float64Array(csize); + const output = new Float64Array(csize); + + const results = new Array>(data.length); + + for (let j = 0; j < data.length; j++) { + const { re, im } = data[j]; + + for (let i = 0; i < csize; i += 2) { + complexArray[i] = re[i >>> 1]; + complexArray[i + 1] = im[i >>> 1]; + } + + const outRe = inPlace ? re : new Float64Array(size); + const outIm = inPlace ? im : new Float64Array(size); + + if (inverse) { + const input = applyZeroShift + ? zeroShift(complexArray, true) + : complexArray; + fft.inverseTransform(output, input); + for (let i = 0; i < csize; i += 2) { + outRe[i >>> 1] = output[i]; + outIm[i >>> 1] = output[i + 1]; + } + } else { + fft.transform(output, complexArray); + const source = applyZeroShift ? zeroShift(output) : output; + for (let i = 0; i < csize; i += 2) { + outRe[i >>> 1] = source[i]; + outIm[i >>> 1] = source[i + 1]; + } + } + + results[j] = { re: outRe as Float64Array, im: outIm as Float64Array }; + } + + return results; +} diff --git a/tsconfig.json b/tsconfig.json index a68fc7e1..c2cf2bd5 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -5,5 +5,5 @@ "outDir": "lib", "types": ["node"] }, - "include": ["src", "vite*.ts"] + "include": ["src", "benchmark", "vite*.ts"] }