feat: ThreeJS 'Compute Snow Particles' example in tgpu

apps/typegpu-docs/src/examples/threejs/compute-particles-snow/entities.ts

@@ -0,0 +1,113 @@
+import * as THREE from 'three/webgpu';
+import * as TSL from 'three/tsl';
+import { TeapotGeometry } from 'three/addons/geometries/TeapotGeometry.js';
+
+import * as t3 from '@typegpu/three';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+
+export const dirLight = (() => {
+  const dirLight = new THREE.DirectionalLight(0xf9ff9b, 9);
+  dirLight.castShadow = true;
+  dirLight.position.set(10, 10, 0);
+  dirLight.castShadow = true;
+  Object.assign(dirLight.shadow.camera, {
+    near: 1,
+    far: 30,
+    right: 30,
+    left: -30,
+    top: 30,
+    bottom: -30,
+  });
+  dirLight.shadow.mapSize.width = 2048;
+  dirLight.shadow.mapSize.height = 2048;
+  dirLight.shadow.bias = -0.009;
+  return dirLight;
+})();
+
+export const fog = new THREE.Fog(0x0f3c37, 4, 40);
+
+export const hemisphereLight = new THREE.HemisphereLight(
+  0x0f3c37,
+  0x080d10,
+  100,
+);
+
+export const floor = (() => {
+  const floorGeometry = new THREE.PlaneGeometry(100, 100);
+  floorGeometry.rotateX(-Math.PI / 2);
+  const floor = new THREE.Mesh(
+    floorGeometry,
+    new THREE.MeshStandardMaterial({
+      color: 0x800000,
+      roughness: 0.5,
+      metalness: 0,
+      transparent: false,
+    }),
+  );
+  floor.position.y = 0;
+  floor.material.opacityNode = t3.toTSL(() => {
+    'use gpu';
+    return std.saturate(
+      std.length(t3.fromTSL(TSL.positionLocal.xz, d.vec2f).$.mul(0.05)),
+    ) - 1;
+  });
+  floor.layers.disableAll();
+  floor.layers.enable(1);
+  floor.layers.enable(2);
+  return floor;
+})();
+
+export const xmasTree = (() => {
+  const count = 8;
+
+  const coneMaterial = new THREE.MeshStandardNodeMaterial({
+    color: 0x004225,
+    roughness: 0.6,
+    metalness: 0,
+  });
+
+  const object = new THREE.Group();
+
+  for (let i = 0; i < count; i++) {
+    const radius = 1 + i;
+
+    const coneGeometry = new THREE.ConeGeometry(
+      radius * 0.95,
+      radius * 1.25,
+      32,
+    );
+
+    const cone = new THREE.Mesh(coneGeometry, coneMaterial);
+    cone.castShadow = true;
+    cone.position.y = ((count - i) * 1.5) + (count * 0.6);
+    object.add(cone);
+  }
+
+  const cylinderGeometry = new THREE.CylinderGeometry(1, 1, count, 32);
+  const cylinder = new THREE.Mesh(cylinderGeometry, coneMaterial);
+  cylinder.position.y = count / 2;
+  object.add(cylinder);
+
+  object.layers.disableAll();
+  object.layers.enable(1);
+  object.layers.enable(2);
+
+  return object;
+})();
+
+export const teapot = (() => {
+  const teapot = new THREE.Mesh(
+    new TeapotGeometry(0.5, 18),
+    new THREE.MeshBasicNodeMaterial({
+      color: 0xfcfb9e,
+    }),
+  );
+
+  teapot.position.y = 18;
+  teapot.layers.disableAll();
+  teapot.layers.enable(1);
+  teapot.layers.enable(2);
+
+  return teapot;
+})();

apps/typegpu-docs/src/examples/threejs/compute-particles-snow/index.html

@@ -0,0 +1 @@
+<canvas data-fit-to-container></canvas>

apps/typegpu-docs/src/examples/threejs/compute-particles-snow/index.ts

