GPGPU particle simulation

import * as THREE from 'three';
const SIZE = 48;
const COUNT = SIZE * SIZE;
const renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setPixelRatio(Math.min(devicePixelRatio, 2));
document.body.appendChild(renderer.domElement);

const initData = new Float32Array(COUNT * 4);
for (let i = 0; i < COUNT; i++) {
  const r = 0.4 + Math.random() * 1.3;
  const th = Math.random() * Math.PI * 2;
  const ph = Math.acos(2 * Math.random() - 1);
  initData[i * 4] = r * Math.sin(ph) * Math.cos(th);
  initData[i * 4 + 1] = r * Math.sin(ph) * Math.sin(th);
  initData[i * 4 + 2] = r * Math.cos(ph);
  initData[i * 4 + 3] = 1;
}
const initTex = new THREE.DataTexture(initData, SIZE, SIZE, THREE.RGBAFormat, THREE.FloatType);
initTex.minFilter = THREE.NearestFilter;
initTex.magFilter = THREE.NearestFilter;
initTex.needsUpdate = true;

function makeRT() {
  return new THREE.WebGLRenderTarget(SIZE, SIZE, {
    type: THREE.HalfFloatType,
    format: THREE.RGBAFormat,
    minFilter: THREE.NearestFilter,
    magFilter: THREE.NearestFilter,
    depthBuffer: false,
    stencilBuffer: false,
  });
}
let rtRead = makeRT();
let rtWrite = makeRT();

const quadGeo = new THREE.PlaneGeometry(2, 2);
const passCamera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0, 1);
passCamera.position.z = 1;
passCamera.lookAt(0, 0, 0);

const passVertex = 'varying vec2 vUv; void main() { vUv = uv; gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0); }';

const copyMat = new THREE.ShaderMaterial({
  uniforms: { tSrc: { value: initTex } },
  vertexShader: passVertex,
  fragmentShader: 'uniform sampler2D tSrc; varying vec2 vUv; void main() { gl_FragColor = texture2D(tSrc, vUv); }',
});
const copyScene = new THREE.Scene();
copyScene.add(new THREE.Mesh(quadGeo, copyMat));
renderer.setRenderTarget(rtRead);
renderer.render(copyScene, passCamera);
renderer.setRenderTarget(null);

const simFrag = [
  'uniform sampler2D tPrev;',
  'uniform float uTime;',
  'uniform float uDelta;',
  'varying vec2 vUv;',
  'float hash(vec3 p) { return fract(sin(dot(p, vec3(127.1, 311.7, 74.7))) * 43758.5453); }',
  'float vnoise(vec3 p) {',
  '  vec3 i = floor(p); vec3 f = fract(p); f = f * f * (3.0 - 2.0 * f);',
  '  float n000 = hash(i), n100 = hash(i + vec3(1.0,0.0,0.0)), n010 = hash(i + vec3(0.0,1.0,0.0)), n110 = hash(i + vec3(1.0,1.0,0.0));',
  '  float n001 = hash(i + vec3(0.0,0.0,1.0)), n101 = hash(i + vec3(1.0,0.0,1.0)), n011 = hash(i + vec3(0.0,1.0,1.0)), n111 = hash(i + vec3(1.0,1.0,1.0));',
  '  float nx00 = mix(n000, n100, f.x), nx10 = mix(n010, n110, f.x), nx01 = mix(n001, n101, f.x), nx11 = mix(n011, n111, f.x);',
  '  float nxy0 = mix(nx00, nx10, f.y), nxy1 = mix(nx01, nx11, f.y);',
  '  return mix(nxy0, nxy1, f.z);',
  '}',
  'vec3 curl(vec3 p) {',
  '  float e = 0.2;',
  '  float dx = vnoise(p + vec3(e,0.0,0.0)) - vnoise(p - vec3(e,0.0,0.0));',
  '  float dy = vnoise(p + vec3(0.0,e,0.0)) - vnoise(p - vec3(0.0,e,0.0));',
  '  float dz = vnoise(p + vec3(0.0,0.0,e)) - vnoise(p - vec3(0.0,0.0,e));',
  '  return vec3(dy - dz, dz - dx, dx - dy);',
  '}',
  'void main() {',
  '  vec3 pos = texture2D(tPrev, vUv).xyz;',
  '  vec3 vel = curl(pos * 0.55 + uTime * 0.06) * 1.6;',
  '  pos += vel * uDelta;',
  '  float d = length(pos);',
  '  if (d > 2.2) pos *= 0.985;',
  '  gl_FragColor = vec4(pos, 1.0);',
  '}',
].join('\n');
const simMat = new THREE.ShaderMaterial({
  uniforms: { tPrev: { value: rtRead.texture }, uTime: { value: 0 }, uDelta: { value: 0 } },
  vertexShader: passVertex,
  fragmentShader: simFrag,
});
const simScene = new THREE.Scene();
simScene.add(new THREE.Mesh(quadGeo, simMat));

