Public
Edited
May 24, 2024
2 stars
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
renderable = {
const AGENT_COUNT = 1_000_000
const { device, context } = await gpu.init(S, S)

// observable-specific cleanup:
invalidation?.then(() => device.destroy())

const DISPATCH_SIZE = [Math.ceil(S / 16), Math.ceil(S / 16), 1]

const texInput = gpu.canvasTexture(device, sample.canvas, { usage: gpu.Usage.StorageTexture })

const texOutput = device.createTexture({
size: [S, S],
format: 'rgba8unorm',
usage:
GPUTextureUsage.STORAGE_BINDING |
GPUTextureUsage.COPY_SRC |
GPUTextureUsage.TEXTURE_BINDING
})

// Uniforms
const uniforms = new Float32Array([0, ...Object.values(config)]);
const uniformBuffer = gpu.uniformBuffer(device, uniforms);

// Agents

const agent = new Float32Array(
util.flatArr(AGENT_COUNT, i => {
const species = [0, 0, 0, 0]
species[util.randi(species.length)] = 1
return [
S / 2 + 20 * Math.cos(i / AGENT_COUNT * util.TWO_PI),
S / 2 + 20 * Math.sin(i / AGENT_COUNT * util.TWO_PI),
0, // padding
i / AGENT_COUNT * util.TWO_PI,
...species,
]
}),
);

const agentBufferPing = gpu.storageBuffer(device, agent);
const agentBufferPong = gpu.storageBuffer(device, agent);

// Trails
const trailCountsBuffer = gpu.storageBuffer(device, new Uint32Array(S * S));

const trailTexturePing = texInput;
const trailTexturePong = texOutput;
const bindGroupLayout = device.createBindGroupLayout({
label: 'everything else bind group layout',
entries: [
{
binding: 0,
visibility: GPUShaderStage.FRAGMENT | GPUShaderStage.COMPUTE,
buffer: {},
},
// agents in
{
binding: 1,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: { type: 'read-only-storage' },
},
// agents out
{
binding: 2,
visibility: GPUShaderStage.COMPUTE,
buffer: { type: 'storage' },
},
// trail counts
{
binding: 3,
visibility: GPUShaderStage.COMPUTE,
buffer: { type: 'storage' },
},
// trail texture IN
{
binding: 4,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
texture: {},
},
// trail texture OUT (blur + decay + counts)
{
binding: 5,
visibility: GPUShaderStage.COMPUTE,
storageTexture: { format: 'rgba8unorm' },
},
],
});
const agentBindGroupPing = device.createBindGroup({
label: 'agent bind group ping',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: uniformBuffer.buffer } },
{ binding: 1, resource: { buffer: agentBufferPing.buffer } },
{ binding: 2, resource: { buffer: agentBufferPong.buffer } },
{ binding: 3, resource: { buffer: trailCountsBuffer.buffer } },
{ binding: 4, resource: trailTexturePing.createView() },
{ binding: 5, resource: trailTexturePong.createView() },
],
});

const agentBindGroupPong = device.createBindGroup({
label: 'agent bind group pong',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: uniformBuffer.buffer } },
{ binding: 1, resource: { buffer: agentBufferPong.buffer } },
{ binding: 2, resource: { buffer: agentBufferPing.buffer } },
{ binding: 3, resource: { buffer: trailCountsBuffer.buffer } },
{ binding: 4, resource: trailTexturePong.createView() },
{ binding: 5, resource: trailTexturePing.createView() },
],
});

const pipelineLayout = device.createPipelineLayout({
label: 'pipeline layout',
bindGroupLayouts: [bindGroupLayout],
});
const agentPipeline = device.createComputePipeline({
label: 'agent pipeline',
layout: pipelineLayout,
compute: {
module: device.createShaderModule({ code: shaders.agent }),
entryPoint: 'cs',
},
});

const trailCombinePipeline = device.createComputePipeline({
label: 'agent pipeline',
layout: pipelineLayout,
compute: {
module: device.createShaderModule({ code: shaders.trailCombine }),
entryPoint: 'cs',
},
});

const debug = new TextureRenderer(device, texOutput, {
// colorScheme: gpu.canvasTexture(device, colorScheme),
style: { background: '#000' },
})

