Fix issues

src/index.scss

@@ -18,12 +18,10 @@ html {
 
 body {
   height: 100%;
-  background-color: hotpink;
   display: flex;
 }
 
 canvas {
-  background: hotpink;
   height: 100%;
   width: 100%;
 }

src/pipelines/agents/agent-pipeline.ts

@@ -1,4 +1,4 @@
-import { AGENT_SIZE, Agent } from './agent';
+import { AGENT_SIZE_IN_BYTES, Agent } from './agent';
 import shader from './agent.wgsl';
 
 export class AgentPipeline {
@@ -8,11 +8,16 @@ export class AgentPipeline {
   private readonly pipeline: GPUComputePipeline;
   private readonly uniforms: GPUBuffer;
   private readonly agentsBuffer: GPUBuffer;
+
   private bindGroup?: GPUBindGroup;
   private previousTrailMapIn?: GPUTexture;
   private previousTrailMapOut?: GPUTexture;
 
   public constructor(private readonly device: GPUDevice, agents: Array<Agent>) {
+    if (agents.length === 0) {
+      throw new Error('No agents provided');
+    }
+
     this.pipeline = device.createComputePipeline({
       layout: 'auto',
       compute: {
@@ -29,16 +34,16 @@ export class AgentPipeline {
     });
 
     const serializedAgents = new Float32Array(
-      new ArrayBuffer(agents.length * AGENT_SIZE)
+      new ArrayBuffer(agents.length * AGENT_SIZE_IN_BYTES)
     );
-    agents.forEach((agent, index) => {
+    agents.forEach((agent, i) => {
-      serializedAgents[index * 4 + 0] = agent.position[0];
+      serializedAgents[i * 4 + 0] = agent.position[0];
-      serializedAgents[index * 4 + 1] = agent.position[1];
+      serializedAgents[i * 4 + 1] = agent.position[1];
-      serializedAgents[index * 4 + 2] = agent.angle;
+      serializedAgents[i * 4 + 2] = agent.angle;
     });
 
     this.agentsBuffer = device.createBuffer({
-      size: agents.length * AGENT_SIZE,
+      size: agents.length * AGENT_SIZE_IN_BYTES,
       usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
     });
 
@@ -97,7 +102,9 @@ export class AgentPipeline {
     passEncoder.setPipeline(this.pipeline);
     passEncoder.setBindGroup(0, this.bindGroup);
     passEncoder.dispatchWorkgroups(
-      Math.ceil(this.agentsBuffer.size / AGENT_SIZE / AgentPipeline.WORKGROUP_SIZE)
+      Math.ceil(
+        this.agentsBuffer.size / AGENT_SIZE_IN_BYTES / AgentPipeline.WORKGROUP_SIZE
+      )
     );
     passEncoder.end();
   }

src/pipelines/agents/agent.ts

@@ -5,4 +5,4 @@ export interface Agent {
   angle: number;
 }
 
-export const AGENT_SIZE = 4;
+export const AGENT_SIZE_IN_BYTES = 4 * 4;

src/pipelines/agents/agent.wgsl

@@ -46,8 +46,6 @@ fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
 
   var agent = agents[id];
 
-
-
   var random = f32(hash(
     u32(
       agent.position.y * f32(settings.width) + agent.position.x

src/pipelines/diffusion/diffuse.wgsl

@@ -1,34 +1,28 @@
-struct VertexOutput {
+struct Settings {
-  @builtin(position) Position : vec4<f32>,
+  size : vec2<f32>,
-  @location(0) fragUV : vec2<f32>,
+  diffusionRate : f32,
-}
+  decayRate : f32,
-
+  deltaTime : f32,
-@vertex
+};
-fn vertex(@builtin(vertex_index) i : u32) -> VertexOutput {
-  var pos = array<vec2<f32>, 4>(
-    vec2(-1.0, 1.0),
-    vec2(-1.0, -1.0),
-    vec2(1.0, 1.0),
-    vec2(1.0, -1.0),
-  );
-
-  var output : VertexOutput;
-  output.Position = vec4<f32>(pos[i], 0.0, 1.0);
-  output.fragUV = output.Position.xy * 0.5 + 0.5;
-  return output;
-}
-
-
-@group(0) @binding(0) var mySampler: sampler;
-@group(0) @binding(1) var TargetTexture : texture_2d<f32>;
 
+@group(0) @binding(0) var<uniform> settings : Settings;
+@group(0) @binding(1) var Sampler: sampler;
+@group(0) @binding(2) var trailMap : texture_2d<f32>;
 
 @fragment
-fn fragment(@location(0) fragUV: vec2<f32>) -> @location(0) vec4<f32> {
+fn fragment(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
-  // return vec4(1.0, 0.0, 0.0, 1.0);
+  let current = textureSample(trailMap, Sampler, uv).rgb;
-  return mix(
+
-    vec4(textureSample(TargetTexture, mySampler, fragUV).rgb, 0.1),
+  let neighbours: vec3<f32> = (
-    vec4(1.0, 1.0, 1.0, 1.0), 
+      textureSample(trailMap, Sampler, uv + vec2<f32>(0, 1) / settings.size).rgb
-    0.01
+    + textureSample(trailMap, Sampler, uv + vec2<f32>(0, -1) / settings.size).rgb
+    + textureSample(trailMap, Sampler, uv + vec2<f32>(-1, 0) / settings.size).rgb
+    + textureSample(trailMap, Sampler, uv + vec2<f32>(1, 0) / settings.size).rgb
   );
+
+  return vec4(mix(
+    current,
+    neighbours / 4.0,
+    settings.diffusionRate
+  ) * (1.0 - settings.decayRate), 1.0);
 }

src/pipelines/diffusion/diffusion-pipeline.ts

@@ -1,19 +1,23 @@
+import { setUpFullScreenQuad } from '../../utils/full-screen-quad';
 import shader from './diffuse.wgsl';
 
 export class DiffusionPipeline {
+  private static readonly UNIFORM_COUNT = 6;
+
   private readonly pipeline: GPURenderPipeline;
+  private readonly uniforms: GPUBuffer;
+  private readonly quadVertexBuffer: GPUBuffer;
+
   private bindGroup?: GPUBindGroup;
   private previousTrailMapIn?: GPUTexture;
 
   public constructor(private readonly device: GPUDevice) {
+    const { buffer, vertex } = setUpFullScreenQuad(device);
+    this.quadVertexBuffer = buffer;
+
     this.pipeline = device.createRenderPipeline({
       layout: 'auto',
-      vertex: {
+      vertex,
-        module: device.createShaderModule({
-          code: shader,
-        }),
-        entryPoint: 'vertex',
-      },
       fragment: {
         module: device.createShaderModule({
           code: shader,
@@ -29,6 +33,31 @@ export class DiffusionPipeline {
         topology: 'triangle-strip',
       },
     });
+
+    this.uniforms = this.device.createBuffer({
+      size: DiffusionPipeline.UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+  }
+
+  public setParameters({
+    width,
+    height,
+    diffusionRate,
+    decayRate,
+    deltaTime,
+  }: {
+    width: number;
+    height: number;
+    diffusionRate: number;
+    decayRate: number;
+    deltaTime: number;
+  }) {
+    this.device.queue.writeBuffer(
+      this.uniforms,
+      0,
+      new Float32Array([width, height, diffusionRate, decayRate, deltaTime])
+    );
   }
 
   public execute(
@@ -49,11 +78,12 @@ export class DiffusionPipeline {
       ],
     };
 
-    const renderPassEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
+    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
-    renderPassEncoder.setBindGroup(0, this.bindGroup!);
+    passEncoder.setPipeline(this.pipeline);
-    renderPassEncoder.setPipeline(this.pipeline);
+    passEncoder.setVertexBuffer(0, this.quadVertexBuffer);
-    renderPassEncoder.draw(4, 1);
+    passEncoder.setBindGroup(0, this.bindGroup);
-    renderPassEncoder.end();
+    passEncoder.draw(4, 1);
+    passEncoder.end();
   }
 
   private ensureBindGroupExists(trailMapIn: GPUTexture) {
@@ -63,13 +93,19 @@ export class DiffusionPipeline {
         entries: [
           {
             binding: 0,
+            resource: {
+              buffer: this.uniforms,
+            },
+          },
+          {
+            binding: 1,
             resource: this.device.createSampler({
               magFilter: 'linear',
               minFilter: 'linear',
             }),
           },
           {
-            binding: 1,
+            binding: 2,
             resource: trailMapIn.createView(),
           },
         ],

src/pipelines/render/render-pipeline.ts

@@ -1,7 +1,10 @@
+import { setUpFullScreenQuad } from '../../utils/full-screen-quad';
 import shader from './render.wgsl';
 
 export class RenderPipeline {
   private readonly pipeline: GPURenderPipeline;
+  private readonly quadVertexBuffer: GPUBuffer;
+
   private bindGroup?: GPUBindGroup;
   private previousColorTexture?: GPUTexture;
 
@@ -10,14 +13,12 @@ export class RenderPipeline {
     private readonly device: GPUDevice,
     preferredCanvasFormat: GPUTextureFormat
   ) {
+    const { buffer, vertex } = setUpFullScreenQuad(device);
+    this.quadVertexBuffer = buffer;
+
     this.pipeline = device.createRenderPipeline({
       layout: 'auto',
-      vertex: {
+      vertex,
-        module: device.createShaderModule({
-          code: shader,
-        }),
-        entryPoint: 'vertex',
-      },
       fragment: {
         module: device.createShaderModule({
           code: shader,
@@ -42,17 +43,18 @@ export class RenderPipeline {
       colorAttachments: [
         {
           view: this.context.getCurrentTexture().createView(),
-          clearValue: { r: 0.0, g: 0.0, b: 0.0, a: 1.0 },
+          clearValue: { r: 1.0, g: 1.0, b: 1.0, a: 1.0 },
           loadOp: 'clear',
           storeOp: 'store',
         },
       ],
     };
-    const renderPassEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
+    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
-    renderPassEncoder.setBindGroup(0, this.bindGroup);
+    passEncoder.setPipeline(this.pipeline);
-    renderPassEncoder.setPipeline(this.pipeline);
+    passEncoder.setVertexBuffer(0, this.quadVertexBuffer);
-    renderPassEncoder.draw(4, 1);
+    passEncoder.setBindGroup(0, this.bindGroup);
-    renderPassEncoder.end();
+    passEncoder.draw(4, 1);
+    passEncoder.end();
   }
 
   private ensureBindGroupExists(colorTexture: GPUTexture) {

src/pipelines/render/render.wgsl

@@ -1,30 +1,8 @@
-struct VertexOutput {
-  @builtin(position) Position : vec4<f32>,
-  @location(0) fragUV : vec2<f32>,
-}
-
-@vertex
-fn vertex(@builtin(vertex_index) i : u32) -> VertexOutput {
-  var pos = array<vec2<f32>, 4>(
-    vec2(-1.0, 1.0),
-    vec2(-1.0, -1.0),
-    vec2(1.0, 1.0),
-    vec2(1.0, -1.0),
-  );
-
-  var output : VertexOutput;
-  output.Position = vec4<f32>(pos[i], 0.0, 1.0);
-  output.fragUV = output.Position.xy * 0.5 + 0.5;
-  return output;
-}
-
-
 @group(0) @binding(0) var mySampler: sampler;
 @group(0) @binding(1) var TargetTexture : texture_2d<f32>;
 
 
 @fragment
-fn fragment(@location(0) fragUV: vec2<f32>) -> @location(0) vec4<f32> {
+fn fragment(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
-//   return vec4(1.0, 0.0, 0.0, 1.0);
+  return textureSample(TargetTexture, mySampler, uv) * 10.0;
-  return textureSample(TargetTexture, mySampler, fragUV);
 }

src/renderer.ts

@@ -15,7 +15,7 @@ export default class Renderer {
 
   private agentPipeline: AgentPipeline;
   private renderPipeline: RenderPipeline;
-  private diffusionPipeline: any;
+  private diffusionPipeline: DiffusionPipeline;
 
   private preferredCanvasFormat: GPUTextureFormat;
   private trailMapA?: GPUTexture;
@@ -34,8 +34,8 @@ export default class Renderer {
     this.resize();
     window.addEventListener('resize', this.resize.bind(this));
 
-    const agents: Array<Agent> = new Array(settings.numAgents).fill(0).map(() => ({
+    const agents: Array<Agent> = new Array(settings.agentCount).fill(0).map(() => ({
-      position: vec2.fromValues(randomBetween(0, 500), randomBetween(0, 500)),
+      position: vec2.fromValues(randomBetween(0, 1000), randomBetween(0, 1000)),
       angle: randomBetween(0, Math.PI * 2),
     }));
 
@@ -54,21 +54,14 @@ export default class Renderer {
     this.canvas.height = this.canvas.clientHeight * devicePixelRatio;
 
     this.trailMapA?.destroy();
-    this.trailMapA = this.device.createTexture({
+    this.trailMapA = this.createTrailMap();
-      size: {
-        width: this.canvas.width,
-        height: this.canvas.height,
-        depthOrArrayLayers: 1,
-      },
-      format: 'rgba16float',
-      usage:
-        GPUTextureUsage.STORAGE_BINDING |
-        GPUTextureUsage.TEXTURE_BINDING |
-        GPUTextureUsage.RENDER_ATTACHMENT,
-    });
 
     this.trailMapB?.destroy();
-    this.trailMapB = this.device.createTexture({
+    this.trailMapB = this.createTrailMap();
+  }
+
+  private createTrailMap(): GPUTexture {
+    return this.device.createTexture({
       size: {
         width: this.canvas.width,
         height: this.canvas.height,
@@ -110,11 +103,17 @@ export default class Renderer {
       sensorAngleDegrees: 45,
       ...settings,
     });
+    this.diffusionPipeline.setParameters({
+      width: this.canvas.width,
+      height: this.canvas.height,
+      deltaTime: 0.016,
+      ...settings,
+    });
     const commandEncoder = this.device.createCommandEncoder();
 
     this.agentPipeline.execute(commandEncoder, this.trailMapA, this.trailMapB);
     this.diffusionPipeline.execute(commandEncoder, this.trailMapB, this.trailMapA);
-    this.renderPipeline.execute(commandEncoder, this.trailMapB);
+    this.renderPipeline.execute(commandEncoder, this.trailMapA);
     [this.trailMapA, this.trailMapB] = [this.trailMapB, this.trailMapA];
 
     this.queue.submit([commandEncoder.finish()]);

src/settings.ts

@@ -2,11 +2,11 @@ const SpawnMode = { Random: 0, Point: 1, InwardCircle: 2, RandomCircle: 3 };
 
 interface Settings {
   stepsPerFrame: number;
-  numAgents: number;
+  agentCount: number;
   spawnMode: number;
   trailWeight: number;
   decayRate: number;
-  diffuseRate: number;
+  diffusionRate: number;
   moveSpeed: number;
   turnSpeed: number;
   sensorAngleSpacing: number;
@@ -16,11 +16,12 @@ interface Settings {
 
 export const settings: Settings = {
   stepsPerFrame: 2,
-  numAgents: 250000,
+  agentCount: 500000,
   spawnMode: SpawnMode.InwardCircle,
   trailWeight: 5,
-  decayRate: 0.2,
+
-  diffuseRate: 3,
+  decayRate: 0.05,
+  diffusionRate: 0.1,
 
   moveSpeed: 20,
   turnSpeed: 2,

src/utils/full-screen-quad.ts

@@ -0,0 +1,52 @@
+import shader from './full-screen-quad.wgsl';
+
+export const setUpFullScreenQuad = (
+  device: GPUDevice
+): {
+  buffer: GPUBuffer;
+  vertex: GPUVertexState;
+} => {
+  const buffer = device.createBuffer({
+    size: 4 * 4 * 4, // 4x vec4<f32>
+    usage: GPUBufferUsage.VERTEX,
+    mappedAtCreation: true,
+  });
+  // prettier-ignore
+  const vertexData = [
+    // posX  posY U    V
+      -1.0, -1.0, 0.0, 1.0,
+      +1.0, -1.0, 1.0, 1.0,
+      -1.0, +1.0, 0.0, 0.0,
+      +1.0, +1.0, 1.0, 0.0,
+  ];
+  new Float32Array(buffer.getMappedRange()).set(vertexData);
+  buffer.unmap();
+
+  return {
+    buffer,
+    vertex: {
+      module: device.createShaderModule({
+        code: shader,
+      }),
+      entryPoint: 'vertex',
+      buffers: [
+        {
+          arrayStride: 4 * 4,
+          stepMode: 'vertex',
+          attributes: [
+            {
+              shaderLocation: 0,
+              offset: 0,
+              format: 'float32x2',
+            },
+            {
+              shaderLocation: 1,
+              offset: 8,
+              format: 'float32x2',
+            },
+          ],
+        },
+      ],
+    },
+  };
+};

src/utils/full-screen-quad.wgsl

@@ -0,0 +1,12 @@
+struct VertexOutput {
+  @builtin(position) position : vec4<f32>,
+  @location(0) uv : vec2<f32>,
+}
+
+@vertex
+fn vertex(
+  @location(0) position : vec2<f32>,
+  @location(1) uv : vec2<f32>
+) -> VertexOutput {
+  return VertexOutput(vec4(position, 0.0, 1.0), uv);
+}