WebGPU-FFT-Ocean-Demo/tests/OceanTests.js

export class OceanTests {

	constructor( getPrimaryPipeline ) {

		this.getPrimaryPipeline	= getPrimaryPipeline;

	}


	async testFft( ) {

		try {

			const pipeline		= this.getPrimaryPipeline( );

			if ( !pipeline ) {

				console.warn( "No primary pipeline for testFft()" );
				return;

			}

			const N				= pipeline.gridSize;

			const block			= pipeline.getBlockByName( "ocean" );

			const rowPass		= block.getPass( "RowFFT" );

			const colPass		= block.getPass( "ColFFT" );

			const inputReal		= new Float32Array( N * N );
			const inputImag		= new Float32Array( N * N );

			for ( let y = 0; y < N; y++ ) {

				for ( let x = 0; x < N; x++ ) {

					const idx			= y * N + x;

					inputReal[ idx ]	= Math.sin( 2 * Math.PI * x / N ) + 0.5 * Math.cos( 2 * Math.PI * y / N );
					inputImag[ idx ]	= 0;

				}

			}

			pipeline.memory.spectrumReal.set( inputReal );
			pipeline.memory.spectrumImag.set( inputImag );

			const spectrumPass	= block.getPass( "Spectrum" );

			spectrumPass.shader.setVariable( "spectrumReal", pipeline.memory.spectrumReal );
			spectrumPass.shader.setVariable( "spectrumImag", pipeline.memory.spectrumImag );

			await rowPass.execute( );
			await colPass.execute( );

			const gpuHeight		= await colPass.shader.debugBuffer( "heightField" );

			const cpuResult		= this.cpuFft2D( inputReal, inputImag, N );

			let maxDiff			= 0;
			let rms				= 0;

			for ( let i = 0; i < gpuHeight.length; i++ ) {

				const diff	= gpuHeight[ i ] - cpuResult.real[ i ];

				maxDiff		= Math.max( maxDiff, Math.abs( diff ) );
				rms			+= diff * diff;

			}

			rms = Math.sqrt( rms / gpuHeight.length );

			console.log( "FFT test completed. maxDiff:", maxDiff, "rms:", rms );

		} catch ( e ) {

			console.error( "FFT test failed:", e );

		}

	}


	async testSpectrum( ) {

		try {

			const pipeline		= this.getPrimaryPipeline( );

			if ( !pipeline ) {

				console.warn( "No primary pipeline for testSpectrum()" );
				return;

			}

			const N				= pipeline.gridSize;

			const block			= pipeline.getBlockByName( "ocean" );

			const spectrumPass	= block.getPass( "Spectrum" );

			const h0Real		= pipeline.memory.h0Real;
			const h0Imag		= pipeline.memory.h0Imag;

			const time			= 0;
			const heightScale	= pipeline.memory.computeParams[ 2 ];

			const cpuSpec		= this.cpuSpectrum2D( h0Real, h0Imag, N, time, heightScale );

			pipeline.memory.computeParams[ 0 ]	= time;
			pipeline.memory.computeParams[ 1 ]	= N;
			pipeline.memory.computeParams[ 2 ]	= heightScale;

			spectrumPass.shader.setVariable( "params", pipeline.memory.computeParams );

			await spectrumPass.execute( );

			const gpuReal		= await spectrumPass.shader.debugBuffer( "spectrumReal" );
			const gpuImag		= await spectrumPass.shader.debugBuffer( "spectrumImag" );

			let maxDiff			= 0;
			let rms				= 0;

			for ( let i = 0; i < gpuReal.length; i++ ) {

				const dr		= gpuReal[ i ] - cpuSpec.real[ i ];
				const di		= gpuImag[ i ] - cpuSpec.imag[ i ];

				const magDiff	= Math.sqrt( dr * dr + di * di );

				if ( magDiff > maxDiff ) {

					maxDiff	= magDiff;

				}

				rms			+= dr * dr + di * di;

			}

			rms = Math.sqrt( rms / gpuReal.length );

			console.log( "Spectrum test completed. maxDiff:", maxDiff, "rms:", rms );

		} catch ( e ) {

			console.error( "Spectrum test failed:", e );

		}

	}

