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();

}