WebGPU-FFT-Ocean-Demo/pipelines/OceanPipeline.js

import { RenderPipeline }	from "../framework/RenderPipeline.js";

import { Block }			from "../framework/Block.js";

import Camera				from "../framework/Camera.js";

import EventManager			from "../framework/eventManager.js";

import { SpectrumPass }		from "../passes/SpectrumPass.js";

import { RowFFTPass }		from "../passes/RowFFTPass.js";

import { ColFFTPass }		from "../passes/ColFFTPass.js";

import { OceanRenderPass }	from "../passes/OceanRenderPass.js";

import { OceanSolidRenderPass }	from "../passes/OceanSolidRenderPass.js";

import { SkySpherePass }	from "../passes/SkySpherePass.js";


export class OceanPipeline extends RenderPipeline {

	constructor( engine, canvas ) {

		super( engine );

		this.canvas				= canvas;
		this.canvasConfigured	= false;

		this.gridSize			= 64;
		this.meshResolution		= 64;

		this.offsetX			= 0;
		this.offsetZ			= 0;

		this.wavelengthScale	= 1.0;

		this.tiling				= 1;

		this.handleInput		= true;

		this.renderMode			= "solid";
		this.shadingMode		= "lighting";

		this.isPaused			= false;
		this.elapsedTime		= 0;

	}


	async create( ) {

		this.patchSize		= 80;
		this.heightScale	= 54.0;
		this.timeScale		= 0.35;

		this.startTime		= performance.now();

		const simSize		= this.gridSize;

		this.memory.set( "h0Real", new Float32Array( simSize * simSize ) );

		this.memory.set( "h0Imag", new Float32Array( simSize * simSize ) );

		this.memory.set( "spectrumReal", new Float32Array( simSize * simSize ) );

		this.memory.set( "spectrumImag", new Float32Array( simSize * simSize ) );

		this.memory.set( "rowReal", new Float32Array( simSize * simSize ) );

		this.memory.set( "rowImag", new Float32Array( simSize * simSize ) );

		this.memory.set( "colReal", new Float32Array( simSize * simSize ) );

		this.memory.set( "colImag", new Float32Array( simSize * simSize ) );

		this.memory.set( "heightField", new Float32Array( simSize * simSize ) );

		this.memory.set( "computeParams", new Float32Array( 4 ) );

		this.memory.set( "renderParams", new Float32Array( [ simSize, this.patchSize, this.heightScale, 0, this.meshResolution, this.tiling, this.wavelengthScale ] ) );

		this._seedSpectrum();

		const geometry	= this._buildGridGeometry( this.meshResolution, this.patchSize );

		this.memory.set( "positions", geometry.positions );

		this.memory.set( "uvs", geometry.uvs );

		this.memory.set( "indices", geometry.indices );

		this.memory.set( "lineIndices", geometry.lineIndices );

		const skySphere	= this._buildSkySphereGeometry( 220, 32, 64 );

		this.memory.set( "skyPositions", skySphere.positions );

		this.memory.set( "skyNormals", skySphere.normals );

		this.memory.set( "skyIndices", skySphere.indices );

		const cubeGeometry			= this._buildCubeGeometry();

		this.memory.set( "cubePositions", cubeGeometry.positions );

		this.memory.set( "cubeColors", cubeGeometry.colors );

		this.memory.set( "cubeIndices", cubeGeometry.indices );


		const block				= new Block( "ocean", this );

		const spectrumPass		= new SpectrumPass( );

		const rowFFTPass		= new RowFFTPass( );

		const colFFTPass		= new ColFFTPass( );

		const renderWirePass	= new OceanRenderPass( );

		const renderSolidPass	= new OceanSolidRenderPass( );

	 	const skySpherePass		= new SkySpherePass( );

		block.addPass( "Spectrum", spectrumPass );

		block.addPass( "RowFFT", rowFFTPass );

		block.addPass( "ColFFT", colFFTPass );

		block.addPass( "RenderWire", renderWirePass );

		block.addPass( "RenderSolid", renderSolidPass );

		block.addPass( "SkySphere", skySpherePass );

		this.addBlock( block );

		this.camera				= new Camera( [ 0, 90, 120 ], [ 0, 0, 0 ], [ 0, 1, 0 ] );

		this.camera.far		= 6000.0;

		this.camera.pitch		= Math.PI / 3;

		this.camera.update();

		this.eventManager		= new EventManager( );

		if ( this.canvas && this.handleInput ) {

			this.eventManager.setup( this.canvas, this.camera );

			this.eventManager.registerEventListeners();

		}

		await super.create();

	}


