Симуляція рідини на GPU за допомогою WebGL

Рівняння імпульсу для нестисливої рідини Нав'є-Стокса має вигляд:

∂u/∂t = −(u·∇)u − (1/ρ)∇p + ν∇²u + f

де u — поле швидкостей, p — тиск, ν — в'язкість, f — зовнішні сили. З умовою нестисливості ∇·u = 0.

Кожного кадру GPU виконує ці проходи по порядку:

Шейдер не може одночасно читати й записувати ту саму текстуру. Ping-pong використовує два FBO — з одного читають, поки в інший пишуть, а потім вони міняються місцями:

import * as THREE from 'https://cdn.jsdelivr.net/npm/three@0.168/build/three.module.js';

function createFBO(width, height) {
  return new THREE.WebGLRenderTarget(width, height, {
    type: THREE.FloatType,           // 32-бітний float на канал
    format: THREE.RGBAFormat,
    minFilter: THREE.LinearFilter,
    magFilter: THREE.LinearFilter,
    wrapS: THREE.ClampToEdgeWrapping,
    wrapT: THREE.ClampToEdgeWrapping,
  });
}

class DoubleFBO {
  constructor(w, h) {
    this.read  = createFBO(w, h);
    this.write = createFBO(w, h);
  }
  swap() { [this.read, this.write] = [this.write, this.read]; }
}

const SIM_RES = 256;
const velocity = new DoubleFBO(SIM_RES, SIM_RES);
const pressure = new DoubleFBO(SIM_RES, SIM_RES);
const dye      = new DoubleFBO(SIM_RES, SIM_RES);

Напівлагранжева адвекція: «звідки прийшла ця частинка один крок часу тому?» — рухаємося полем швидкостей назад і семплимо там:

// фрагментний шейдер адвекції
const advectFrag = /* glsl */`
  uniform sampler2D uVelocity; // поле швидкостей, ЯКИМ адвектуємо
  uniform sampler2D uSource;   // поле, яке адвектуємо (швидкість або барвник)
  uniform float uDt;
  uniform vec2  uTexelSize;    // 1/SIM_RES для кожної осі

  varying vec2 vUv;

  void main() {
    // Читаємо швидкість у поточній комірці
    vec2 velocity = texture2D(uVelocity, vUv).xy;

    // Трасуємо назад у часі: prev_pos = current - velocity * dt
    vec2 prevUv = vUv - velocity * uDt * uTexelSize;

    // Семплимо вихідне поле в попередній позиції
    gl_FragColor = texture2D(uSource, prevUv);
  }
`;

Щоб виконати прохід шейдера: рендеримо повноекранний квад із цим шейдером, записуючи в цільовий FBO:

function runPass(material, targetFBO) {
  quad.material = material;
  renderer.setRenderTarget(targetFBO);
  renderer.render(quadScene, quadCamera);
  renderer.setRenderTarget(null);
}

advectMat.uniforms.uVelocity.value = velocity.read.texture;
advectMat.uniforms.uSource.value   = velocity.read.texture;
advectMat.uniforms.uDt.value       = dt;
runPass(advectMat, velocity.write);
velocity.swap();

Обчислюємо дивергенцію (наскільки розширюється поле швидкостей), а потім ітеративно розв'язуємо рівняння для тиску методом Якобі:

// Шейдер дивергенції — читає швидкість, записує скаляр div у канал .r
const divergenceFrag = /* glsl */`
  uniform sampler2D uVelocity;
  uniform vec2 uTexelSize;
  varying vec2 vUv;
  void main() {
    vec2 ts = uTexelSize;
    float L = texture2D(uVelocity, vUv - vec2(ts.x, 0)).x;
    float R = texture2D(uVelocity, vUv + vec2(ts.x, 0)).x;
    float B = texture2D(uVelocity, vUv - vec2(0, ts.y)).y;
    float T = texture2D(uVelocity, vUv + vec2(0, ts.y)).y;
    float div = 0.5 * (R - L + T - B);
    gl_FragColor = vec4(div, 0, 0, 1);
  }
`;