function render() {
// TIME
// let p = performance.now()
// let now = p - Math.floor(p);
const now = Date.now();
uniforms.set([now], 0);
uniformBuffer.write();

const encoder = device.createCommandEncoder();

const agentPass = encoder.beginComputePass();
agentPass.setPipeline(agentPipeline);
agentPass.setBindGroup(0, agentBindGroupPing);
agentPass.dispatchWorkgroups(Math.ceil(AGENT_COUNT / 64), 1, 1);
agentPass.setBindGroup(0, agentBindGroupPong);
agentPass.dispatchWorkgroups(Math.ceil(AGENT_COUNT / 64), 1, 1);
agentPass.end();
const trailCombinePass = encoder.beginComputePass();
trailCombinePass.setPipeline(trailCombinePipeline);

// for (let i = 0; i < 2; i++) {
trailCombinePass.setBindGroup(0, agentBindGroupPing);
trailCombinePass.dispatchWorkgroups(...DISPATCH_SIZE);

trailCombinePass.setBindGroup(0, agentBindGroupPong);
trailCombinePass.dispatchWorkgroups(...DISPATCH_SIZE);
// }
trailCombinePass.end();

debug.render()
device.queue.submit([encoder.finish()]);
}

render();

return { render, display: debug.canvas }
}
Insert cell
{
while (true) {
renderable.render();
yield null;
}
}
Insert cell
shaders = ({
agent: `
const PI: f32 = 3.141592653589793;
const TWO_PI: f32 = 6.283185307179586;

struct Uniforms {
time: f32,
sensor_angle: f32,
sensor_distance: f32,
rotation_angle: f32,
step_distance: f32,
decay_rate: f32,
deposit_amount: f32,
blur_strength: f32,
}

struct Agent {
pos: vec2f, // offset = 0, size = 8, alignment = 8
_pad: f32, // offset = 8, size = 4, alignment = 4 (to align next f32)
heading: f32, // offset = 12, size = 4, alignment = 4
species_mask: vec4f, // offset = 16, size = 16, alignment = 16
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage> agentsIn: array<Agent>;
@group(0) @binding(2) var<storage, read_write> agentsOut: array<Agent>;
@group(0) @binding(3) var<storage, read_write> trailCounts: array<atomic<u32>>;
@group(0) @binding(4) var trails_in: texture_2d<f32>;
fn grid_index_from_agent_pos(pos: vec2u) -> u32 {
return (pos.x % ${S}) + (pos.y % ${S}) * ${S};
}

// On generating random numbers, with help of y= [(a+x)sin(bx)] mod 1", W.J.J. Rey,
// 22nd European Meeting of Statisticians 1998
fn rand11(n: f32) -> f32 { return fract(sin(n) * 43758.5453123); }
fn rand22(n: vec2f) -> f32 { return fract(sin(dot(n, vec2f(12.9898, 4.1414))) * 43758.5453); }

// fn hash11(_p: f32) -> f32 {
// var p = fract(_p * .1031);
// p *= p + 33.33;
// p *= p + p;
// return fract(p);
// }

// fn hash12(p: vec2f) -> f32 {
// var p3 = fract(vec3f(p.xyx) * .1031);
// p3 += dot(p3, p3.yzx + 33.33);
// return fract((p3.x + p3.y) * p3.z);
// }

// fn hash13(_p3: vec3f) -> f32 {
// var p3 = fract(_p3 * .1031);
// p3 += dot(p3, p3.zyx + 31.32);
// return fract((p3.x + p3.y) * p3.z);
// }

fn shft(x: f32, size: f32) -> f32 {
return x + select(0.0, size, x < 0) + select(0.0, -size, x >= size);
}

fn next_agent_state(agent: Agent, texSize: vec2u) -> Agent {
let p = agent.pos;
let a = agent.heading;
let sa = uniforms.sensor_angle;
let sd = uniforms.sensor_distance;
let ra = uniforms.rotation_angle;
let step_dist = uniforms.step_distance;

let xc = p.x + cos(a) * sd;
let yc = p.y + sin(a) * sd;
let xl = p.x + cos(a - sa) * sd;
let yl = p.y + sin(a - sa) * sd;
let xr = p.x + cos(a + sa) * sd;
let yr = p.y + sin(a + sa) * sd;

// good enough
let C = textureLoad(trails_in, vec2u(u32(xc % ${S}), u32(yc % ${S})), 0).r;
let L = textureLoad(trails_in, vec2u(u32(xl % ${S}), u32(yl % ${S})), 0).r;
let R = textureLoad(trails_in, vec2u(u32(xr % ${S}), u32(yr % ${S})), 0).r;

var d = 0.0;
if (C > L && C > R) {
d = 0.0;
} else if (C < L && C < R) {
d = select(-1.0, 1.0, rand11(uniforms.time) > 0.5);
} else if (L < R) {
d = 1.0;
} else if (R < L) {
d = -1.0;
}

var rnd = rand11(uniforms.time) * (0.3 * TWO_PI) * select(-1.0, 1.0, rand11(uniforms.time) > 0.5);
// better random:
// var rnd = hash13(vec3f(p.x / ${S}, p.y / ${S}, uniforms.time)) * (0.01 * TWO_PI) * select(-1.0, 1.0, hash11(uniforms.time) > 0.5);

let da = ra * d;
var next_angle = shft(a + da + rnd, TWO_PI);
var next_pos = vec2f(
p.x + cos(next_angle) * step_dist,
p.y + sin(next_angle) * step_dist,
);

// TODO optimize
if (next_pos.x < 0.0) {
next_pos.x = 0.0;
next_angle += PI; // + (rand11(rnd) - 0.5) * 0.1;
} else if (next_pos.y < 0.0) {
next_pos.y = 0.0;
next_angle += PI;
} else if (next_pos.x >= f32(texSize.x)) {
next_pos.x = f32(texSize.x - 1);
next_angle += PI;
} else if (next_pos.y >= f32(texSize.y)) {
next_pos.y = f32(texSize.y - 1);
next_angle += PI;
}
return Agent(next_pos, 0, next_angle % TWO_PI, vec4f(0));
}

@compute
@workgroup_size(64, 1, 1)
fn cs(
@builtin(global_invocation_id) gid: vec3u,
@builtin(local_invocation_id) lid: vec3u,
) {
let agent_index = gid.x;
if (agent_index > arrayLength(&agentsIn)) {
return;
}
let agent = agentsIn[agent_index];

let texSize = textureDimensions(trails_in);

let next_agent = next_agent_state(agent, texSize);
agentsOut[agent_index].pos = next_agent.pos;
agentsOut[agent_index].heading = next_agent.heading;

// add to trail counts
let trail_index = grid_index_from_agent_pos(vec2u(u32(next_agent.pos.x), u32(next_agent.pos.y)));
atomicAdd(&trailCounts[trail_index], 1);
}
`,
trailCombine: `
struct Uniforms {
time: f32,
sensor_angle: f32,
sensor_distance: f32,
rotation_angle: f32,
step_distance: f32,
decay_rate: f32,
deposit_amount: f32,
blur_strength: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(3) var<storage, read_write> trailCounts: array<atomic<u32>>;
@group(0) @binding(4) var inputTexture: texture_2d<f32>;
@group(0) @binding(5) var outputTexture: texture_storage_2d<rgba8unorm, write>;

const tileSize: vec2u = vec2u(16, 16);
const blurRadius: i32 = 1;
const count: f32 = pow(f32(blurRadius) * 2 + 1, 2);

@compute
@workgroup_size(16, 16, 1)
fn cs(@builtin(global_invocation_id) global_id: vec3u) {
let texSize: vec2u = textureDimensions(inputTexture).xy;

let tileNW: vec2u = (global_id.xy / tileSize) * tileSize;
let localPos: vec2u = global_id.xy % tileSize;
let texPos: vec2u = tileNW + localPos;

let index = texPos.x + texPos.y * texSize.x;

if (texPos.x < texSize.x && texPos.y < texSize.y) {
var colorSum: vec4f = vec4f(0.0, 0.0, 0.0, 0.0);

for (var y = -blurRadius; y <= blurRadius; y++) {
for (var x = -blurRadius; x <= blurRadius; x++) {
let samplePos = vec2i(texPos) + vec2i(x, y);
if (
samplePos.x >= 0 && samplePos.x < i32(texSize.x) &&
samplePos.y >= 0 && samplePos.y < i32(texSize.y)
) {
colorSum = colorSum + textureLoad(inputTexture, samplePos, 0);
}
}
}
let blurred: vec4f = colorSum / count;

let count = f32(atomicLoad(&trailCounts[index])) * uniforms.deposit_amount;
var color = blurred - vec4f(uniforms.decay_rate) + count;
color = clamp(color, vec4f(0), vec4f(1));
textureStore(outputTexture, texPos, color);
}

// reset
atomicStore(&trailCounts[index], 0);
}
`
})
Insert cell
class TextureRenderer {
device = null
texture = null
context = null
sampler = null
pipeline = null
bindGroup = null
constructor(
device /* : GPUDevice */,
texture /* : GPUTexture */,
{
colorScheme /* ?: GPUTexture */,
binding,
format = null,
alphaMode = 'premultiplied',
style = {},
} = {}
) {
format ??= gpu.format()
this.device = device
this.context = gpu.context(texture.width, texture.height);
this.context.configure({ device, format, alphaMode })
for (const key in style) {
this.context.canvas.style[key] = style[key]
}
const module = device.createShaderModule({
code: TextureRenderer.Shader
})

this.sampler = gpu.sampler(device, gpu.LinearSampler)

this.pipeline = device.createRenderPipeline({
layout: 'auto',
vertex: {
module,
entryPoint: 'vs',
},
fragment: {
module,
entryPoint: 'fs',
targets: [{ format }],
},
})

this.bindGroup = device.createBindGroup({
layout: this.pipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: this.sampler },
{ binding: 1, resource: texture.createView() },
...(colorScheme ? [{ binding: 2, resource: colorScheme.createView() }] : [])
],
})
}