@@ -0,0 +1,244 @@
+/*
+ * Based on: https://github.com/mrdoob/three.js/blob/master/examples/webgpu_compute_particles_snow.html
+ */
+import * as THREE from 'three/webgpu';
+import * as TSL from 'three/tsl';
+import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
+import { gaussianBlur } from 'three/addons/tsl/display/GaussianBlurNode.js';
+
+import * as t3 from '@typegpu/three';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+import { randf } from '@typegpu/noise';
+
+import * as e from './entities.ts';
+
+const maxParticleCount = 100000;
+
+const canvas = document.querySelector('canvas') as HTMLCanvasElement;
+const renderer = new THREE.WebGPURenderer({ canvas, antialias: true });
+renderer.setClearColor(0x000000);
+renderer.setPixelRatio(window.devicePixelRatio);
+renderer.setSize(canvas.clientWidth, canvas.clientHeight, false);
+renderer.toneMapping = THREE.ACESFilmicToneMapping;
+
+const camera = new THREE.PerspectiveCamera(
+  60,
+  canvas.clientWidth / canvas.clientHeight,
+  0.1,
+  1000,
+);
+camera.position.set(20, 2, 20);
+camera.lookAt(0, 40, 0);
+camera.layers.enable(2);
+
+const collisionCamera = new THREE.OrthographicCamera(
+  -50,
+  50,
+  50,
+  -50,
+  0.1,
+  1000,
+);
+collisionCamera.position.y = 50;
+collisionCamera.lookAt(0, 0, 0);
+collisionCamera.layers.enable(1);
+
+const collisionPosRT = new THREE.RenderTarget(2048, 2048, {
+  type: THREE.FloatType,
+  minFilter: THREE.NearestFilter,
+  magFilter: THREE.NearestFilter,
+  colorSpace: THREE.NoColorSpace,
+  generateMipmaps: false,
+});
+
+const collisionPosMaterial = new THREE.MeshBasicNodeMaterial();
+collisionPosMaterial.fog = false;
+collisionPosMaterial.toneMapped = false;
+collisionPosMaterial.outputNode = TSL.vec4(
+  TSL.positionWorld.y,
+  TSL.positionWorld.y,
+  TSL.positionWorld.y,
+  1,
+);
+
+const positionBuffer = t3.instancedArray(
+  maxParticleCount,
+  d.vec3f,
+);
+const scaleBuffer = t3.instancedArray(maxParticleCount, d.vec3f);
+const staticPositionBuffer = t3.instancedArray(maxParticleCount, d.vec3f);
+const dataBuffer = t3.instancedArray(maxParticleCount, d.vec4f);
+
+const computeInit = t3.toTSL(() => {
+  'use gpu';
+  const instanceIdx = t3.instanceIndex.$;
+
+  randf.seed(instanceIdx / maxParticleCount);
+  const rand = d.vec3f(randf.sample(), randf.sample(), randf.sample());
+  const randPos = rand.mul(d.vec3f(100, 500, 100))
+    .add(d.vec3f(-50, 3, -50));
+
+  positionBuffer.$[instanceIdx] = d.vec3f(randPos);
+  scaleBuffer.$[instanceIdx] = d.vec3f(randf.sample() * 0.8 + 0.2);
+  staticPositionBuffer.$[instanceIdx] = d.vec3f(1000, 10000, 1000);
+
+  dataBuffer.$[instanceIdx] = d.vec4f(
+    randPos.x,
+    -0.01 * rand.y - 0.02,
+    randPos.z,
+    rand.x,
+  );
+}).compute(maxParticleCount).setName('Init Partciles');
+
+const surfaceOffset = 0.2;
+const speed = 0.4;
+const terrainTextureAccess = t3.fromTSL(
+  TSL.texture(
+    collisionPosRT.texture,
+    positionBuffer.node.element(TSL.instanceIndex).xz.add(50).div(100),
+  ),
+  d.vec4f,
+);
+
+const computeUpdate = t3.toTSL(() => {
+  'use gpu';
+  const instanceIdx = t3.instanceIndex.$;
+  const velocity = dataBuffer.$[instanceIdx].y;
+  const random = dataBuffer.$[instanceIdx].w;
+
+  const terrainHeight = terrainTextureAccess.$.y +
+    scaleBuffer.$[instanceIdx].x * surfaceOffset;
+
+  if (positionBuffer.$[instanceIdx].y > terrainHeight) {
+    positionBuffer.$[instanceIdx].x = dataBuffer.$[instanceIdx].x +
+      3 * std.sin(t3.time.$ * (random * random) * speed);
+    positionBuffer.$[instanceIdx].z = dataBuffer.$[instanceIdx].z +
+      10 * random * std.cos(t3.time.$ * speed);
+    positionBuffer.$[instanceIdx].y += velocity;
+  } else {
+    staticPositionBuffer.$[instanceIdx] = positionBuffer.$[instanceIdx];
+  }
+}).compute(maxParticleCount).setName('Update Particles');
+
+const sphereGeometry = new THREE.SphereGeometry(surfaceOffset, 5, 5);
+function particles(isStatic: boolean = false) {
+  const posBuffer = isStatic ? staticPositionBuffer : positionBuffer;
+  const layer = isStatic ? 1 : 2;
+
+  const material = new THREE.MeshStandardNodeMaterial({
+    color: 0xeeeeee,
+    roughness: 0.9,
+    metalness: 0,
+  });
+
+  material.positionNode = t3.toTSL(() => {
+    'use gpu';
+    return t3.fromTSL(TSL.positionLocal, d.vec3f).$
+      .mul(scaleBuffer.$[t3.instanceIndex.$])
+      .add(posBuffer.$[t3.instanceIndex.$]);
+  });
+
+  const rainParticles = new THREE.Mesh(sphereGeometry, material);
+  rainParticles.count = maxParticleCount;
+  rainParticles.castShadow = true;
+  rainParticles.layers.disableAll();
+  rainParticles.layers.enable(layer);
+
+  return rainParticles;
+}
+
+const scene = new THREE.Scene();
+scene.fog = e.fog;
+scene.add(e.dirLight);
+scene.add(e.hemisphereLight);
+
+const dynamicParticles = particles();
+const staticParticles = particles(true);
+scene.add(dynamicParticles);
+scene.add(staticParticles);
+
+scene.add(e.floor);
+scene.add(e.xmasTree);
+scene.add(e.teapot);
+
+scene.backgroundNode = t3.toTSL(() => {
+  'use gpu';
+  const ratio = std.saturate(
+    std.distance(t3.fromTSL(TSL.screenUV, d.vec2f).$, d.vec2f(0.5)) / 0.5,
+  );
+  // 0.25 multiplier is empirical
+  return std.mix(
+    d.vec4f(d.vec3f(0.059, 0.255, 0.251).mul(0.25), 1),
+    d.vec4f(d.vec3f(0.024, 0.039, 0.051).mul(0.25), 1),
+    ratio,
+  );
+});
+
+const controls = new OrbitControls(camera, canvas);
+controls.target.set(0, 10, 0);
+controls.minDistance = 20;
+controls.maxDistance = 35;
+controls.maxPolarAngle = Math.PI / 1.7;
+controls.autoRotate = true;
+controls.autoRotateSpeed = -0.7;
+controls.update();
+
+const scenePass = TSL.pass(scene, camera);
+const scenePassColor = scenePass.getTextureNode();
+const vignette = t3.toTSL(() => {
+  'use gpu';
+  return 1 - std.saturate(
+    std.distance(t3.fromTSL(TSL.screenUV, d.vec2f).$, d.vec2f(0.5)) * 1.35,
+  );
+});
+
+const teapotPass = TSL.pass(e.teapot, camera).getTextureNode();
+const teapotPassBlurred = gaussianBlur(teapotPass, TSL.vec2(1), 6);
+teapotPassBlurred.resolutionScale = 0.2;
+
+const scenePassColorBlurred = gaussianBlur(scenePassColor);
+scenePassColorBlurred.resolutionScale = 0.5;
+scenePassColorBlurred.directionNode = TSL.vec2(1);
+
+const totalPass = scenePass
+  .add(scenePassColorBlurred.mul(0.1))
+  .mul(vignette)
+  .add(teapotPass.mul(10).add(teapotPassBlurred));
+
+const postProcessing = new THREE.PostProcessing(renderer);
+postProcessing.outputNode = totalPass;
+
+const resizeObserver = new ResizeObserver(() => {
+  camera.aspect = canvas.clientWidth / canvas.clientHeight;
+  camera.updateProjectionMatrix();
+  renderer.setSize(canvas.clientWidth, canvas.clientHeight, false);
+});
+resizeObserver.observe(canvas);
+
+function animate() {
+  controls.update();
+
+  scene.overrideMaterial = collisionPosMaterial;
+  renderer.setRenderTarget(collisionPosRT);
+  renderer.render(scene, collisionCamera);
+  renderer.compute(computeUpdate);
+
+  scene.overrideMaterial = null;
+  renderer.setRenderTarget(null);
+
+  postProcessing.render();
+}
+
+await renderer.init();
+renderer.compute(computeInit);
+renderer.setAnimationLoop(animate);
+
+// #region Example controls and cleanup
+
+export function onCleanup() {
+  renderer.dispose();
+  resizeObserver.unobserve(canvas);
+}
+
+// #endregion

apps/typegpu-docs/src/examples/threejs/compute-particles-snow/meta.json

@@ -0,0 +1,5 @@
+{
+  "title": "Three.js - compute / particles / snow",
+  "category": "threejs",
+  "tags": ["experimental"]
+}

packages/typegpu-three/src/typegpu-node.ts

@@ -194,8 +194,11 @@ export class TSLAccessor<T extends d.AnyWgslData, TNode extends THREE.Node> {
     this.node = node;
     this.#dataType = dataType;
 
-    // @ts-expect-error: The properties exist on the node
-    if (!node.isStorageBufferNode && !node.isUniformNode) {
+    // node.isTextureNode - temporary workaround for textures
+    if (
+      // @ts-expect-error: The properties exist on the node
+      (!node.isStorageBufferNode && !node.isUniformNode) || node.isTextureNode
+    ) {
       this.var = tgpu.privateVar(dataType);
     }
   }