webgpu-app/flat.js

import { getPath } from "./common.js";

function test(position)
{
  const x = position[0] * 4 * Math.PI;
  const y = position[1] * 4 * Math.PI;
  return Math.cos(y - 3 + Math.sin(x)) * 0.5 + 0.5;
}

function length(vector)
{
  return Math.sqrt(vector[0] * vector[0] + vector[1] * vector[1]);
}

function sub(a, b)
{
  return [
    a[0] - b[0],
    a[1] - b[1],
  ];
}

function smoothing(radius, distance)
{
  const volume = Math.PI * Math.pow(radius, 8) / 4;
  const value = Math.max(0, radius * radius - distance * distance);
  return value * value * value / volume;
}

function smoothingSlope(radius, distance)
{
  if (distance >= radius)
    return 0;
  const value = radius * radius - distance * distance;
  const scale = -24.0 / (Math.PI * Math.pow(radius, 8));
  return scale * distance * value * value;
}

function densityToPressure(density, targetDensity, pressure)
{
  return (density - targetDensity) * pressure;
}

class FlatRenderer {
  positionProperty(position, circleRadius, particleCount)
  {
    var property = 0.0;
    for (var i = 0; i < particleCount; i++) {
      const particlePosition = [this.particles[i * this.particleStride + 0],
                                this.particles[i * this.particleStride + 1]];
      const particleProperty = this.particles[i * this.particleStride + 4];
      const particleDensity = this.particles[i * this.particleStride + 5];

      const distance = length(sub(position, particlePosition));
      const influence = smoothing(circleRadius, distance);

      property += particleProperty * influence / particleDensity;
    }
    return property;
  }

  positionGradient(position, circleRadius, particleCount)
  {
    var property = new Float32Array([0, 0]);
    for (var i = 0; i < particleCount; i++) {
      const particlePosition = [this.particles[i * this.particleStride + 0],
                                this.particles[i * this.particleStride + 1]];
      const signedDistance = sub(position, particlePosition);
      const distance = length(signedDistance);

      const direction = [signedDistance[0] / distance, signedDistance[1] / distance];
      const particleProperty = this.particles[i * this.particleStride + 4];
      const particleDensity = this.particles[i * this.particleStride + 5];

      const slope = smoothingSlope(circleRadius, distance);

      const weight = particleProperty * slope / particleDensity;
      property[0] += direction[0] * weight;
      property[1] += direction[1] * weight;
    }
    return property;
  }

  positionPressureForce(index, circleRadius, particleCount, targetDensity, pressure, mass)
  {
    var pressureForce = new Float32Array([0, 0]);

    const position = [this.particles[index * this.particleStride + 0],
                      this.particles[index * this.particleStride + 1]];

    for (var i = 0; i < particleCount; i++) {
      if (i === index) continue;

      const particlePosition = [this.particles[i * this.particleStride + 0],
                                this.particles[i * this.particleStride + 1]];
      const signedDistance = sub(position, particlePosition);
      const distance = length(signedDistance);

      var direction;
      if (distance === 0) {
        const x = Math.random();
        const y = Math.random();
        const l = length([x, y]);
        direction = [x / l, y / l];
      } else {
        direction = [signedDistance[0] / distance, signedDistance[1] / distance];
      }
      const particleDensity = this.particles[i * this.particleStride + 5];

      const slope = smoothingSlope(circleRadius, distance);

      const sharedPressure = (
        densityToPressure(this.particles[index * this.particleStride * 5], targetDensity, pressure) +
        densityToPressure(particleDensity, targetDensity, pressure)
      ) / 2;

      const weight = sharedPressure * slope * mass / particleDensity;

      pressureForce[0] += direction[0] * weight;
      pressureForce[1] += direction[1] * weight;
    }
    return pressureForce;
  }

  particleDensity(aPosition, circleRadius, particleCount)
  {
    var density = 0.0;
    for (var i = 0; i < particleCount; i++) {
      const bPosition = [this.particles[i * this.particleStride + 0],
                         this.particles[i * this.particleStride + 1]];
      const distance = length(sub(aPosition, bPosition));
      const influence = smoothing(circleRadius, distance);
      density += influence;
    }
    return density;
  }

