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

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 kx : u32 = gid.x;
	let ky : u32 = gid.y;

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

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

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

	if ( kLength == 0.0 ) {
		spectrumReal[ idx ] = 0.0;
		spectrumImag[ idx ] = 0.0;
		return;
	}

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

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

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

	let omega   : f32 = sqrt( GRAVITY * kLength );
	let time    : f32 = params[ 0u ];
	let heightA : f32 = params[ 2u ];

	let theta : f32 = omega * time;

	let cosT : f32 = cos( theta );
	let sinT : f32 = sin( theta );

	// h0(k) * e^{ i * omega * t }
	let hPosR : f32 = h0R * cosT - h0I * sinT;
	let hPosI : f32 = h0R * sinT + h0I * cosT;

	// h0(-k) * e^{ -i * omega * t }
	let hNegR : f32 = h0mR * cosT + h0mI * sinT;
	let hNegI : f32 = h0mI * cosT - h0mR * sinT;

	let outR : f32 = ( hPosR + hNegR ) * heightA;
	let outI : f32 = ( hPosI + hNegI ) * heightA;

	spectrumReal[ idx ] = outR;
	spectrumImag[ idx ] = outI;

}