xxxxxxxxxx
248
const N = 200;
const dt = 0.00005;
const mindist2 = 1 / (4 * N);
const G = 10.;
let vpos = new Float32Array(2*N);
let vvel = new Float32Array(2*N);
let vacc = new Float32Array(2*N);
let vforces = new Float32Array(2*N);
let vmass = new Float32Array(N);
let vcol = new Int32Array(N*3);
let options = new Float32Array(2);
let device, commandBuffer, vposBuffer, vmassBuffer, vforcesBuffer, vforcesStagingBuffer, optionsBuffer, pipeline, bindGroup;
function initStars() {
for(let i = 0; i < N; ++i) {
do {
x = random();
y = random();
} while(x*x + y*y < mindist2);
vpos[2*i] = x;
vpos[2*i + 1] = y;
if(i == 0) {
vmass[i] = 1000;
} else {
vmass[i] = random()*1000;
}
vcol[3*i] = random(255);
vcol[3*i + 1] = random(255);
vcol[3*i + 2] = random(255);
}
}
function drawStars() {
for(let i = 0; i < N; ++i) {
// fill(vcol[3*i], vcol[3*i + 1],
let speed = sqrt(vvel[2*i]*vvel[2*i] + vvel[2*i+1]*vvel[2*i+1]);
fill(speed*20, 100/ (speed), speed*10);
circle(vpos[2*i]*width, vpos[2*i + 1]*height, sqrt(vmass[i]) / 10);
}
}
// function computeGravitationalForces() {
// for(let j = 0; j < N; ++j) {
// fx = 0;
// fy = 0;
// for(let i = 0; i < N; ++i) {
// if(i == j) continue;
// diff_px = vpos[2*i] - vpos[2*j];
// diff_py = vpos[2*i + 1] - vpos[2*j + 1];
// l3 = pow(sqrt(diff_px*diff_px + diff_py*diff_py), 2);
// fx += vmass[i] * diff_px / l3;
// fy += vmass[i] * diff_py / l3;
// }
// fx *= G;
// fy *= G;
// vforces[2*j] = fx;
// vforces[2*j + 1] = fy;
// }
// }
async function updateStars() {
for(let i = 0; i < N; ++i) {
vpos[2*i] += vvel[2*i]*dt + 0.5 * vacc[2*i] * dt*dt;
vpos[2*i + 1] += vvel[2*i + 1]*dt + 0.5 * vacc[2*i + 1] * dt*dt;
}
await submitCommand();
//computeGravitationalForces();
for(let i = 0; i < N; ++i) {
accx = vforces[2*i] / vmass[i];
vvel[2*i] += 0.5 * (vacc[2*i] + accx) * dt;
vacc[2*i] = accx;
if(vpos[2*i] < 0) {
vpos[2*i] = 0;
vvel[2*i] *= -1;
vacc[2*i] *= -1;
} else if(vpos[2*i] > 1) {
vpos[2*i] = 1;
vvel[2*i] *= -1;
vacc[2*i] *= -1;
}
accy = vforces[2*i + 1] / vmass[i]
vvel[2*i + 1] += 0.5 * (vacc[2*i + 1] + accy) * dt;
vacc[2*i + 1] = accy;
if(vpos[2*i + 1] < 0) {
vpos[2*i + 1] = 0;
vvel[2*i + 1] *= -1;
vacc[2*i + 1] *= -1;
} else if(vpos[2*i + 1] > 1) {
vpos[2*i + 1] = 1;
vvel[2*i + 1] *= -1;
vacc[2*i + 1] *= -1;
}
}
}
async function initWebGPU() {
const adapter = await navigator.gpu?.requestAdapter();
device = await adapter?.requestDevice();
if (!device) {
fail('need a browser that supports WebGPU');
return;
}
const module = device.createShaderModule({
label: 'gravity ComputeShaderModule',
code: `
@group(0) @binding(0) var<storage, read_write> vmass: array<f32>;
@group(0) @binding(1) var<storage, read_write> vpos: array<f32>;
@group(0) @binding(2) var<storage, read_write> vforces: array<f32>;
@group(0) @binding(3) var<storage, read_write> options: array<f32>;
@compute @workgroup_size(1) fn brief(
@builtin(global_invocation_id) id: vec3<u32>
) {
let j = id.x; // star id
// If I don't do this, the program crashes. Why?
vmass[j] = vmass[j];
vpos[2*j] = vpos[2*j];
vpos[2*j + 1] = vpos[2*j + 1];
vforces[2*j] = vforces[2*j];
vforces[2*j + 1] = vforces[2*j + 1];
options[0] = options[0];
var fx: f32 = 0.0;
var fy: f32 = 0.0;
for(var i: u32 = 0; i < u32(options[0]); i += 1) {
let diff_px = vpos[2*i] - vpos[2*j];
let diff_py = vpos[2*i + 1] - vpos[2*j + 1];
let l2 = diff_px*diff_px + diff_py*diff_py + 1e-10;
fx += vmass[i] * diff_px / l2;
fy += vmass[i] * diff_py / l2;
}
fx *= options[1];
fy *= options[1];
vforces[2*j] = fx;
vforces[2*j + 1] = fy;
}
`,
});
pipeline = device.createComputePipeline({
label: 'gravity compute',
layout: 'auto',
compute: {
module,
entryPoint: 'brief',
},
});
vposBuffer = device.createBuffer({
label: 'pos',
size: vpos.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
vmassBuffer = device.createBuffer({
label: 'mass',
size: vmass.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
vforcesBuffer = device.createBuffer({
label: 'forces buffer',
size: vforces.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
vforcesStagingBuffer = device.createBuffer({
label: 'forces staging buffer',
size: vforces.byteLength,
usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});
optionsBuffer = device.createBuffer({
label: 'options buffer',
size: options.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
});
bindGroup = device.createBindGroup({
label: 'bindGroup for work buffer',
layout: pipeline.getBindGroupLayout(0),
entries: [
{ binding: 0, resource: { buffer: vmassBuffer } },
{ binding: 1, resource: { buffer: vposBuffer } },
{ binding: 2, resource: { buffer: vforcesBuffer } },
{ binding: 3, resource: { buffer: optionsBuffer } }
],
});
}
async function submitCommand() {
const encoder = device.createCommandEncoder({
label: 'gravity encoder',
});
const pass = encoder.beginComputePass({
label: 'gravity compute pass',
});
pass.setPipeline(pipeline);
pass.setBindGroup(0, bindGroup);
pass.dispatchWorkgroups(N);
pass.end();
encoder.copyBufferToBuffer(vforcesBuffer, 0, vforcesStagingBuffer, 0, vforcesStagingBuffer.size);
// Finish encoding and submit the commands
commandBuffer = encoder.finish();
device.queue.writeBuffer(vposBuffer, 0, vpos);
device.queue.writeBuffer(vmassBuffer, 0, vmass);
device.queue.writeBuffer(optionsBuffer, 0, options);
device.queue.submit([commandBuffer]);
await vforcesStagingBuffer.mapAsync(GPUMapMode.READ);
vforces = new Float32Array(vforcesStagingBuffer.getMappedRange().slice());
await vforcesStagingBuffer.unmap()
}
async function gpu() {
options[0] = N;
options[1] = G;
await initWebGPU();
while(true) {
await updateStars();
}
}
function setup() {
createCanvas(windowWidth, windowHeight);
noStroke();
initStars();
gpu();
}
function draw() {
background(10);
drawStars();
}