  createParticles(device)
  {
    this.maxDim = 32;
    this.maxParticles = this.maxDim * this.maxDim;
    this.particleStride = 4 * 3; // in elements
    this.particles = new Float32Array(this.maxParticles * this.particleStride);

    this.maxGradientsDim = 32;
    this.maxGradients = this.maxGradientsDim * this.maxGradientsDim;
    this.gradientStride = 4 * 2;
    this.gradients = new Float32Array(this.maxGradients * this.gradientStride);

    this.particleBuffers = [];
    this.gradientBuffers = [];
    this.particleBindGroups = [];

    this.particleConfiguration = new Float32Array(4 * 4);
    this.particleConfigurationBuffers = [];

    for (let i = 0; i < 2; i++) {
      const particleBuffer = device.createBuffer({
        label: `particle buffer ${i}`,
        size: this.particles.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
      });
      this.particleBuffers.push(particleBuffer);

      const gradientBuffer = device.createBuffer({
        label: `gradient buffer ${i}`,
        size: this.gradients.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
      });
      this.gradientBuffers.push(gradientBuffer);

      const configBuffer = device.createBuffer({
        label: `particle configuration buffer ${i}`,
        size: this.particleConfiguration.byteLength,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
      });
      this.particleConfigurationBuffers.push(configBuffer);

      const bindGroup = device.createBindGroup({
        label: `particle buffer bind group ${i}`,
        layout: this.bindGroupLayout.particle,
        entries: [{
          binding: 0,
          resource: { buffer: configBuffer },
        }, {
          binding: 1,
          resource: { buffer: particleBuffer },
        }, {
          binding: 2,
          resource: { buffer: gradientBuffer },
        }],
      });
      this.particleBindGroups.push(bindGroup);
    }

    for (let i = 0; i < this.maxParticles; i++) {
      // position
      const position = [Math.random(), Math.random()];
      this.particles[i * this.particleStride + 0] = position[0];
      this.particles[i * this.particleStride + 1] = position[1];
      this.particles[i * this.particleStride + 2] = 0.0;
      this.particles[i * this.particleStride + 3] = 1.0;
      // value.x
      this.particles[i * this.particleStride + 4] = test(position);
    }
  }

  constructor(device, canvasFormat, viewUniformBuffer, shaderModule)
  {
    const label = "flat";

    //////////////////////////////////////////////////////////////////////
    // buffer
    //////////////////////////////////////////////////////////////////////

    const array = new Uint16Array(6);
    array[0] = 0;
    array[1] = 1;
    array[2] = 2;
    array[3] = 0;
    array[4] = 3;
    array[5] = 1;

    this.buffer = device.createBuffer({
      label: `${label} renderer buffer`,
      size: array.byteLength,
      usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST,
    });

    device.queue.writeBuffer(this.buffer, 0, array, 0, array.length);

    //////////////////////////////////////////////////////////////////////
    // pipeline
    //////////////////////////////////////////////////////////////////////

    this.bindGroupLayout = {
      view: device.createBindGroupLayout({
        label: `${label} bind group layout view`,
        entries: [{
          binding: 0,
          visibility: GPUShaderStage.VERTEX,
          buffer: { type: "uniform" }
        }]
      }),
      particle: device.createBindGroupLayout({
        label: `${label} bind group layout particle`,
        entries: [{
          binding: 0,
          visibility: GPUShaderStage.FRAGMENT,
          buffer: { type: "uniform" }
        }, {
          binding: 1,
          visibility: GPUShaderStage.FRAGMENT,
          buffer: { type: "read-only-storage" }
        }, {
          binding: 2,
          visibility: GPUShaderStage.FRAGMENT,
          buffer: { type: "read-only-storage" }
        }]
      }),
    };

    const pipelineLayout = device.createPipelineLayout({
      label: `${label} pipeline layout`,
      bindGroupLayouts: [
        this.bindGroupLayout.view,
        this.bindGroupLayout.particle
      ],
    });

    this.renderPipeline = device.createRenderPipeline({
      label: `${label} pipeline`,
      layout: pipelineLayout,
      vertex: {
        module: shaderModule,
        entryPoint: "vertexMain",
      },
      fragment: {
        module: shaderModule,
        entryPoint: "fragmentMain",
        targets: [{
          format: canvasFormat,
        }]
      },
      primitive: {
        topology: "triangle-list",
      },
      depthStencil: {
        depthWriteEnabled: true,
        depthCompare: "less",
        format: "depth24plus",
      },
    });

    //////////////////////////////////////////////////////////////////////
    // particles
    //////////////////////////////////////////////////////////////////////

    this.createParticles(device);

    //////////////////////////////////////////////////////////////////////
    // bind group
    //////////////////////////////////////////////////////////////////////

    this.bindGroups = [
      device.createBindGroup({
        label: `${label} bind group`,
        layout: this.bindGroupLayout.view,
        entries: [{
          binding: 0,
          resource: { buffer: viewUniformBuffer },
        }],
      }),
    ];
  }

  updateGradients(device, frameNumber, configuration)
  {
    const radius = configuration.circleRadius;
    const count = configuration.particleCount;

    // calculate particle gradients
    for (let y = 0; y < this.maxGradientsDim; y++) {
      for (let x = 0; x < this.maxGradientsDim; x++) {
        const i = y * this.maxGradientsDim + x;
        const position = [(x / this.maxGradientsDim) + 0.5 / this.maxGradientsDim,
                          (y / this.maxGradientsDim) + 0.5 / this.maxGradientsDim];

        const gradient = this.positionGradient(position, radius, count);

        this.gradients[i * this.gradientStride + 0] = position[0];
        this.gradients[i * this.gradientStride + 1] = position[1];

        this.gradients[i * this.gradientStride + 4] = gradient[0];
        this.gradients[i * this.gradientStride + 5] = gradient[1];
      }
    }

    device.queue.writeBuffer(this.gradientBuffers[frameNumber], 0,
                             this.gradients, 0, this.gradients.length);
  }