get canvas() {
return this.context.canvas
}
render() {
const encoder = this.device.createCommandEncoder();
const pass = encoder.beginRenderPass({
colorAttachments: [
{
view: this.context.getCurrentTexture().createView(),
clearValue: [0, 0, 0, 1],
loadOp: 'clear',
storeOp: 'store',
},
],
});
pass.setPipeline(this.pipeline);
pass.setBindGroup(0, this.bindGroup);
pass.draw(6);
pass.end();
this.device.queue.submit([encoder.finish()]);
}

static Shader = `
struct VertexOut {
@builtin(position) position: vec4f,
@location(0) texcoord: vec2f,
}
@group(0) @binding(0) var tex_sampler: sampler;
@group(0) @binding(1) var tex: texture_2d<f32>;
// @group(0) @binding(2) var colorScheme: texture_2d<f32>;

${shaderFragments.hsl2rgb}

@vertex
fn vs(
@builtin(vertex_index) vertexIndex : u32
) -> VertexOut {
let quad = array(
vec2f(0.0, 0.0),
vec2f(1.0, 0.0),
vec2f(0.0, 1.0),
vec2f(0.0, 1.0),
vec2f(1.0, 0.0),
vec2f(1.0, 1.0),
);
let texcoord = quad[vertexIndex];
return VertexOut(
vec4f((texcoord - 0.5) * 2, 0, 1),
texcoord
);
}

// this is just a dumping ground for color combinations
@fragment
fn fs(vs_out: VertexOut) -> @location(0) vec4f {
let color = textureSample(tex, tex_sampler, vs_out.texcoord);
return color;
// return vec4f(hsl2rgb(vec3f(mix(90.0 / 360.0, 350.0 / 360.0, color.r), 1.0, 0.65)), 1.0);
// return vec4f(color.r, 0.0, 0.0, 1.0);
// return vec4f(color.r, 0.0, 0.0, 1.0) * color.a; // nice
// return vec4f(hsl2rgb(vec3f(color.r, 1.0, 0.65)), 1.0);
// return vec4f(hsl2rgb(vec3f(mix(200.0 / 360.0, 0.95, color.r), 1.0, 0.65)), 1.0);
// return vec4f(hsl2rgb(vec3f(mix(200.0 / 360.0, 0.95, color.r), 1.0, mix(0.4, 1.0, color.r))), 1.0);
// return vec4f(hsl2rgb(vec3f(mix(200.0 / 360.0, 0.95, color.r), mix(0.3, 1.0, 1.0 - color.r), mix(0.1, 0.7, color.r))), 1.0);
// return vec4f(hsl2rgb(vec3f(mix(60.0 / 360.0, 320.0 / 360.0, color.r), 1.0, mix(0.1, 0.7, color.r))), 1.0);
// return vec4f(hsl2rgb(vec3f(1.0 - mix(200.0 / 360.0, 0.95, color.r), 1.0, 0.65)), color.a);
// return vec4f(hsl2rgb(vec3f(1.0 - mix(200.0 / 360.0, 0.95, color.r), 1.0, 0.65)), 1.0);
// return vec4f(hsl2rgb(vec3f(1.0 - mix(0.85, 0.2, color.r), 1.0, 0.5)), color.a);
// return vec4f(hsl2rgb(vec3f(1.0 - mix(0.85, 0.2, color.r), 1.0, 0.5)), 1.0);
// return vec4f(hsl2rgb(vec3f(color.r, 1.0, 1.0 - color.r)), 1.0);
// return vec4f(hsl2rgb(vec3f(color.r, 1.0, 1.0 - color.r)), 1.0) * color.a;
// return vec4f(color.r, 0.0, 1.0 - color.r, color.a);
}
`
}
Insert cell
sample = {
const size = 64
const count = 1024 / size
const ctx = gpu.context(1024, 1024, '2d')
ctx.fillStyle = '#ff0000'
// ctx.fillStyle = '#000'

ctx.font = 'bold 300px Helvetica'
ctx.textAlign = 'center'
ctx.fillText('slime', S / 2, S / 2)
// for (let y = 0; y < count; y++) {
// for (let x = -1; x < count; x += 2) {
// ctx.fillRect(x * size + y % 2 * size, y * size, size, size)
// }
// }

// for (let i = 0; i < 500; i++) {
// ctx.beginPath()
// ctx.arc(util.randi(S), util.randi(S), util.rand(30) + 5, 0, util.TWO_PI)
// ctx.fill()
// // ctx.fillRect(util.randi(S), util.randi(S), util.rand(100) + 5, util.rand(100) + 5)
// }
return ctx
}
Insert cell
Insert cell
Insert cell
Insert cell

Purpose-built for displays of data

Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
Learn more