007: Reading Back
2026-10-12
Building on Falling Atoms, the same box, but now we ask it questions. Everything so far has been one-way: JavaScript uploads buffers, the GPU simulates and draws, and the numbers never come back. To plot anything we need them back. WebGPU makes that deliberately awkward. A buffer the shaders write cannot be read directly: you copy it into a staging buffer, ask for it to be mapped with mapAsync, and wait, because the copy only happens once the GPU has finished the work queued before it. So we do it every few frames, never every frame, and the charts lag the picture a little.
Two charts are drawn with Observable Plot. The left one is the energy per atom: kinetic, potential and total. Turn the thermostat off and the box is an isolated system, so the total should stay flat apart from small wiggles. That is the test of the integrator. Make the timestep bigger and watch the total drift; make it much bigger and the simulation blows up. The right one is the density of atoms against height. The dashed line is the barometric profile of an ideal gas, started at the top of the liquid pool: the vapour above follows it, while the pool below does not, because a liquid is nearly incompressible.
Show code
{
// Shared WebGPU boilerplate: setup, buffers and the animation loop (lib/gpu.js)
const gpu = await import(new URL("../lib/gpu.js", document.baseURI).href);
const md = await import(new URL("../lib/md.js", document.baseURI).href);
const ui = await import(new URL("../lib/ui.js", document.baseURI).href);
// The sliders are read live from their input elements (referencing `viewof x`
// gives the element, not the value), so moving them does not restart the
// simulation.
const temperatureInput = viewof temperature;
const gravityInput = viewof gravity;
const thermostatInput = viewof thermostat;
const timestepInput = viewof timestep;
// Setup, as before: the canvas is 16:9 and full width, and the drawing buffer
// follows its displayed size
const { canvas, device, context, format } = await gpu.setup({
width: 640,
height: 360,
resolution: "auto",
});
// Reduced units: sigma = epsilon = mass = k_B = 1. The box is 16:9 so it fills
// the canvas, and a particle (diameter sigma) is roughly 10 pixels across. At
// this density (about 0.34) the gas and the liquid can coexist.
const Lx = 112;
const Ly = 63;
const count = 2400;
const stepsPerFrame = 5;
const friction = 0.5; // thermostat coupling: how fast velocities are refreshed
const sampleEvery = 6; // read back every this many frames
const keep = 300; // samples kept in the energy chart
// Potential energy is shifted so it is zero at the 2.5 sigma cutoff; otherwise
// energy would jump every time a pair crossed it.
const cutoff = 2.5;
const shift = -4 * (Math.pow(1 / cutoff, 12) - Math.pow(1 / cutoff, 6));
// Initial state: a gas of random, non-overlapping positions, with Gaussian
// velocities at the starting temperature (see lib/md.js)
const { positions, velocities } = md.randomGas({
count,
box: [Lx, Ly],
temperature: temperatureInput.value,
periodic: [true, false],
margin: 0.5, // centres stay half a particle off the walls
});
// Storage buffers: positions, velocities and the accelerations between passes
const positionBuffer = gpu.storageBuffer(device, positions);
const velocityBuffer = gpu.storageBuffer(device, velocities);
const accelerationBuffer = gpu.storageBuffer(device, new Float32Array(count * 2));
// One random-number-generator state (a u32) per particle, seeded from the CPU
const rngBuffer = gpu.storageBuffer(device, crypto.getRandomValues(new Uint32Array(count)));
// Per-particle potential energy, written by the force pass and read back
const potentialBuffer = gpu.storageBuffer(device, new Float32Array(count));
// Staging buffer for readback: positions, velocities and potentials, back to
// back. A buffer that can be mapped for reading cannot also be a storage buffer,
// so the data is copied here first.
const positionBytes = count * 2 * 4;
const velocityBytes = count * 2 * 4;
const potentialBytes = count * 4;
const stagingBuffer = device.createBuffer({
size: positionBytes + velocityBytes + potentialBytes,
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});
// Uniform buffer: the thermostat, gravity and timestep, which follow the
// controls and are rewritten every frame
const paramBuffer = device.createBuffer({
size: 16, // uniform buffers are happiest in multiples of 16 bytes
usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});
// Compute shaders. Two entry points, so each pass only touches what it owns:
// `forces` reads every position but writes only its own acceleration, and
// `integrate` reads and writes only its own particle. No thread ever reads
// something another thread is writing in the same dispatch.
const computeModule = device.createShaderModule({
code: `
const L = vec2f(${Lx}, ${Ly}); // box size
const RC2 = 6.25; // cutoff radius 2.5 sigma, squared
const WALL = 0.5; // floor and ceiling: half a particle from the edge
const VSHIFT = ${shift}; // potential energy shift at the cutoff
// noise: thermostat kick size; gravity: acceleration down; decay: velocity
// kept each step by the thermostat (1 when it is off); dt: the timestep
struct Params { noise: f32, gravity: f32, decay: f32, dt: f32 }
@group(0) @binding(0) var<storage, read_write> positions: array<vec2f>;
@group(0) @binding(1) var<storage, read_write> velocities: array<vec2f>;
@group(0) @binding(2) var<storage, read_write> accelerations: array<vec2f>;
@group(0) @binding(3) var<storage, read_write> rng: array<u32>;
@group(0) @binding(4) var<uniform> params: Params;
@group(0) @binding(5) var<storage, read_write> potentials: array<f32>;
${md.rngWGSL}
@compute @workgroup_size(64)
fn forces(@builtin(global_invocation_id) id: vec3u) {
let i = id.x;
if (i >= arrayLength(&positions)) { return; }
let here = positions[i];
var a = vec2f(0.0);
var energy = 0.0;
for (var j = 0u; j < arrayLength(&positions); j++) {
if (j == i) { continue; }
var d = here - positions[j];
// Minimum image in x only: the box is periodic sideways, not vertically
d.x -= L.x * round(d.x / L.x);
let r2 = dot(d, d);
if (r2 < RC2) {
let s2 = 1.0 / r2; // (sigma / r)^2
let s6 = s2 * s2 * s2;
// force / r = 24 (2 s^12 - s^6) / r^2
a += d * (24.0 * s2 * s6 * (2.0 * s6 - 1.0));
// Pair energy, shared between the two particles
energy += 0.5 * (4.0 * (s6 * s6 - s6) + VSHIFT);
}
}
accelerations[i] = a;
potentials[i] = energy;
}
@compute @workgroup_size(64)
fn integrate(@builtin(global_invocation_id) id: vec3u) {
let i = id.x;
if (i >= arrayLength(&positions)) { return; }
// Leapfrog: kick the velocity, with gravity added to the pair forces...
let a = accelerations[i] + vec2f(0.0, -params.gravity);
var v = velocities[i] + a * params.dt;
// ...thermostat: slow it a little and add a random kick, which holds the
// mean kinetic energy at the target temperature...
var state = rng[i];
v = v * params.decay + params.noise * gaussian2(&state);
rng[i] = state;
// ...then drift the position
var p = positions[i] + v * params.dt;
p.x -= L.x * floor(p.x / L.x); // wrap sideways
// Floor and ceiling: reflect the position and flip the velocity
if (p.y < WALL) { p.y = 2.0 * WALL - p.y; v.y = -v.y; }
if (p.y > L.y - WALL) { p.y = 2.0 * (L.y - WALL) - p.y; v.y = -v.y; }
p.y = clamp(p.y, WALL, L.y - WALL); // a huge timestep must not escape the box
velocities[i] = v;
positions[i] = p;
}
`,
});
// With layout "auto" each pipeline only gets the bindings its entry point uses,
// so each needs its own bind group.
const forcePipeline = device.createComputePipeline({
layout: "auto",
compute: { module: computeModule, entryPoint: "forces" },
});
const integratePipeline = device.createComputePipeline({
layout: "auto",
compute: { module: computeModule, entryPoint: "integrate" },
});
const forceBindGroup = device.createBindGroup({
layout: forcePipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: { buffer: positionBuffer } },
{ binding: 2, resource: { buffer: accelerationBuffer } },
{ binding: 5, resource: { buffer: potentialBuffer } },
],
});
const integrateBindGroup = device.createBindGroup({
layout: integratePipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: { buffer: positionBuffer } },
{ binding: 1, resource: { buffer: velocityBuffer } },
{ binding: 2, resource: { buffer: accelerationBuffer } },
{ binding: 3, resource: { buffer: rngBuffer } },
{ binding: 4, resource: { buffer: paramBuffer } },
],
});
// Render: a disc per particle, coloured by its speed (see lib/md.js)
const discs = md.discs(device, { format, positionBuffer, velocityBuffer, box: [Lx, Ly] });
// Charts: one Observable Plot per host, redrawn after every readback
const [energyHost, densityHost] = charts.children;
const hot = "#FF4834";
const cold = "#5a80cc";
const samples = []; // {t, kind, e}: energy per atom against simulated time
const bins = 30;
const binHeight = Ly / bins;
let density = new Float32Array(bins); // atoms per sigma^2, smoothed over samples
function drawCharts(temperature, g) {
const width = energyHost.clientWidth || 480;
const style = ui.plotStyle;
energyHost.replaceChildren(
Plot.plot({
width,
height: 280,
style,
x: { label: "time (reduced units)" },
y: { label: "energy per atom", grid: true },
color: {
domain: ["kinetic", "potential", "total"],
range: [hot, cold, "currentColor"],
legend: true,
},
marks: [Plot.lineY(samples, { x: "t", y: "e", stroke: "kind" })],
})
);
// Barometric profile, n(y) = n_ref exp(-g (y - y_ref) / T), anchored at the
// surface of the pool: the lowest bin whose density is gas-like. Below that the
// liquid is nearly incompressible, so it does not follow the curve.
const surface = density.findIndex((n) => n > 0 && n < 0.25);
const yRef = (surface + 0.5) * binHeight;
const curve =
surface < 0
? []
: Array.from({ length: 61 }, (_, k) => {
const y = yRef + ((Ly - 0.5 - yRef) * k) / 60;
return { y, n: density[surface] * Math.exp((-g * (y - yRef)) / temperature) };
});
const rows = Array.from(density, (n, k) => ({
y0: k * binHeight,
y1: (k + 1) * binHeight,
n,
}));
densityHost.replaceChildren(
Plot.plot({
width,
height: 280,
style,
x: { label: "density (atoms per area)", grid: true, domain: [0, 1] },
y: { label: "height", domain: [0, Ly] },
marks: [
Plot.rectX(rows, { y1: "y0", y2: "y1", x: "n", fill: cold, fillOpacity: 0.85 }),
Plot.line(curve, { x: "n", y: "y", stroke: "currentColor", strokeDasharray: "4 3" }),
],
})
);
}
// Turn one readback into chart data
function record(data, time) {
const vel = data.subarray(count * 2, count * 4);
const pot = data.subarray(count * 4, count * 5);
const g = gravityInput.value;
let kinetic = 0;
let potential = 0;
const counts = new Float32Array(bins);
for (let i = 0; i < count; i++) {
kinetic += 0.5 * (vel[2 * i] ** 2 + vel[2 * i + 1] ** 2);
const y = data[2 * i + 1];
potential += pot[i] + g * (y - 0.5); // pair energy plus height
counts[Math.min(bins - 1, Math.floor(y / binHeight))]++;
}
samples.push(
{ t: time, kind: "kinetic", e: kinetic / count },
{ t: time, kind: "potential", e: potential / count },
{ t: time, kind: "total", e: (kinetic + potential) / count }
);
if (samples.length > keep * 3) samples.splice(0, 3);
for (let k = 0; k < bins; k++) {
const n = counts[k] / (Lx * binHeight);
density[k] = density[k] === 0 ? n : 0.9 * density[k] + 0.1 * n;
}
// In 2D the mean kinetic energy per atom is k_B T, so this is the measured T
drawCharts(kinetic / count, g);
}
// One frame: `stepsPerFrame` simulation steps in a single compute pass, then
// draw. Every `sampleEvery` frames the buffers are also copied for readback.
const workgroups = Math.ceil(count / 64);
let simTime = 0;
let frameCount = 0;
let pending = false; // a mapAsync is in flight: never copy into a mapped buffer
function frame() {
const dt = timestepInput.value;
const on = thermostatInput.value;
const decay = on ? Math.exp(-friction * dt) : 1;
const noise = on ? Math.sqrt(temperatureInput.value * (1 - decay * decay)) : 0;
device.queue.writeBuffer(
paramBuffer,
0,
new Float32Array([noise, gravityInput.value, decay, dt])
);
const encoder = device.createCommandEncoder();
const compute = encoder.beginComputePass();
for (let s = 0; s < stepsPerFrame; s++) {
compute.setPipeline(forcePipeline);
compute.setBindGroup(0, forceBindGroup);
compute.dispatchWorkgroups(workgroups);
compute.setPipeline(integratePipeline);
compute.setBindGroup(0, integrateBindGroup);
compute.dispatchWorkgroups(workgroups);
}
compute.end();
simTime += stepsPerFrame * dt;
const pass = gpu.beginRenderPass(encoder, context);
discs.draw(pass, count);
pass.end();
// Readback, step 1: copy what we want into the staging buffer
const sample = !pending && ++frameCount % sampleEvery === 0;
if (sample) {
encoder.copyBufferToBuffer(positionBuffer, 0, stagingBuffer, 0, positionBytes);
encoder.copyBufferToBuffer(
velocityBuffer, 0, stagingBuffer, positionBytes, velocityBytes);
encoder.copyBufferToBuffer(
potentialBuffer, 0, stagingBuffer, positionBytes + velocityBytes, potentialBytes);
}
device.queue.submit([encoder.finish()]);
// Readback, step 2: map it (resolves once the GPU is done), copy out, unmap
if (sample) {
pending = true;
const time = simTime;
stagingBuffer
.mapAsync(GPUMapMode.READ)
.then(() => {
const data = new Float32Array(stagingBuffer.getMappedRange().slice(0));
stagingBuffer.unmap();
pending = false;
record(data, time);
})
.catch(() => {}); // the device was destroyed because the cell re-ran
}
}
gpu.run(invalidation, device, frame);
// The canvas with the controls laid over its corner
yield ui.stage(canvas, [temperatureInput, thermostatInput, timestepInput, gravityInput]);
}