WebGPU-FFT-Ocean-Demo/shaders/ocean_spectrum.wgsl

const TWO_PI : f32 = 6.28318530718;
const GRAVITY : f32 = 9.81;

@group(0) @binding(0) var<storage, read> h0Real : array<f32>;
@group(0) @binding(1) var<storage, read> h0Imag : array<f32>;
@group(0) @binding(2) var<storage, read_write> spectrumReal : array<f32>;
@group(0) @binding(3) var<storage, read_write> spectrumImag : array<f32>;
@group(0) @binding(4) var<storage, read> params : array<f32>;

fn indexFromCoord( x : u32, y : u32, size : u32 ) -> u32 {

	return y * size + x;

}

fn mirrorIndex( x : u32, y : u32, size : u32 ) -> u32 {

	let mx : u32 = ( size - x ) % size;
	let my : u32 = ( size - y ) % size;

	return indexFromCoord( mx, my, size );

}

@compute @workgroup_size( 8, 8 )
fn main( @builtin(global_invocation_id) gid : vec3<u32> ) {

	let size : u32 = u32( params[ 1u ] );

	if ( gid.x >= size || gid.y >= size ) {
		return;
	}

	let nFloat : f32 = f32( size );
	let xPos : f32   = f32( gid.x );
	let yPos : f32   = f32( gid.y );

	var realValue : f32 = 0.0;
	var imagValue : f32 = 0.0;

	for ( var ky : u32 = 0u; ky < size; ky = ky + 1u ) {

		let kyShift : f32 = f32( i32( ky ) - i32( size ) / 2 );

		for ( var kx : u32 = 0u; kx < size; kx = kx + 1u ) {

			let kxShift : f32 = f32( i32( kx ) - i32( size ) / 2 );

			let idx : u32 = indexFromCoord( kx, ky, size );

			let mirrorIdx : u32 = mirrorIndex( kx, ky, size );

			let kLength : f32 = sqrt( kxShift * kxShift + kyShift * kyShift );

			if ( kLength == 0.0 ) {
				continue;
			}

			let h0R : f32 = h0Real[ idx ];
			let h0I : f32 = h0Imag[ idx ];

			let h0mR : f32 = h0Real[ mirrorIdx ];
			let h0mI : f32 = -h0Imag[ mirrorIdx ]; // conjugate

			let omega : f32 = sqrt( GRAVITY * kLength );

			let phase : f32 =
				omega * params[ 0u ] +
				( kxShift * TWO_PI * xPos / nFloat ) +
				( kyShift * TWO_PI * yPos / nFloat );

			let sinP : f32 = sin( phase );
			let cosP : f32 = cos( phase );

			let hPos : f32 = h0R * cosP - h0I * sinP;

			let hNeg : f32 = h0mR * cosP + h0mI * sinP;

			let hPosI : f32 = h0R * sinP + h0I * cosP;
			let hNegI : f32 = h0mR * sinP - h0mI * cosP;

			realValue = realValue + hPos + hNeg;
			imagValue = imagValue + hPosI + hNegI;

		}
	}

	let idxOut : u32 = indexFromCoord( gid.x, gid.y, size );

	spectrumReal[ idxOut ] = realValue * params[ 2u ];
	spectrumImag[ idxOut ] = imagValue * params[ 2u ];

}