  updateParticles(device, frameNumber, configuration)
  {
    for (var i = 0; i < configuration.particleCount; i++) {
      const position = [this.particles[i * this.particleStride + 0],
                        this.particles[i * this.particleStride + 1]];
      const density = this.particleDensity(position,
                                           configuration.circleRadius,
                                           configuration.particleCount);
      // value.y
      this.particles[i * this.particleStride + 5] = density;
    }

    device.queue.writeBuffer(this.particleBuffers[frameNumber], 0,
                             this.particles, 0, this.particles.length);
  }

  updateSimulation(device, frameNumber, configuration)
  {
    for (var i = 0; i < configuration.particleCount; i++) {
      const force = this.positionPressureForce(i,
                                               configuration.circleRadius,
                                               configuration.particleCount,
                                               configuration.targetDensity,
                                               configuration.pressure,
                                               configuration.particleMass1);
      if (force[0] == NaN)
        throw new Error("force");

      const particleDensity = this.particles[i * this.particleStride + 5];

      const acceleration = [force[0] / particleDensity, force[1] / particleDensity];

      this.particles[i * this.particleStride + 8] += acceleration[0];
      this.particles[i * this.particleStride + 9] += acceleration[1] + 0.001;

      // collisions
    }

    for (var i = 0; i < configuration.particleCount; i++) {
      this.particles[i * this.particleStride + 0] -= this.particles[i * this.particleStride + 8];
      this.particles[i * this.particleStride + 1] -= this.particles[i * this.particleStride + 9];

      const position = [this.particles[i * this.particleStride + 0],
                        this.particles[i * this.particleStride + 1]];

      for (var c = 0; c < 2; c++) {
        if (position[c] < 0) {
          this.particles[i * this.particleStride + c] = 0.01;
          this.particles[i * this.particleStride + c + 8] *= -1 * 0.9;
        }

        if (position[c] > 1) {
          this.particles[i * this.particleStride + c] = 0.99;
          this.particles[i * this.particleStride + c + 8] *= -1 * 0.9;
        }
      }
    }
  }

  updateConfiguration(device, frameNumber, configuration)
  {
    const PARTICLE = 0;
    const CIRCLE = 1;
    const MOUSE = 2;
    const D = 3;

    const config = this.particleConfiguration;
    config[PARTICLE * 4 + 0] = configuration.particleCount;
    config[PARTICLE * 4 + 1] = configuration.particleSize;
    config[PARTICLE * 4 + 2] = configuration.particleMass1;
    config[PARTICLE * 4 + 3] = configuration.particleMass2;

    config[CIRCLE * 4 + 0] = configuration.circleRadius;
    config[CIRCLE * 4 + 1] = configuration.circleThickness;
    config[CIRCLE * 4 + 2] = configuration.testIntensity;
    config[CIRCLE * 4 + 3] = configuration.mass12Mix;

    config[MOUSE * 4 + 0] = configuration.mousePosition[0];
    config[MOUSE * 4 + 1] = configuration.mousePosition[1];
    config[MOUSE * 4 + 2] = configuration.lineThickness;
    config[MOUSE * 4 + 3] = configuration.lineLength;

    config[D * 4 + 0] = configuration.targetDensity;
    config[D * 4 + 1] = configuration.pressure;
    config[D * 4 + 2] = 0;
    config[D * 4 + 3] = 0;

    device.queue.writeBuffer(this.particleConfigurationBuffers[frameNumber], 0,
                             this.particleConfiguration, 0, this.particleConfiguration.length)
  }

  render(device, renderPass, frameNumber, configuration)
  {
    // particle update must be before gradient update
    this.updateParticles(device, frameNumber, configuration);
    this.updateGradients(device, frameNumber, configuration);
    this.updateSimulation(device, frameNumber, configuration);
    this.updateConfiguration(device, frameNumber, configuration);

    renderPass.setPipeline(this.renderPipeline);
    renderPass.setIndexBuffer(this.buffer, "uint16");
    renderPass.setBindGroup(0, this.bindGroups[0]);
    renderPass.setBindGroup(1, this.particleBindGroups[frameNumber]);
    renderPass.drawIndexed(6);
  }
}

async function loadFlat(device, canvasFormat, viewUniformBuffer)
{
  const flatWgsl = await getPath("flat.wgsl");
  const shaderModule = device.createShaderModule({
    label: "flag shader",
    code: flatWgsl,
  });

  const renderer = new FlatRenderer(device, canvasFormat, viewUniformBuffer, shaderModule);
  return renderer;
}

export { loadFlat };