	setHeightScale( value ) {

		this.heightScale					= value;

		this.memory.computeParams[ 2 ]		= value;
		this.memory.renderParams[ 2 ]		= value;

	}


	setWavelengthScale( value ) {

		if ( value <= 0 ) {

			value = 0.25;

		}

		this.wavelengthScale	= value;

		if ( this.memory && this.memory.renderParams ) {

			this.memory.renderParams[ 6 ] = this.wavelengthScale;

		}

		this._seedSpectrum( );

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

		if ( block ) {

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

			if ( spectrumPass && spectrumPass.shader ) {

				spectrumPass.shader.setVariable( "h0Real", this.memory.h0Real );

				spectrumPass.shader.setVariable( "h0Imag", this.memory.h0Imag );

			}

		}

	}


	setRenderMode( mode ) {

		if ( mode !== "wireframe" && mode !== "solid" ) {

			return;

		}

		this.renderMode	= mode;

	}


	setShadingMode( mode ) {

		let code = 0;

		if ( mode === "normals" ) {

			code = 1;

		} else if ( mode === "solid" ) {

			code = 2;

		} else if ( mode === "height" ) {

			code = 3;

		} else if ( mode === "realistic" ) {

			code = 4;

		} else {

			mode = "lighting";
			code = 0;

		}

		this.shadingMode				= mode;
		this.memory.renderParams[ 3 ]	= code;

	}


	setTiling( value ) {

		if ( value < 1 ) {

			value = 1;

		}

		this.tiling					= value;
		this.memory.renderParams[ 5 ]	= value;

	}

	setOffset( x, z ) {

		this.offsetX					= x;
		this.offsetZ					= z;

		// Offsets are now handled in the vertex shader via instancing.

	}


	setPaused( paused ) {

		if ( paused === this.isPaused ) {

			return;

		}

		this.isPaused	= paused;

		if ( !paused ) {

			this.startTime	= performance.now() - this.elapsedTime * 1000.0;

		}

	}


	async stepOnce( stepSeconds = 1 / 60 ) {

		if ( !this.isPaused ) {

			return;

		}

		this.elapsedTime					+= stepSeconds;

		this.memory.computeParams[ 0 ]		= this.elapsedTime * this.timeScale;
		this.memory.computeParams[ 1 ]		= this.gridSize;

		await this.bindBuffers( );
		await this.execute( );

	}


	async bindBuffers( ) {

		const nowSeconds			= ( performance.now() - this.startTime ) / 1000;

		if ( !this.isPaused ) {

			this.elapsedTime		= nowSeconds;

		}

		this.memory.computeParams[ 0 ]	= this.elapsedTime * this.timeScale;
		this.memory.computeParams[ 1 ]	= this.gridSize;


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

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

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

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

		const renderWirePass	= block.getPass( "RenderWire" );

		const renderSolidPass	= block.getPass( "RenderSolid" );

		const skyPass			= block.getPass( "SkySphere" );


		await spectrumPass.bindBuffers( );

		await rowFFTPass.bindBuffers( );

		await colFFTPass.bindBuffers( );


		if ( this.renderMode === "solid" ) {

			await renderSolidPass.bindBuffers( );

		} else {

			await renderWirePass.bindBuffers( );

		}

		if ( skyPass ) {

			await skyPass.bindBuffers( );

		}

	}


	async execute( ) {

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

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

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

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


		await spectrumPass.execute( );

		await rowFFTPass.execute( );

		await colFFTPass.execute( );

 }


