webgpu_compute_water.html 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505
  1. <!DOCTYPE html>
  2. <html lang="en">
  3. <head>
  4. <title>three.js webgpu - compute - water</title>
  5. <meta charset="utf-8">
  6. <meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
  7. <link type="text/css" rel="stylesheet" href="main.css">
  8. </head>
  9. <body>
  10. <div id="info">
  11. <a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> - <span id="waterSize"></span> webgpu compute water<br/>
  12. Move mouse to disturb water.
  13. </div>
  14. <script type="importmap">
  15. {
  16. "imports": {
  17. "three": "../build/three.webgpu.js",
  18. "three/tsl": "../build/three.webgpu.js",
  19. "three/addons/": "./jsm/"
  20. }
  21. }
  22. </script>
  23. <script type="module">
  24. import * as THREE from 'three';
  25. import { color, instanceIndex, If, varyingProperty, uint, int, negate, floor, float, length, clamp, vec2, cos, vec3, vertexIndex, Fn, uniform, storageObject, min, max, positionLocal, transformNormalToView } from 'three/tsl';
  26. import { SimplexNoise } from 'three/addons/math/SimplexNoise.js';
  27. import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
  28. import Stats from 'three/addons/libs/stats.module.js';
  29. // Dimensions of simulation grid.
  30. const WIDTH = 128;
  31. // Water size in system units.
  32. const BOUNDS = 512;
  33. const BOUNDS_HALF = BOUNDS * 0.5;
  34. const waterMaxHeight = 10;
  35. let container, stats;
  36. let camera, scene, renderer;
  37. let mouseMoved = false;
  38. const mouseCoords = new THREE.Vector2();
  39. const raycaster = new THREE.Raycaster();
  40. let effectController;
  41. let waterMesh, meshRay;
  42. let computeHeight, computeSmooth, computeSphere;
  43. const NUM_SPHERES = 100;
  44. const simplex = new SimplexNoise();
  45. init();
  46. function noise( x, y ) {
  47. let multR = waterMaxHeight;
  48. let mult = 0.025;
  49. let r = 0;
  50. for ( let i = 0; i < 15; i ++ ) {
  51. r += multR * simplex.noise( x * mult, y * mult );
  52. multR *= 0.53 + 0.025 * i;
  53. mult *= 1.25;
  54. }
  55. return r;
  56. }
  57. function init() {
  58. container = document.createElement( 'div' );
  59. document.body.appendChild( container );
  60. camera = new THREE.PerspectiveCamera( 75, window.innerWidth / window.innerHeight, 1, 3000 );
  61. camera.position.set( 0, 200, 350 );
  62. camera.lookAt( 0, 0, 0 );
  63. scene = new THREE.Scene();
  64. const sun = new THREE.DirectionalLight( 0xFFFFFF, 3.0 );
  65. sun.position.set( 300, 400, 175 );
  66. scene.add( sun );
  67. const sun2 = new THREE.DirectionalLight( 0x40A040, 2.0 );
  68. sun2.position.set( - 100, 350, - 200 );
  69. scene.add( sun2 );
  70. //
  71. effectController = {
  72. mousePos: uniform( new THREE.Vector2( 10000, 10000 ) ).label( 'mousePos' ),
  73. mouseSize: uniform( 30.0 ).label( 'mouseSize' ),
  74. viscosity: uniform( 0.95 ).label( 'viscosity' ),
  75. spheresEnabled: true,
  76. wireframe: false
  77. };
  78. // Initialize height storage buffers
  79. const heightArray = new Float32Array( WIDTH * WIDTH );
  80. const prevHeightArray = new Float32Array( WIDTH * WIDTH );
  81. let p = 0;
  82. for ( let j = 0; j < WIDTH; j ++ ) {
  83. for ( let i = 0; i < WIDTH; i ++ ) {
  84. const x = i * 128 / WIDTH;
  85. const y = j * 128 / WIDTH;
  86. const height = noise( x, y );
  87. heightArray[ p ] = height;
  88. prevHeightArray[ p ] = height;
  89. p ++;
  90. }
  91. }
  92. const heightBufferAttribute = new THREE.StorageBufferAttribute( heightArray, 1 );
  93. const prevHeightBufferAttribute = new THREE.StorageBufferAttribute( prevHeightArray, 1 );
  94. const heightStorage = storageObject( heightBufferAttribute, 'float', heightBufferAttribute.count ).label( 'Height' );
  95. const prevHeightStorage = storageObject( prevHeightBufferAttribute, 'float', prevHeightBufferAttribute.count ).label( 'PrevHeight' );
  96. const heightRead = storageObject( heightBufferAttribute, 'float', heightBufferAttribute.count ).toReadOnly().label( 'HeightRead' );
  97. // Get Indices of Neighbor Values of an Index in the Simulation Grid
  98. const getNeighborIndicesTSL = ( index ) => {
  99. const width = uint( WIDTH );
  100. // Get 2-D compute coordinate from one-dimensional instanceIndex. The calculation will
  101. // still work even if you dispatch your compute shader 2-dimensionally, since within a compute
  102. // context, instanceIndex is a 1-dimensional value derived from the workgroup dimensions.
  103. // Cast to int to prevent unintended index overflow upon subtraction.
  104. const x = int( index.modInt( WIDTH ) );
  105. const y = int( index.div( WIDTH ) );
  106. // The original shader accesses height via texture uvs. However, unlike with textures, we can't
  107. // access areas that are out of bounds. Accordingly, we emulate the Clamp to Edge Wrapping
  108. // behavior of accessing a DataTexture with out of bounds uvs.
  109. const leftX = max( 0, x.sub( 1 ) );
  110. const rightX = min( x.add( 1 ), width.sub( 1 ) );
  111. const bottomY = max( 0, y.sub( 1 ) );
  112. const topY = min( y.add( 1 ), width.sub( 1 ) );
  113. const westIndex = y.mul( width ).add( leftX );
  114. const eastIndex = y.mul( width ).add( rightX );
  115. const southIndex = bottomY.mul( width ).add( x );
  116. const northIndex = topY.mul( width ).add( x );
  117. return { northIndex, southIndex, eastIndex, westIndex };
  118. };
  119. // Get simulation index neighbor values
  120. const getNeighborValuesTSL = ( index, store ) => {
  121. const { northIndex, southIndex, eastIndex, westIndex } = getNeighborIndicesTSL( index );
  122. const north = store.element( northIndex );
  123. const south = store.element( southIndex );
  124. const east = store.element( eastIndex );
  125. const west = store.element( westIndex );
  126. return { north, south, east, west };
  127. };
  128. // Get new normals of simulation area.
  129. const getNormalsFromHeightTSL = ( index, store ) => {
  130. const { north, south, east, west } = getNeighborValuesTSL( index, store );
  131. const normalX = ( west.sub( east ) ).mul( WIDTH / BOUNDS );
  132. const normalY = ( south.sub( north ) ).mul( WIDTH / BOUNDS );
  133. return { normalX, normalY };
  134. };
  135. computeHeight = Fn( () => {
  136. const { viscosity, mousePos, mouseSize } = effectController;
  137. const height = heightStorage.element( instanceIndex ).toVar();
  138. const prevHeight = prevHeightStorage.element( instanceIndex ).toVar();
  139. const { north, south, east, west } = getNeighborValuesTSL( instanceIndex, heightStorage );
  140. const neighborHeight = north.add( south ).add( east ).add( west );
  141. neighborHeight.mulAssign( 0.5 );
  142. neighborHeight.subAssign( prevHeight );
  143. const newHeight = neighborHeight.mul( viscosity );
  144. // Get 2-D compute coordinate from one-dimensional instanceIndex.
  145. const x = float( instanceIndex.modInt( WIDTH ) ).mul( 1 / WIDTH );
  146. const y = float( instanceIndex.div( WIDTH ) ).mul( 1 / WIDTH );
  147. // Mouse influence
  148. const centerVec = vec2( 0.5 );
  149. // Get length of position in range [ -BOUNDS / 2, BOUNDS / 2 ], offset by mousePos, then scale.
  150. const mousePhase = clamp( length( ( vec2( x, y ).sub( centerVec ) ).mul( BOUNDS ).sub( mousePos ) ).mul( Math.PI ).div( mouseSize ), 0.0, Math.PI );
  151. newHeight.addAssign( cos( mousePhase ).add( 1.0 ).mul( 0.28 ) );
  152. prevHeightStorage.element( instanceIndex ).assign( height );
  153. heightStorage.element( instanceIndex ).assign( newHeight );
  154. } )().compute( WIDTH * WIDTH );
  155. computeSmooth = Fn( () => {
  156. const height = heightStorage.element( instanceIndex ).toVar();
  157. const prevHeight = prevHeightStorage.element( instanceIndex ).toVar();
  158. // Get neighboring height values.
  159. const { north: northH, south: southH, east: eastH, west: westH } = getNeighborValuesTSL( instanceIndex, heightStorage );
  160. // Get neighboring prev height values.
  161. const { north: northP, south: southP, east: eastP, west: westP } = getNeighborValuesTSL( instanceIndex, prevHeightStorage );
  162. height.addAssign( northH.add( southH ).add( eastH ).add( westH ) );
  163. prevHeight.addAssign( northP.add( southP ).add( eastP ).add( westP ) );
  164. heightStorage.element( instanceIndex ).assign( height.div( 5 ) );
  165. prevHeightStorage.element( instanceIndex ).assign( height.div( 5 ) );
  166. } )().compute( WIDTH * WIDTH/*, [ 8, 8 ]*/ );
  167. // Water Geometry corresponds with buffered compute grid.
  168. const waterGeometry = new THREE.PlaneGeometry( BOUNDS, BOUNDS, WIDTH - 1, WIDTH - 1 );
  169. // material: make a THREE.ShaderMaterial clone of THREE.MeshPhongMaterial, with customized position shader.
  170. const waterMaterial = new THREE.MeshPhongNodeMaterial();
  171. waterMaterial.lights = true;
  172. waterMaterial.colorNode = color( 0x0040C0 );
  173. waterMaterial.specularNode = color( 0x111111 );
  174. waterMaterial.shininess = Math.max( 50, 1e-4 );
  175. waterMaterial.positionNode = Fn( () => {
  176. // To correct the lighting as our mesh undulates, we have to reassign the normals in the position shader.
  177. const { normalX, normalY } = getNormalsFromHeightTSL( vertexIndex, heightRead );
  178. varyingProperty( 'vec3', 'v_normalView' ).assign( transformNormalToView( vec3( normalX, negate( normalY ), 1.0 ) ) );
  179. return vec3( positionLocal.x, positionLocal.y, heightRead.element( vertexIndex ) );
  180. } )();
  181. waterMesh = new THREE.Mesh( waterGeometry, waterMaterial );
  182. waterMesh.rotation.x = - Math.PI / 2;
  183. waterMesh.matrixAutoUpdate = false;
  184. waterMesh.updateMatrix();
  185. scene.add( waterMesh );
  186. // THREE.Mesh just for mouse raycasting
  187. const geometryRay = new THREE.PlaneGeometry( BOUNDS, BOUNDS, 1, 1 );
  188. meshRay = new THREE.Mesh( geometryRay, new THREE.MeshBasicMaterial( { color: 0xFFFFFF, visible: false } ) );
  189. meshRay.rotation.x = - Math.PI / 2;
  190. meshRay.matrixAutoUpdate = false;
  191. meshRay.updateMatrix();
  192. scene.add( meshRay );
  193. // Create sphere THREE.InstancedMesh
  194. const sphereGeometry = new THREE.SphereGeometry( 4, 24, 12 );
  195. const sphereMaterial = new THREE.MeshPhongMaterial( { color: 0xFFFF00 } );
  196. // Initialize sphere mesh instance position and velocity.
  197. const spherePositionArray = new Float32Array( NUM_SPHERES * 3 );
  198. // Only hold velocity in x and z directions.
  199. // The sphere is wedded to the surface of the water, and will only move vertically with the water.
  200. const sphereVelocityArray = new Float32Array( NUM_SPHERES * 2 );
  201. for ( let i = 0; i < NUM_SPHERES; i ++ ) {
  202. spherePositionArray[ i * 3 + 0 ] = ( Math.random() - 0.5 ) * BOUNDS * 0.7;
  203. spherePositionArray[ i * 3 + 1 ] = 0;
  204. spherePositionArray[ i * 3 + 2 ] = ( Math.random() - 0.5 ) * BOUNDS * 0.7;
  205. }
  206. sphereVelocityArray.fill( 0.0 );
  207. // Sphere Instance Storage
  208. const sphereInstancePositionAttribute = new THREE.StorageInstancedBufferAttribute( spherePositionArray, 3 );
  209. const sphereInstancePositionStorage = storageObject( sphereInstancePositionAttribute, 'vec3', sphereInstancePositionAttribute.count ).label( 'SpherePosition' );
  210. const sphereInstancePositionRead = storageObject( sphereInstancePositionAttribute, 'vec3', sphereInstancePositionAttribute.count ).toReadOnly();
  211. const sphereVelocityAttribute = new THREE.StorageInstancedBufferAttribute( sphereVelocityArray, 2 );
  212. const sphereVelocityStorage = storageObject( sphereVelocityAttribute, 'vec2', sphereVelocityAttribute.count ).label( 'SphereVelocity' );
  213. computeSphere = Fn( () => {
  214. const instancePosition = sphereInstancePositionStorage.element( instanceIndex );
  215. const velocity = sphereVelocityStorage.element( instanceIndex );
  216. // Bring position from range of [ -BOUNDS/2, BOUNDS/2 ] to [ 0, BOUNDS ]
  217. const tempX = instancePosition.x.add( BOUNDS_HALF );
  218. const tempZ = instancePosition.z.add( BOUNDS_HALF );
  219. // Bring position from range [ 0, BOUNDS ] to [ 0, WIDTH ]
  220. // ( i.e bring geometry range into 'heightmap' range )
  221. // WIDTH = 128, BOUNDS = 512... same as dividing by 4
  222. tempX.mulAssign( WIDTH / BOUNDS );
  223. tempZ.mulAssign( WIDTH / BOUNDS );
  224. // Can only access storage buffers with uints
  225. const xCoord = uint( floor( tempX ) );
  226. const zCoord = uint( floor( tempZ ) );
  227. // Get one dimensional index
  228. const heightInstanceIndex = zCoord.mul( WIDTH ).add( xCoord );
  229. // Set to read-only to be safe, even if it's not strictly necessary for compute access.
  230. const height = heightRead.element( heightInstanceIndex );
  231. // Assign height to sphere position
  232. instancePosition.y.assign( height );
  233. // Calculate normal of the water mesh at this location.
  234. const { normalX, normalY } = getNormalsFromHeightTSL( heightInstanceIndex, heightRead );
  235. normalX.mulAssign( 0.1 );
  236. normalY.mulAssign( 0.1 );
  237. const waterNormal = vec3( normalX, 0.0, negate( normalY ) );
  238. const newVelocity = vec3( velocity.x, 0.0, velocity.y ).add( waterNormal );
  239. newVelocity.mulAssign( 0.998 );
  240. const newPosition = instancePosition.add( newVelocity ).toVar();
  241. // Reverse velocity and reset position when exceeding bounds.
  242. If( newPosition.x.lessThan( - BOUNDS_HALF ), () => {
  243. newPosition.x = float( - BOUNDS_HALF ).add( 0.001 );
  244. newVelocity.x.mulAssign( - 0.3 );
  245. } ).ElseIf( newPosition.x.greaterThan( BOUNDS_HALF ), () => {
  246. newPosition.x = float( BOUNDS_HALF ).sub( 0.001 );
  247. newVelocity.x.mulAssign( - 0.3 );
  248. } );
  249. If( newPosition.z.lessThan( - BOUNDS_HALF ), () => {
  250. newPosition.z = float( - BOUNDS_HALF ).add( 0.001 );
  251. newVelocity.z.mulAssign( - 0.3 );
  252. } ).ElseIf( newPosition.z.greaterThan( BOUNDS_HALF ), () => {
  253. newPosition.z = float( BOUNDS_HALF ).sub( 0.001 );
  254. newVelocity.z.mulAssign( - 0.3 );
  255. } );
  256. instancePosition.assign( newPosition );
  257. velocity.assign( vec2( newVelocity.x, newVelocity.z ) );
  258. } )().compute( NUM_SPHERES );
  259. sphereMaterial.positionNode = Fn( () => {
  260. const instancePosition = sphereInstancePositionRead.element( instanceIndex );
  261. const newPosition = positionLocal.add( instancePosition );
  262. return newPosition;
  263. } )();
  264. const sphereMesh = new THREE.InstancedMesh( sphereGeometry, sphereMaterial, NUM_SPHERES );
  265. scene.add( sphereMesh );
  266. renderer = new THREE.WebGPURenderer( { antialias: true } );
  267. renderer.setPixelRatio( window.devicePixelRatio );
  268. renderer.setSize( window.innerWidth, window.innerHeight );
  269. renderer.setAnimationLoop( animate );
  270. container.appendChild( renderer.domElement );
  271. stats = new Stats();
  272. container.appendChild( stats.dom );
  273. container.style.touchAction = 'none';
  274. container.addEventListener( 'pointermove', onPointerMove );
  275. window.addEventListener( 'resize', onWindowResize );
  276. const gui = new GUI();
  277. gui.add( effectController.mouseSize, 'value', 1.0, 100.0, 1.0 ).name( 'Mouse Size' );
  278. gui.add( effectController.viscosity, 'value', 0.9, 0.999, 0.001 ).name( 'viscosity' );
  279. const buttonCompute = {
  280. smoothWater: function () {
  281. renderer.compute( computeSmooth );
  282. }
  283. };
  284. gui.add( buttonCompute, 'smoothWater' );
  285. gui.add( effectController, 'spheresEnabled' ).onChange( () => {
  286. sphereMesh.visible = effectController.spheresEnabled;
  287. } );
  288. gui.add( effectController, 'wireframe' ).onChange( () => {
  289. waterMesh.material.wireframe = ! waterMesh.material.wireframe;
  290. waterMesh.material.needsUpdate = true;
  291. } );
  292. }
  293. function onWindowResize() {
  294. camera.aspect = window.innerWidth / window.innerHeight;
  295. camera.updateProjectionMatrix();
  296. renderer.setSize( window.innerWidth, window.innerHeight );
  297. }
  298. function setMouseCoords( x, y ) {
  299. mouseCoords.set( ( x / renderer.domElement.clientWidth ) * 2 - 1, - ( y / renderer.domElement.clientHeight ) * 2 + 1 );
  300. mouseMoved = true;
  301. }
  302. function onPointerMove( event ) {
  303. if ( event.isPrimary === false ) return;
  304. setMouseCoords( event.clientX, event.clientY );
  305. }
  306. function animate() {
  307. render();
  308. stats.update();
  309. }
  310. function render() {
  311. if ( mouseMoved ) {
  312. raycaster.setFromCamera( mouseCoords, camera );
  313. const intersects = raycaster.intersectObject( meshRay );
  314. if ( intersects.length > 0 ) {
  315. const point = intersects[ 0 ].point;
  316. effectController.mousePos.value.set( point.x, point.z );
  317. } else {
  318. effectController.mousePos.value.set( 10000, 10000 );
  319. }
  320. mouseMoved = false;
  321. } else {
  322. effectController.mousePos.value.set( 10000, 10000 );
  323. }
  324. renderer.compute( computeHeight );
  325. if ( effectController.spheresEnabled ) {
  326. renderer.compute( computeSphere );
  327. }
  328. renderer.render( scene, camera );
  329. }
  330. </script>
  331. </body>
  332. </html>