	cpuFft1D( realIn, imagIn, N ) {

		const outR	= new Float32Array( N );
		const outI	= new Float32Array( N );

		for ( let k = 0; k < N; k++ ) {

			let sumR	= 0;
			let sumI	= 0;

			for ( let n = 0; n < N; n++ ) {

				const angle	= -2 * Math.PI * k * n / N;

				const c		= Math.cos( angle );
				const s		= Math.sin( angle );

				const xr	= realIn[ n ];
				const xi	= imagIn[ n ];

				sumR		+= xr * c - xi * s;
				sumI		+= xr * s + xi * c;

			}

			const invScale	= 1 / N;

			outR[ k ]		= sumR * invScale;
			outI[ k ]		= sumI * invScale;

		}

		return { real: outR, imag: outI };

	}


	cpuFft2D( realIn, imagIn, N ) {

		const realTmp	= new Float32Array( realIn );
		const imagTmp	= new Float32Array( imagIn );

		for ( let row = 0; row < N; row++ ) {

			const rowOffset	= row * N;

			const rSlice	= realTmp.subarray( rowOffset, rowOffset + N );
			const iSlice	= imagTmp.subarray( rowOffset, rowOffset + N );

			const res		= this.cpuFft1D( rSlice, iSlice, N );

			for ( let x = 0; x < N; x++ ) {

				realTmp[ rowOffset + x ]	= res.real[ x ];
				imagTmp[ rowOffset + x ]	= res.imag[ x ];

			}

		}

		const realOut	= new Float32Array( N * N );
		const imagOut	= new Float32Array( N * N );

		for ( let col = 0; col < N; col++ ) {

			const rCol	= new Float32Array( N );
			const iCol	= new Float32Array( N );

			for ( let row = 0; row < N; row++ ) {

				const idx	= row * N + col;

				rCol[ row ]	= realTmp[ idx ];
				iCol[ row ]	= imagTmp[ idx ];

			}

			const res	= this.cpuFft1D( rCol, iCol, N );

			for ( let row = 0; row < N; row++ ) {

				const idx		= row * N + col;

				realOut[ idx ]	= res.real[ row ];
				imagOut[ idx ]	= res.imag[ row ];

			}

		}

		return { real: realOut, imag: imagOut };

	}


	cpuSpectrum2D( h0Real, h0Imag, size, time, heightScale ) {

		const GRAVITY	= 9.81;

		const realOut	= new Float32Array( size * size );
		const imagOut	= new Float32Array( size * size );

		for ( let ky = 0; ky < size; ky++ ) {

			for ( let kx = 0; kx < size; kx++ ) {

				const kxShift	= kx - size / 2;
				const kyShift	= ky - size / 2;

				const kLength	= Math.sqrt( kxShift * kxShift + kyShift * kyShift );

				const idx		= ky * size + kx;

				if ( kLength === 0 ) {

					continue;

				}

				const mx		= ( size - kx ) % size;
				const my		= ( size - ky ) % size;

				const mirrorIdx	= my * size + mx;

				const h0R		= h0Real[ idx ];
				const h0I		= h0Imag[ idx ];

				const h0mR		= h0Real[ mirrorIdx ];
				const h0mI		= h0Imag[ mirrorIdx ];

				const omega		= Math.sqrt( GRAVITY * kLength );

				const theta		= omega * time;

				const cosT		= Math.cos( theta );
				const sinT		= Math.sin( theta );

				const hPosR		= h0R * cosT - h0I * sinT;
				const hPosI		= h0R * sinT + h0I * cosT;

				const hNegR		= h0mR * cosT + h0mI * sinT;
				const hNegI		= h0mI * cosT - h0mR * sinT;

				realOut[ idx ]	= ( hPosR + hNegR ) * heightScale;
				imagOut[ idx ]	= ( hPosI + hNegI ) * heightScale;

			}

		}

		return { real: realOut, imag: imagOut };

	}

}