const scene = new THREE.Scene();
scene.background = new THREE.Color(0x05060c);
const camera = new THREE.PerspectiveCamera(45, 1, 0.1, 100);
camera.position.set(0, 0, 5.4);
camera.lookAt(0, 0, 0);

const ref = new Float32Array(COUNT * 2);
const dummyPos = new Float32Array(COUNT * 3);
let p = 0;
for (let j = 0; j < SIZE; j++) {
  for (let i = 0; i < SIZE; i++) {
    ref[p * 2] = (i + 0.5) / SIZE;
    ref[p * 2 + 1] = (j + 0.5) / SIZE;
    p++;
  }
}
const pointsGeo = new THREE.BufferGeometry();
pointsGeo.setAttribute('position', new THREE.BufferAttribute(dummyPos, 3));
pointsGeo.setAttribute('ref', new THREE.BufferAttribute(ref, 2));
const pointsMat = new THREE.ShaderMaterial({
  uniforms: { tPosition: { value: rtRead.texture } },
  vertexShader: [
    'uniform sampler2D tPosition;',
    'attribute vec2 ref;',
    'varying float vDist;',
    'void main() {',
    '  vec3 pos = texture2D(tPosition, ref).xyz;',
    '  vDist = length(pos);',
    '  vec4 mv = modelViewMatrix * vec4(pos, 1.0);',
    '  gl_PointSize = 32.0 / -mv.z;',
    '  gl_Position = projectionMatrix * mv;',
    '}',
  ].join('\n'),
  fragmentShader: [
    'varying float vDist;',
    'void main() {',
    '  vec2 c = gl_PointCoord - 0.5;',
    '  if (dot(c, c) > 0.25) discard;',
    '  vec3 colA = vec3(0.35, 0.55, 0.98);',
    '  vec3 colB = vec3(0.95, 0.4, 0.8);',
    '  gl_FragColor = vec4(mix(colA, colB, clamp(vDist / 2.2, 0.0, 1.0)), 0.92);',
    '}',
  ].join('\n'),
  transparent: true,
  depthWrite: false,
});
const points = new THREE.Points(pointsGeo, pointsMat);
points.frustumCulled = false;
scene.add(points);

function resize() { renderer.setSize(innerWidth, innerHeight); camera.aspect = innerWidth / innerHeight; camera.updateProjectionMatrix(); }
addEventListener('resize', resize); resize();

let last = 0;
renderer.setAnimationLoop((t) => {
  const time = t * 0.001;
  const delta = last === 0 ? 0.016 : Math.min(0.05, time - last);
  last = time;
  simMat.uniforms.tPrev.value = rtRead.texture;
  simMat.uniforms.uTime.value = time;
  simMat.uniforms.uDelta.value = delta;
  renderer.setRenderTarget(rtWrite);
  renderer.render(simScene, passCamera);
  renderer.setRenderTarget(null);
  const tmp = rtRead; rtRead = rtWrite; rtWrite = tmp;
  pointsMat.uniforms.tPosition.value = rtRead.texture;
  points.rotation.y = time * 0.1;
  renderer.render(scene, camera);
});