particle system as high dim visualizers

ps.js

// ==========================================
// 3000 Coefficient Direct Mapping Visualizer
// Each coefficient = one particle
// No projection, no pairing
// ==========================================

let N = 3000;
let gridW;
let spacing;

let coeffs = [];
let grads = [];
let velocity = [];
let directions = [];

let optimizer = "adam";

let lr = 0.01;
let beta = 0.9;
let beta1 = 0.9;
let beta2 = 0.999;
let epsilon = 1e-8;

let m = [];
let v = [];
let t = 1;

function setup() {
  createCanvas(800, 800);
  pixelDensity(1);

  gridW = ceil(sqrt(N));
  spacing = width / gridW;

  for (let i = 0; i < N; i++) {
    coeffs[i] = random(-5, 5);
    grads[i] = 0;
    velocity[i] = 0;
    m[i] = 0;
    v[i] = 0;

    // fixed random direction per particle
    directions[i] = p5.Vector.random2D();
  }

  noStroke();
}

function draw() {
  background(10);

  computeGradients();
  updateCoefficients();
  drawParticles();

  t++;
}

// 3000D Rastrigin
function computeGradients() {
  for (let i = 0; i < N; i++) {
    grads[i] = 2 * coeffs[i] + 20 * PI * sin(2 * PI * coeffs[i]);
  }
}

function updateCoefficients() {

  for (let i = 0; i < N; i++) {

    if (optimizer === "sgd") {
      velocity[i] = -lr * grads[i];
      coeffs[i] += velocity[i];
    }

    else if (optimizer === "momentum") {
      velocity[i] = beta * velocity[i] - lr * grads[i];
      coeffs[i] += velocity[i];
    }

    else if (optimizer === "adam") {

      m[i] = beta1 * m[i] + (1 - beta1) * grads[i];
      v[i] = beta2 * v[i] + (1 - beta2) * grads[i] * grads[i];

      let m_hat = m[i] / (1 - pow(beta1, t));
      let v_hat = v[i] / (1 - pow(beta2, t));

      velocity[i] = -lr * m_hat / (sqrt(v_hat) + epsilon);
      coeffs[i] += velocity[i];
    }
  }
}

function drawParticles() {

  for (let i = 0; i < N; i++) {

    let gx = i % gridW;
    let gy = floor(i / gridW);

    let baseX = gx * spacing + spacing/2;
    let baseY = gy * spacing + spacing/2;

    // coefficient-driven displacement
    let disp = directions[i].copy().mult(coeffs[i] * 3);

    let x = baseX + disp.x;
    let y = baseY + disp.y;

    let g = grads[i];
    let intensity = constrain(abs(g) * 10, 0, 255);

    if (g > 0) fill(intensity, 80, 100);
    else fill(80, 100, intensity);

    let size = map(abs(velocity[i]), 0, 0.2, 2, 6, true);

    circle(x, y, size);
  }
}

function keyPressed() {
  if (key === '1') optimizer = "sgd";
  if (key === '2') optimizer = "momentum";
  if (key === '3') optimizer = "adam";
}