	_seedSpectrum( ) {

		const size			= this.gridSize;

		const phillipsA		= 0.0006;

		const windSpeed		= 24.0;

		const windDir		= { x: 0.8, y: 0.6 };

		const baseL			= windSpeed * windSpeed / 9.81;

		let wlScale			= this.wavelengthScale;

		if ( wlScale <= 0 ) {

			wlScale = 0.25;

		}

		const L				= baseL;

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

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

				const kx			= ( x - size / 2 ) * ( Math.PI * 2 / this.patchSize );

				const ky			= ( y - size / 2 ) * ( Math.PI * 2 / this.patchSize );

				const kLength		= Math.sqrt( kx * kx + ky * ky );

				if ( kLength === 0 ) {

					continue;

				}

				const kxNorm		= kx / kLength;

				const kyNorm		= ky / kLength;

				const windDotK		= kxNorm * windDir.x + kyNorm * windDir.y;

				const kScaled		= kLength / wlScale;

				const phillips		= phillipsA * Math.exp( -1 / ( kScaled * kScaled * L * L ) ) / Math.pow( kScaled, 4 ) * ( windDotK * windDotK );

				const gaussianR		= this._gaussianRandom();

				const gaussianI		= this._gaussianRandom();

				const amplitude		= Math.sqrt( phillips ) * Math.SQRT1_2;

				const idx			= y * size + x;

				this.memory.h0Real[ idx ]	= gaussianR * amplitude;

				this.memory.h0Imag[ idx ]	= gaussianI * amplitude;

			}
		}

		// enforce conjugate symmetry: h0(-k) = conj( h0(k) ) for a real height field
		for ( let y = 0; y < size; y++ ) {

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

				const idx			= y * size + x;

				const mx			= ( size - x ) % size;

				const my			= ( size - y ) % size;

				const mirrorIdx		= my * size + mx;

				this.memory.h0Real[ mirrorIdx ]	= this.memory.h0Real[ idx ];

				this.memory.h0Imag[ mirrorIdx ]	= -this.memory.h0Imag[ idx ];

			}
		}

		this.memory.computeParams[ 0 ] = 0;
		this.memory.computeParams[ 1 ] = this.gridSize;
		this.memory.computeParams[ 2 ] = this.heightScale;
		this.memory.computeParams[ 3 ] = this.patchSize;
		this.memory.renderParams[ 2 ]  = this.heightScale;

	}


	_gaussianRandom( ) {

		let u = 0, v = 0;

		while ( u === 0 ) u = Math.random();

		while ( v === 0 ) v = Math.random();

		return Math.sqrt( -2.0 * Math.log( u ) ) * Math.cos( 2.0 * Math.PI * v );

	}


	_buildGridGeometry( size, span ) {

		const vertexCount	= size * size;

		const positions		= new Float32Array( vertexCount * 3 );

		const uvs			= new Float32Array( vertexCount * 2 );

		const quadCount		= ( size - 1 ) * ( size - 1 );

		const indices		= new Uint32Array( quadCount * 6 );

		const horizontalLines	= size * ( size - 1 );

		const verticalLines		= ( size - 1 ) * size;

		const lineIndices		= new Uint32Array( ( horizontalLines + verticalLines ) * 2 );

		let pi = 0;
		let ui = 0;

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

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

				const fx	= x / ( size - 1 );

				const fy	= y / ( size - 1 );

				const worldX = ( fx - 0.5 ) * span;

				const worldZ = ( fy - 0.5 ) * span;

				positions[ pi++ ] = worldX;
				positions[ pi++ ] = 0;
				positions[ pi++ ] = worldZ;

				uvs[ ui++ ] = x;
				uvs[ ui++ ] = y;

			}
		}

		let ii = 0;

		let li = 0;

		for ( let y = 0; y < size - 1; y++ ) {

			for ( let x = 0; x < size - 1; x++ ) {

				const topLeft		= y * size + x;

				const topRight		= topLeft + 1;

				const bottomLeft	= topLeft + size;

				const bottomRight	= bottomLeft + 1;

				indices[ ii++ ] = topLeft;
				indices[ ii++ ] = bottomLeft;
				indices[ ii++ ] = topRight;

				indices[ ii++ ] = topRight;
				indices[ ii++ ] = bottomLeft;
				indices[ ii++ ] = bottomRight;

			}

			// horizontal line for row y at each segment
			for ( let x = 0; x < size - 1; x++ ) {

				const a = y * size + x;

				const b = a + 1;

				lineIndices[ li++ ] = a;
				lineIndices[ li++ ] = b;

			}

		}

		// vertical lines
		for ( let x = 0; x < size; x++ ) {

			for ( let y = 0; y < size - 1; y++ ) {

				const a = y * size + x;

				const b = a + size;

				lineIndices[ li++ ] = a;
				lineIndices[ li++ ] = b;

			}
		}

		return {
			positions,
			uvs,
			indices,
			lineIndices
		};

	}


	_buildSkySphereGeometry( radius, latSegments, lonSegments ) {

		const latCount	= latSegments;

		const lonCount	= lonSegments;

		const vertexCount	= ( latCount + 1 ) * ( lonCount + 1 );

		const positions	= new Float32Array( vertexCount * 3 );

		const normals		= new Float32Array( vertexCount * 3 );

		const uvs			= new Float32Array( vertexCount * 2 );

		const indices		= new Uint32Array( latCount * lonCount * 6 );

		let pi	= 0;

		let ni	= 0;

		let ui	= 0;

		for ( let lat = 0; lat <= latCount; lat++ ) {

			const theta	= lat * Math.PI / latCount;

			const sinTheta	= Math.sin( theta );

			const cosTheta	= Math.cos( theta );

			for ( let lon = 0; lon <= lonCount; lon++ ) {

				const phi	= lon * Math.PI * 2.0 / lonCount;

				const sinPhi	= Math.sin( phi );

				const cosPhi	= Math.cos( phi );

				const x		= sinTheta * cosPhi;

				const y		= cosTheta;

				const z		= sinTheta * sinPhi;

				positions[ pi++ ] = x * radius;

				positions[ pi++ ] = y * radius;

				positions[ pi++ ] = z * radius;

				const nx		= -x;

				const ny		= -y;

				const nz		= -z;

				normals[ ni++ ] = nx;

				normals[ ni++ ] = ny;

				normals[ ni++ ] = nz;

				uvs[ ui++ ] = lon / lonCount;

				uvs[ ui++ ] = lat / latCount;

			}

		}

		let ii = 0;

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

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

				const i0 = y * ( lonCount + 1 ) + x;

				const i1 = i0 + 1;

				const i2 = ( y + 1 ) * ( lonCount + 1 ) + x;

				const i3 = i2 + 1;

				indices[ ii++ ] = i0;

				indices[ ii++ ] = i2;

				indices[ ii++ ] = i1;

				indices[ ii++ ] = i1;

				indices[ ii++ ] = i2;

				indices[ ii++ ] = i3;

			}

		}

		return {

			positions,

			normals,

			uvs,

			indices

		};

	}


	_buildCubeGeometry( ) {

		const positions = new Float32Array( [
			// Front
			-1, -1,  1,   1, -1,  1,   1,  1,  1,  -1,  1,  1,
			// Back
			-1, -1, -1,  -1,  1, -1,   1,  1, -1,   1, -1, -1,
			// Top
			-1,  1, -1,  -1,  1,  1,   1,  1,  1,   1,  1, -1,
			// Bottom
			-1, -1, -1,   1, -1, -1,   1, -1,  1,  -1, -1,  1,
			// Right
			 1, -1, -1,   1,  1, -1,   1,  1,  1,   1, -1,  1,
			// Left
			-1, -1, -1,  -1, -1,  1,  -1,  1,  1,  -1,  1, -1
		] );

		const colors = new Float32Array( [
			// Front
			0.1, 0.6, 1.0,  0.1, 0.6, 1.0,  0.1, 0.6, 1.0,  0.1, 0.6, 1.0,
			// Back
			0.1, 0.2, 0.8,  0.1, 0.2, 0.8,  0.1, 0.2, 0.8,  0.1, 0.2, 0.8,
			// Top
			0.2, 0.8, 0.5,  0.2, 0.8, 0.5,  0.2, 0.8, 0.5,  0.2, 0.8, 0.5,
			// Bottom
			0.9, 0.7, 0.2,  0.9, 0.7, 0.2,  0.9, 0.7, 0.2,  0.9, 0.7, 0.2,
			// Right
			0.8, 0.3, 0.4,  0.8, 0.3, 0.4,  0.8, 0.3, 0.4,  0.8, 0.3, 0.4,
			// Left
			0.4, 0.9, 0.7,  0.4, 0.9, 0.7,  0.4, 0.9, 0.7,  0.4, 0.9, 0.7
		] );

		const indices = new Uint32Array( [
			0, 1, 2,   0, 2, 3,       // Front
			4, 5, 6,   4, 6, 7,       // Back
			8, 9, 10,  8, 10, 11,     // Top
			12, 13, 14, 12, 14, 15,   // Bottom
			16, 17, 18, 16, 18, 19,   // Right
			20, 21, 22, 20, 22, 23    // Left
		] );

		return {
			positions,
			colors,
			indices
		};

	}

}