Public
Edited
Mar 11
1 fork
16 stars
async function renderable(initialSimUniforms /*: Record<string, number>*/) {
const { device, context, format } = await gpu.init(w, h);
const CONFIG_UNIFORMS_ORDER = Object.keys(initialSimUniforms);
const module = device.createShaderModule({
label: 'vert & frag shader module',
code: /* wgsl */ `
struct Uniforms {
size: f32,
elapsed: f32,
random: f32,
padding1: f32, // align next mat4x4
matrix: mat4x4f, // viewProj matrix
resolution: vec2f,
padding2: vec2f, // struct size multiple of 16 bytes
}
struct Agent {
position: vec4f,
direction: vec4f,
}
struct VertexOut {
@builtin(position) position: vec4f,
@location(1) norm_index: f32,
@location(2) pos: vec3f,
@location(3) direction: vec3f,
@location(4) quad_pos: vec2f,
};
${ShaderFragment.hsl2rgb}
@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage, read> particles: array<Agent>;
@vertex
fn vs(
@builtin(instance_index) instance_index : u32,
@builtin(vertex_index) vertex_index : 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 position = particles[instance_index].position;
let clip_pos = uniforms.matrix * vec4f(position.xyz, 1);
let point_pos = vec4f((quad[vertex_index] - 0.5) * (uniforms.size / uniforms.resolution), 0, 0);
let out_pos = clip_pos + point_pos;
return VertexOut(
out_pos,
f32(instance_index) / f32(arrayLength(&particles)),
position.xyz,
particles[instance_index].direction.xyz,
quad[vertex_index] - 0.5,
);
}
@fragment
fn fs(vout: VertexOut) -> @location(0) vec4f {
if (distance(vout.quad_pos, vec2f(0)) > 0.5) {
discard;
}
// let color = vec4f(
// hsl2rgb(vec3f((1.0 - dist) * 0.4 + 0.6, 1.0, 0.7)),
// 1.0,
// );
let position_color = vec4f(
vout.pos.x * 0.5 + 0.5,
vout.pos.y * 0.5 + 0.5,
vout.pos.z * 0.5 + 0.5,
1.0,
);
let direction_color = vec4f(
vout.direction.x * 0.5 + 0.5,
vout.direction.y * 0.5 + 0.5,
vout.direction.z * 0.5 + 0.5,
1.0,
);
let dist = distance(vout.pos, vec3f(0));
let color = position_color * direction_color * (1.0 - dist);
// let color = position_color * direction_color * (1.0 - dist * 0.5);
// let color = vec4f(hsl2rgb(vec3f(dist, 1.0, 0.65)), 1.0) * position_color * dist * 0.75;
return color;
}
`,
});
const computeShader = /* wgsl */ `
struct Uniforms {
size: f32,
elapsed: f32,
random: f32,
padding1: f32, // align next mat4x4
matrix: mat4x4f,
resolution: vec2f,
padding2: vec2f, // struct size multiple of 16 bytes
}
struct SimulationUniforms {
${CONFIG_UNIFORMS_ORDER.map(u => `${u}: f32,`).join('\n')}
}
struct Agent {
position: vec4f,
direction: vec4f,
}
struct Sensed {
position: vec3f,
direction: vec3f,
}
@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage> particles_in: array<Agent>;
@group(0) @binding(2) var<storage, read_write> particles_out: array<Agent>;
@group(0) @binding(3) var<storage, read_write> trail_counts: array<atomic<u32>>;
@group(0) @binding(4) var trail_texture_input: texture_3d<f32>;
@group(0) @binding(6) var<uniform> sim_uniforms: SimulationUniforms;
${ShaderFragment.hash14}
// ported from threejs
fn apply_quaternion(q: vec4f, v: vec3f) -> vec3f {
// quaternion q is assumed to have unit length
let vx = v.x;
let vy = v.y;
let vz = v.z;
let qx = q.x;
let qy = q.y;
let qz = q.z;
let qw = q.w;
// t = 2 * cross( q.xyz, v );
let tx = 2 * (qy * vz - qz * vy);
let ty = 2 * (qz * vx - qx * vz);
let tz = 2 * (qx * vy - qy * vx);
// v + q.w * t + cross( q.xyz, t );
let next = vec3f(
vx + qw * tx + qy * tz - qz * ty,
vy + qw * ty + qz * tx - qx * tz,
vz + qw * tz + qx * ty - qy * tx,
);
return next;
}
// the above function is faster than this one (prob bc of cross())
// fn apply_quaternion(q: vec4f, v: vec3f) -> vec3f {
// let qv = vec3f(q.x, q.y, q.z);
// let t = 2.0 * cross(qv, v);
// return v + q.w * t + cross(qv, t);
// }
fn set_from_axis_angle(axis: vec3f, angle: f32) -> vec4f {
let halfAngle = angle * 0.5;
let s = sin(halfAngle);
return vec4f(axis * s, cos(halfAngle));
}
fn apply_axis_angle(axis: vec3f, angle: f32, v: vec3f) -> vec3f {
if (angle == 0.0) {
return v;
}
let q = set_from_axis_angle(axis, angle);
return apply_quaternion(q, v);
}
fn get_next_position(
center: vec3f,
direction: vec3f,
) -> Sensed {
let norm_direction = normalize(vec3f(direction));
let up = vec3f(direction.xyz);
let right = normalize(cross(norm_direction, up));
let adjustedUp = normalize(cross(right, norm_direction));
var theta_offset_index: f32 = 0.0;
var phi_offset_index: f32 = 0.0;
// experiment with this... -1 and 0
var selected_value: f32 = 0.0; // 0.0
// var selected_value: f32 = -1.0; // 0.0
for (var off: f32 = 0.0; off < 9.0; off += 1.0) {
let i = (off % 3.0) - 1.0;
let j = floor(off / 3.0) - 1.0;
let theta = sim_uniforms.theta_offset_range * i;
let phi = sim_uniforms.phi_offset_range * j;
let a = apply_axis_angle(right, theta, norm_direction);
let b = apply_axis_angle(adjustedUp, phi, a);
let local_direction = normalize(b) * sim_uniforms.sensor_distance;
let position = center + local_direction;
if (all(position >= vec3f(-1)) && all(position <= vec3f(1))) {
let texcoords = vec3u(floor((position * 0.5 + vec3f(0.5)) * 64.0));
let value = textureLoad(trail_texture_input, texcoords, 0).r;
if (
(value > selected_value) ||
// FIXME do a simpler random
(value == selected_value && hash14(vec4f(position.xyz, uniforms.random)) > 0.5)
) {
theta_offset_index = i;
phi_offset_index = j;
selected_value = value;
}
}
}
let theta = sim_uniforms.theta_offset_range_agent * theta_offset_index;
let phi = sim_uniforms.phi_offset_range_agent * phi_offset_index;
let a = apply_axis_angle(right, theta, norm_direction);
let b = apply_axis_angle(adjustedUp, phi, a);
let local_direction_distance = normalize(b) * sim_uniforms.step_distance;
return Sensed(center + local_direction_distance, normalize(local_direction_distance));
}
@compute
@workgroup_size(${AGENT_WORKGROUP_SIZE})
fn cs(@builtin(global_invocation_id) global_invocation_id: vec3u) {
let index = global_invocation_id.x;
if (index > arrayLength(&particles_in)) {
return;
}
let agent = particles_in[index];
var next = get_next_position(agent.position.xyz, agent.direction.xyz);
var next_pos = next.position;
var next_dir = next.direction;
// collide sphere
if (dot(next_pos, next_pos) > 1.0) {
let normal = normalize(next_pos);
next_pos = normal;
next_dir = next_dir - 2.0 * dot(next_dir, normal) * normal;
}
particles_out[index].position = vec4f(next_pos, 1.0);
particles_out[index].direction = vec4f(next_dir, 0.0);
let trail_counts_index = u32(
floor((next_pos.x * 0.5 + 0.5) * 64.0) +
floor((next_pos.y * 0.5 + 0.5) * 64.0) * 64.0 +
floor((next_pos.z * 0.5 + 0.5) * 64.0) * (64.0 * 64.0)
);
// update trail counts with new agent position
atomicAdd(&trail_counts[trail_counts_index], 1);
}
`;
const computeModule = device.createShaderModule({
label: 'compute shader module',
code: computeShader,
});
// define access to resources across all pipelines
const bindGroupLayout = device.createBindGroupLayout({
label: 'bind group layout',
entries: [
// uniforms
{
binding: 0,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: {},
},
// particles in
{
binding: 1,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: { type: 'read-only-storage' },
},
// particles 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: {
viewDimension: '3d',
},
},
// trail texture OUT (blur + decay + counts)
{
binding: 5,
visibility: GPUShaderStage.COMPUTE,
storageTexture: {
format: 'rgba8unorm',
viewDimension: '3d',
},
},
// sim uniforms
{
binding: 6,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: {},
},
],
});
/**
* app state uniforms
*/
// prettier-ignore
const uniforms = new Float32Array([
initialSimUniforms.particle_size, // size
0, // elapsed (delta) time between frames
Math.random(), // random
0, // padding1
...util.MAT4X4_IDENTITY, // mat4x4 as 16 floats
w, h, // resolution
0, 0, // padding2
]);
const uniformBuffer = device.createBuffer({
label: 'uniforms buffer',
size: uniforms.byteLength,
usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(uniformBuffer, 0, uniforms);
const matrixValue = uniforms.subarray(4, 4 + 16);
/**
* sim config uniforms
*/
const simUniforms = new Float32Array(
CONFIG_UNIFORMS_ORDER.map(u => initialSimUniforms[u]),
);
const simUniformBuffer = device.createBuffer({
label: 'sim uniforms buffer',
size: simUniforms.byteLength,
usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(simUniformBuffer, 0, simUniforms);
/**
* trails
*/
const trailCountsBuffer = gpu.storageBuffer(
device,
new Uint32Array(S * S * S),
);
const dims = { w: S, h: S, d: S };
const { texture: trailTexturePing, cleanup: cleanupTrailTexture } = tex3d(
device,
dims,
);
const trailTexturePong = device.createTexture({
dimension: '3d',
size: [S, S, S],
format: 'rgba8unorm',
usage:
GPUTextureUsage.STORAGE_BINDING |
GPUTextureUsage.COPY_SRC |
GPUTextureUsage.TEXTURE_BINDING |
GPUTextureUsage.RENDER_ATTACHMENT,
});
/**
* agent buffers setup
*/
// prettier-ignore
const points = new Float32Array(
util.flatArr(PARTICLE_COUNT, () => [
// ...randomPointInSphere(), 1.0, // position
0, 0, 0, 1.0, // position
util.rands(), util.rands(), util.rands(), // direction
0, // pad
]),
);
const pointsBufferA = device.createBuffer({
label: 'points storage buffer A',
size: points.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(pointsBufferA, 0, points);
const pointsBufferB = device.createBuffer({
label: 'points storage buffer B',
size: points.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(pointsBufferB, 0, points);
const bindGroupA = device.createBindGroup({
label: 'bind group layout A',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: uniformBuffer } },
{ binding: 1, resource: { buffer: pointsBufferA } },
{ binding: 2, resource: { buffer: pointsBufferB } },
{ binding: 3, resource: { buffer: trailCountsBuffer.buffer } },
{
binding: 4,
resource: trailTexturePing.createView({ dimension: '3d' }),
},
{
binding: 5,
resource: trailTexturePong.createView({ dimension: '3d' }),
},
{ binding: 6, resource: { buffer: simUniformBuffer } },
],
});
const bindGroupB = device.createBindGroup({
label: 'bind group layout B',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: uniformBuffer } },
{ binding: 1, resource: { buffer: pointsBufferB } }, // swap
{ binding: 2, resource: { buffer: pointsBufferA } },
{ binding: 3, resource: { buffer: trailCountsBuffer.buffer } },
// swap
{
binding: 4,
resource: trailTexturePong.createView({ dimension: '3d' }),
},
{
binding: 5,
resource: trailTexturePing.createView({ dimension: '3d' }),
},
{ binding: 6, resource: { buffer: simUniformBuffer } },
],
});
const pingpong = [bindGroupA, bindGroupB];
const pipelineLayout = device.createPipelineLayout({
label: 'pipeline layout',
bindGroupLayouts: [bindGroupLayout],
});
const pipeline = device.createRenderPipeline({
label: 'pipeline',
layout: pipelineLayout,
vertex: {
module,
entryPoint: 'vs',
},
fragment: {
module,
entryPoint: 'fs',
// @ts-expect-error
targets: [
{
format,
blend: gpu.Blending.Additive,
},
],
},
});
const computePipeline = device.createComputePipeline({
label: 'compute pipeline',
layout: pipelineLayout,
compute: {
module: computeModule,
entryPoint: 'cs',
},
});
const trailUpdatePipeline = device.createComputePipeline({
label: 'trail update pipeline',
layout: pipelineLayout,
compute: {
module: device.createShaderModule({
code: /* wgsl */ `
struct SimulationUniforms {
${CONFIG_UNIFORMS_ORDER.map(u => `${u}: f32,`).join('\n')}
}
const BLUR_RADIUS: i32 = 1;
const TEX_SIZE: vec3i = vec3i(64);
const COUNT: i32 = i32(pow(f32(BLUR_RADIUS * 2 + 1), 3.0));
// const offsets = array<vec3i, 27>(
// vec3i(-1, -1, -1), vec3i(0, -1, -1), vec3i(1, -1, -1),
// vec3i(-1, 0, -1), vec3i(0, 0, -1), vec3i(1, 0, -1),
// vec3i(-1, 1, -1), vec3i(0, 1, -1), vec3i(1, 1, -1),
// vec3i(-1, -1, 0), vec3i(0, -1, 0), vec3i(1, -1, 0),
// vec3i(-1, 0, 0), vec3i(0, 0, 0), vec3i(1, 0, 0),
// vec3i(-1, 1, 0), vec3i(0, 1, 0), vec3i(1, 1, 0),
// vec3i(-1, -1, 1), vec3i(0, -1, 1), vec3i(1, -1, 1),
// vec3i(-1, 0, 1), vec3i(0, 0, 1), vec3i(1, 0, 1),
// vec3i(-1, 1, 1), vec3i(0, 1, 1), vec3i(1, 1, 1)
// );
const offsets = array<vec3i, 6>(
vec3i( 1, 0, 0),
vec3i( 0, 1, 0),
vec3i( 0, 0, 1),
vec3i(-1, 0, 0),
vec3i( 0, -1, 0),
vec3i( 0, 0, -1),
);
@group(0) @binding(3) var<storage, read_write> trail_lookup: array<atomic<u32>>;
@group(0) @binding(4) var trail_texture_input: texture_3d<f32>;
@group(0) @binding(5) var trail_texture_output: texture_storage_3d<rgba8unorm, write>;
@group(0) @binding(6) var<uniform> sim_uniforms: SimulationUniforms;
@compute
@workgroup_size(${TRAIL_WORKGROUP_SIZE})
fn cs(
@builtin(global_invocation_id) global_id: vec3u,
) {
let texcoord = vec3i(global_id.xyz);
let index = (
global_id.x +
global_id.y * 64 +
global_id.z * 64 * 64
);
var colorSum: vec4f = vec4f(0);
for (var i = 0; i < 6; i += 1) {
let samplePos = texcoord + offsets[i];
if (all(samplePos >= vec3i(0)) && all(samplePos < TEX_SIZE)) {
colorSum = colorSum + textureLoad(trail_texture_input, samplePos, 0);
}
}
let blurred: vec4f = colorSum / 6.0;
var trail_count = f32(atomicLoad(&trail_lookup[index]));
// trail_count = max(trail_count + 1, 5.0);
var color = blurred - vec4f(sim_uniforms.decay_rate) + trail_count * sim_uniforms.deposit_amount * sim_uniforms.trail_count_dampener;
color = clamp(color, vec4f(0), vec4f(1));
textureStore(trail_texture_output, texcoord, color);
// reset
atomicStore(&trail_lookup[index], 0);
}
`,
}),
entryPoint: 'cs',
},
});
let step = 0;
let last = 0;
let elapsed;
// let mouse = [0, 0]
// camera
const projection = mat4.perspective((90 * Math.PI) / 180, w / h, 0.1, 256);
const view = mat4.lookAt(
[0, 0, 1.5], // position
[0, 0, 0], // target
[0, 1, 0], // up
);
const viewProjection = mat4.multiply(projection, view);
mat4.copy(viewProjection, matrixValue);
const renderPassConfig = {
colorAttachments: [
{
// clearValue: [0, 0, 0, 1],
// clearValue: [0, 0, 0.29, 1],
clearValue: [0.1, 0.1, 0.1, 1],
loadOp: 'clear',
storeOp: 'store',
view: undefined,
},
],
};
const handleZoom = e => {
e.preventDefault()
mat4.translate(
viewProjection,
[0, 0, e.wheelDelta * 0.0001],
viewProjection,
);
};
context.canvas.addEventListener('wheel', handleZoom);
function updateSimUniforms(uniformName /*: string */, value /*: any */) {
simUniforms.set([value], CONFIG_UNIFORMS_ORDER.indexOf(uniformName));
device.queue.writeBuffer(simUniformBuffer, 0, simUniforms);
}
let frame /*: number */;
function render(_t /*: number */) {
requestAnimationFrame(render);
let t = performance.now() * 0.001;
elapsed = t - last;
last = t;
uniforms.set([elapsed], 1);
uniforms.set([Math.random()], 2);
mat4.rotateY(viewProjection, t * 0.2, matrixValue);
device.queue.writeBuffer(uniformBuffer, 0, uniforms);
const encoder = device.createCommandEncoder();
const computePass = encoder.beginComputePass();
computePass.setPipeline(computePipeline);
computePass.setBindGroup(0, pingpong[step % 2]);
computePass.dispatchWorkgroups(AGENT_WORKGROUP_DISPATCH_COUNT);
computePass.end();
const trailComputePass = encoder.beginComputePass();
trailComputePass.setPipeline(trailUpdatePipeline);
// for (let i = 0; i < 2; i++) {
trailComputePass.setBindGroup(0, pingpong[step % 2]);
trailComputePass.dispatchWorkgroups(...TRAIL_WORKGROUP_DISPATCH_COUNT);
// step++;
// }
trailComputePass.end();
step++;
// render passes
renderPassConfig.colorAttachments[0].view = context
// @ts-expect-error
.getCurrentTexture()
.createView();
// @ts-expect-error
const pass = encoder.beginRenderPass(renderPassConfig);
pass.setPipeline(pipeline);
pass.setBindGroup(0, pingpong[step % 2]);
pass.draw(6, PARTICLE_COUNT);
pass.end();
device.queue.submit([encoder.finish()]);
}
return {
context,
animate: () => requestAnimationFrame(render),
stop: () => cancelAnimationFrame(frame),
cleanup: () => {
cancelAnimationFrame(frame);
uniformBuffer.destroy();
trailCountsBuffer.buffer.destroy();
pointsBufferA.destroy();
pointsBufferB.destroy();
cleanupTrailTexture();
trailTexturePong.destroy();
device.destroy();
},
updateSimUniforms,
}
}
ShaderFragment = ({
hsl2rgb: /* wgsl */`
fn hsl2rgb(hsl: vec3f) -> vec3f {
let c = vec3f(fract(hsl.x), clamp(hsl.yz, vec2f(0), vec2f(1)));
let rgb = clamp(abs((c.x * 6.0 + vec3f(0.0, 4.0, 2.0)) % 6.0 - 3.0) - 1.0, vec3f(0), vec3f(1));
return c.z + c.y * (rgb - 0.5) * (1.0 - abs(2.0 * c.z - 1.0));
}
`,
// www.cs.ubc.ca/~rbridson/docs/schechter-sca08-turbulence.pdf
hash: /* wgsl */ `
fn hash(ini: u32) -> u32 {
var state = ini;
state ^= 2747636419u;
state *= 2654435769u;
state ^= state >> 16;
state *= 2654435769u;
state ^= state >> 16;
state *= 2654435769u;
return state;
}
`,
// https://www.shadertoy.com/view/4djSRW (MIT License, (c) David Hoskins)
hash14: /* wgsl */ `
fn hash14(_p4: vec4f) -> f32 {
var p4 = fract(_p4 * vec4f(.1031, .1030, .0973, .1099));
p4 += dot(p4, p4.wzxy + 33.33);
return fract((p4.x + p4.y) * (p4.z + p4.w));
}
`,
shift: /* wgsl */ `
fn shift(x: f32, size: f32) -> f32 {
return x + select(0.0, size, x < 0) + select(0.0, -size, x >= size);
}
`,
})
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.
