006: Falling Atoms
2026-10-12
Building on Jiggling Atoms, the box now has a floor and a ceiling, and every particle is pulled down. Gravity is one extra line in the integrate pass: a constant acceleration -g added to each vertical velocity. The sides stay periodic, and the floor and ceiling reflect. The thermostat from last time still holds the temperature.
Cool the box and the liquid settles into a pool at the bottom, with vapour above it thinning out with height. In an isothermal gas the density falls off exponentially, n(y) \propto e^{-mgy/k_BT}, the barometric formula that sets the density of the air. Warm it and the surface evaporates, the pool dissolves into the atmosphere, and with gravity off you are back to droplets and a uniform gas.
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;
// 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 dt = 0.005;
const stepsPerFrame = 5;
const friction = 0.5; // thermostat coupling: how fast velocities are refreshed
const decay = Math.exp(-friction * dt); // per-step velocity damping
// 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)));
// Uniform buffer: the thermostat's kick size and gravity, which follow the
// sliders 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 DT = ${dt};
const RC2 = 6.25; // cutoff radius 2.5 sigma, squared
const WALL = 0.5; // floor and ceiling: half a particle from the edge
const DECAY = ${decay}; // thermostat: velocity kept each step
// noise: thermostat kick size, sqrt(T (1 - DECAY^2)); gravity: acceleration down
struct Params { noise: f32, gravity: 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;
${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);
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));
}
}
accelerations[i] = a;
}
@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 * 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 * DECAY + params.noise * gaussian2(&state);
rng[i] = state;
// ...then drift the position
var p = positions[i] + v * 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; }
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 } },
],
});
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] });
// One frame: `stepsPerFrame` simulation steps in a single compute pass, then
// draw. Dispatches in a pass run in order, and each sees the previous one's writes.
// The sliders are read here every frame.
const workgroups = Math.ceil(count / 64);
function frame() {
const noise = Math.sqrt(temperatureInput.value * (1 - decay * decay));
device.queue.writeBuffer(
paramBuffer,
0,
new Float32Array([noise, gravityInput.value]),
);
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();
const pass = gpu.beginRenderPass(encoder, context);
discs.draw(pass, count);
pass.end();
device.queue.submit([encoder.finish()]);
}
gpu.run(invalidation, device, frame);
// The canvas with the controls laid over its corner
yield ui.stage(canvas, [temperatureInput, gravityInput]);
}