// Шейдер ітерації Якобі для тиску — виконується 30–50 разів
const pressureFrag = /* glsl */`
  uniform sampler2D uPressure;
  uniform sampler2D uDivergence;
  uniform vec2 uTexelSize;
  varying vec2 vUv;
  void main() {
    vec2 ts = uTexelSize;
    float L = texture2D(uPressure, vUv - vec2(ts.x, 0)).r;
    float R = texture2D(uPressure, vUv + vec2(ts.x, 0)).r;
    float B = texture2D(uPressure, vUv - vec2(0, ts.y)).r;
    float T = texture2D(uPressure, vUv + vec2(0, ts.y)).r;
    float div = texture2D(uDivergence, vUv).r;
    // Jacobi: p_new = (L+R+B+T - div) / 4
    float pNew = (L + R + B + T - div) * 0.25;
    gl_FragColor = vec4(pNew, 0, 0, 1);
  }
`;

// Виконуємо 40 ітерацій Якобі
for (let i = 0; i < 40; i++) {
  pressureMat.uniforms.uPressure.value = pressure.read.texture;
  runPass(pressureMat, pressure.write);
  pressure.swap();
}

Віднімаємо градієнт тиску від швидкості, щоб зробити поле бездивергентним:

const projectFrag = /* glsl */`
  uniform sampler2D uVelocity;
  uniform sampler2D uPressure;
  uniform vec2 uTexelSize;
  varying vec2 vUv;
  void main() {
    vec2 ts = uTexelSize;
    float pL = texture2D(uPressure, vUv - vec2(ts.x, 0)).r;
    float pR = texture2D(uPressure, vUv + vec2(ts.x, 0)).r;
    float pB = texture2D(uPressure, vUv - vec2(0, ts.y)).r;
    float pT = texture2D(uPressure, vUv + vec2(0, ts.y)).r;
    vec2 gradP = vec2(pR - pL, pT - pB) * 0.5;
    vec2 vel = texture2D(uVelocity, vUv).xy;
    gl_FragColor = vec4(vel - gradP, 0, 1);
  }
`;

«Бризкаємо» силами та барвником на симуляцію, коли користувач рухає мишею:

const splatFrag = /* glsl */`
  uniform sampler2D uTarget;
  uniform vec2  uPoint;    // нормалізована позиція миші [0,1]
  uniform vec3  uColor;    // вектор сили (xy) або колір барвника (xyz)
  uniform float uRadius;
  varying vec2 vUv;
  void main() {
    float d = distance(vUv, uPoint);
    float splat = exp(-d * d / uRadius);
    vec4 current = texture2D(uTarget, vUv);
    gl_FragColor = current + vec4(uColor * splat, 0);
  }
`;

// Викликаємо на mousemove, щоб додати швидкість + барвник
function onMouseMove(e) {
  const x = e.clientX / innerWidth;
  const y = 1.0 - e.clientY / innerHeight; // інвертуємо y
  const dx = (e.clientX - prevMouseX) / innerWidth  * 15.0; // масштабуємо до швидкості симуляції
  const dy = (e.clientY - prevMouseY) / innerHeight * 15.0;

  splatMat.uniforms.uPoint.value.set(x, y);
  splatMat.uniforms.uRadius.value = 0.0015;

  // Сплеск швидкості
  splatMat.uniforms.uColor.value.set(dx, -dy, 0);
  splatMat.uniforms.uTarget.value = velocity.read.texture;
  runPass(splatMat, velocity.write); velocity.swap();

  // Сплеск барвника
  splatMat.uniforms.uColor.value.set(0.3, 0.7, 1.0); // блакитний барвник
  splatMat.uniforms.uTarget.value = dye.read.texture;
  runPass(splatMat, dye.write); dye.swap();
}

Адвектуємо текстуру барвника бездивергентним полем швидкостей, а потім відображаємо її:

// Кожного кадру:
// 1. Адвекція швидкості
advectMat.uniforms.uSource.value = velocity.read.texture;
runPass(advectMat, velocity.write); velocity.swap();

// 2. Додавання сил (якщо натиснута миша)

// 3. Обчислення дивергенції
runPass(divergenceMat, divFBO);

// 4. Розв'язання тиску (40 ітерацій Якобі)
for (let k = 0; k < 40; k++) { /* ... */ }

// 5. Проєкція
projectMat.uniforms.uVelocity.value = velocity.read.texture;
runPass(projectMat, velocity.write); velocity.swap();

// 6. Адвекція барвника
advectMat.uniforms.uVelocity.value = velocity.read.texture;
advectMat.uniforms.uSource.value   = dye.read.texture;
runPass(advectMat, dye.write); dye.swap();

// 7. Відображення барвника на екран
displayMat.uniforms.uTexture.value = dye.read.texture;
renderer.setRenderTarget(null);
renderer.render(quadScene, quadCamera);