umap-gpu 0.2.9 → 0.2.11

package/README.md

@@ -1,3 +1,5 @@
+# ⚠️ NOT READY YET | WORK IN PROGRESS ⚠️
+
 # umap-gpu
 
 UMAP dimensionality reduction with WebGPU-accelerated SGD and HNSW approximate nearest neighbors.

package/dist/index.js

@@ -1,95 +1,90 @@
 var H = Object.defineProperty;
-var V = (t, e, i) => e in t ? H(t, e, { enumerable: !0, configurable: !0, writable: !0, value: i }) : t[e] = i;
-var E = (t, e, i) => V(t, typeof e != "symbol" ? e + "" : e, i);
+var V = (e, n, r) => n in e ? H(e, n, { enumerable: !0, configurable: !0, writable: !0, value: r }) : e[n] = r;
+var G = (e, n, r) => V(e, typeof n != "symbol" ? n + "" : n, r);
 import { loadHnswlib as C } from "hnswlib-wasm";
-async function Y(t, e, i = {}) {
-  const { M: a = 16, efConstruction: s = 200, efSearch: p = 50 } = i, f = await C(), h = t[0].length, c = t.length, o = new f.HierarchicalNSW("l2", h, "");
-  o.initIndex(c, a, s, 200), o.setEfSearch(Math.max(p, e)), o.addItems(t, !1);
-  const n = [], r = [];
-  for (let l = 0; l < c; l++) {
-    const d = o.searchKnn(t[l], e + 1, void 0), u = d.neighbors.map((g, _) => ({ idx: g, dist: d.distances[_] })).filter(({ idx: g }) => g !== l).slice(0, e);
-    n.push(u.map(({ idx: g }) => g)), r.push(u.map(({ dist: g }) => g));
+async function Y(e, n, r = {}) {
+  const { M: s = 16, efConstruction: t = 200, efSearch: p = 50 } = r, f = await C(), l = e[0].length, c = e.length, a = new f.HierarchicalNSW("l2", l, "");
+  a.initIndex(c, s, t, 200), a.setEfSearch(Math.max(p, n)), a.addItems(e, !1);
+  const o = [], i = [];
+  for (let h = 0; h < c; h++) {
+    const d = a.searchKnn(e[h], n + 1, void 0), u = d.neighbors.map((g, m) => ({ idx: g, dist: d.distances[m] })).filter(({ idx: g }) => g !== h).slice(0, n);
+    o.push(u.map(({ idx: g }) => g)), i.push(u.map(({ dist: g }) => g));
   }
-  return { indices: n, distances: r };
+  return { indices: o, distances: i };
 }
-async function J(t, e, i = {}) {
-  const { M: a = 16, efConstruction: s = 200, efSearch: p = 50 } = i, f = await C(), h = t[0].length, c = t.length, o = new f.HierarchicalNSW("l2", h, "");
-  o.initIndex(c, a, s, 200), o.setEfSearch(Math.max(p, e)), o.addItems(t, !1);
-  const n = [], r = [];
+async function Q(e, n, r = {}) {
+  const { M: s = 16, efConstruction: t = 200, efSearch: p = 50 } = r, f = await C(), l = e[0].length, c = e.length, a = new f.HierarchicalNSW("l2", l, "");
+  a.initIndex(c, s, t, 200), a.setEfSearch(Math.max(p, n)), a.addItems(e, !1);
+  const o = [], i = [];
   for (let d = 0; d < c; d++) {
-    const u = o.searchKnn(t[d], e + 1, void 0), g = u.neighbors.map((_, w) => ({ idx: _, dist: u.distances[w] })).filter(({ idx: _ }) => _ !== d).slice(0, e);
-    n.push(g.map(({ idx: _ }) => _)), r.push(g.map(({ dist: _ }) => _));
+    const u = a.searchKnn(e[d], n + 1, void 0), g = u.neighbors.map((m, w) => ({ idx: m, dist: u.distances[w] })).filter(({ idx: m }) => m !== d).slice(0, n);
+    o.push(g.map(({ idx: m }) => m)), i.push(g.map(({ dist: m }) => m));
   }
-  return { knn: { indices: n, distances: r }, index: {
+  return { knn: { indices: o, distances: i }, index: {
     searchKnn(d, u) {
-      const g = [], _ = [];
+      const g = [], m = [];
       for (const w of d) {
-        const y = o.searchKnn(w, u, void 0), b = y.neighbors.map((M, x) => ({ idx: M, dist: y.distances[x] })).sort((M, x) => M.dist - x.dist).slice(0, u);
-        g.push(b.map(({ idx: M }) => M)), _.push(b.map(({ dist: M }) => M));
+        const b = a.searchKnn(w, u, void 0), y = b.neighbors.map((v, x) => ({ idx: v, dist: b.distances[x] })).sort((v, x) => v.dist - x.dist).slice(0, u);
+        g.push(y.map(({ idx: v }) => v)), m.push(y.map(({ dist: v }) => v));
       }
-      return { indices: g, distances: _ };
+      return { indices: g, distances: m };
     }
   } };
 }
-function D(t, e, i, a = 1) {
-  const s = t.length, { sigmas: p, rhos: f } = L(e, i), h = [], c = [], o = [];
-  for (let r = 0; r < s; r++)
-    for (let l = 0; l < t[r].length; l++) {
-      const d = e[r][l], u = d <= f[r] ? 1 : Math.exp(-((d - f[r]) / p[r]));
-      h.push(r), c.push(t[r][l]), o.push(u);
+function D(e, n, r, s = 1) {
+  const t = e.length, { sigmas: p, rhos: f } = L(n, r), l = [], c = [], a = [];
+  for (let i = 0; i < t; i++)
+    for (let h = 0; h < e[i].length; h++) {
+      const d = n[i][h], u = d <= f[i] ? 1 : Math.exp(-((d - f[i]) / p[i]));
+      l.push(i), c.push(e[i][h]), a.push(u);
     }
-  return { ...X(h, c, o, s, a), nVertices: s };
+  return { ...X(l, c, a, t, s), nVertices: t };
 }
-function Q(t, e, i) {
-  const a = t.length, { sigmas: s, rhos: p } = L(e, i), f = [], h = [], c = [];
-  for (let o = 0; o < a; o++)
-    for (let n = 0; n < t[o].length; n++) {
-      const r = e[o][n], l = r <= p[o] ? 1 : Math.exp(-((r - p[o]) / s[o]));
-      f.push(o), h.push(t[o][n]), c.push(l);
+function J(e, n, r) {
+  const s = e.length, { sigmas: t, rhos: p } = L(n, r), f = [], l = [], c = [];
+  for (let a = 0; a < s; a++)
+    for (let o = 0; o < e[a].length; o++) {
+      const i = n[a][o], h = i <= p[a] ? 1 : Math.exp(-((i - p[a]) / t[a]));
+      f.push(a), l.push(e[a][o]), c.push(h);
     }
   return {
     rows: new Float32Array(f),
-    cols: new Float32Array(h),
+    cols: new Float32Array(l),
     vals: new Float32Array(c),
-    nVertices: a
+    nVertices: s
   };
 }
-function L(t, e) {
-  const a = t.length, s = new Float32Array(a), p = new Float32Array(a);
-  for (let f = 0; f < a; f++) {
-    const h = t[f];
-    p[f] = h.find((l) => l > 0) ?? 0;
-    let c = 0, o = 1 / 0, n = 1;
-    const r = Math.log2(e);
-    for (let l = 0; l < 64; l++) {
+function L(e, n) {
+  const s = e.length, t = new Float32Array(s), p = new Float32Array(s);
+  for (let f = 0; f < s; f++) {
+    const l = e[f];
+    p[f] = l.find((h) => h > 0) ?? 0;
+    let c = 0, a = 1 / 0, o = 1;
+    const i = Math.log2(n);
+    for (let h = 0; h < 64; h++) {
       let d = 0;
-      for (let u = 1; u < h.length; u++)
-        d += Math.exp(-Math.max(0, h[u] - p[f]) / n);
-      if (Math.abs(d - r) < 1e-5) break;
-      d > r ? (o = n, n = (c + o) / 2) : (c = n, n = o === 1 / 0 ? n * 2 : (c + o) / 2);
+      for (let u = 0; u < l.length; u++)
+        d += Math.exp(-Math.max(0, l[u] - p[f]) / o);
+      if (Math.abs(d - i) < 1e-5) break;
+      d > i ? (a = o, o = (c + a) / 2) : (c = o, o = a === 1 / 0 ? o * 2 : (c + a) / 2);
     }
-    s[f] = n;
+    t[f] = o;
   }
-  return { sigmas: s, rhos: p };
+  return { sigmas: t, rhos: p };
 }
-function X(t, e, i, a, s) {
-  const p = /* @__PURE__ */ new Map(), f = (n, r, l) => {
-    const d = n * a + r;
-    p.set(d, (p.get(d) ?? 0) + l);
-  };
-  for (let n = 0; n < t.length; n++)
-    f(t[n], e[n], i[n]), f(e[n], t[n], i[n]);
-  const h = [], c = [], o = [];
-  for (const [n, r] of p.entries()) {
-    const l = Math.floor(n / a), d = n % a;
-    h.push(l), c.push(d), o.push(
-      r > 1 ? s * (2 - r) + (1 - s) * (r - 1) : r
-    );
+function X(e, n, r, s, t) {
+  const p = /* @__PURE__ */ new Map();
+  for (let a = 0; a < e.length; a++)
+    p.set(e[a] * s + n[a], r[a]);
+  const f = [], l = [], c = [];
+  for (const [a, o] of p) {
+    const i = Math.floor(a / s), h = a % s, d = p.get(h * s + i) ?? 0, u = o + d - o * d, g = o * d;
+    f.push(i), l.push(h), c.push(t * u + (1 - t) * g);
   }
   return {
-    rows: new Float32Array(h),
-    cols: new Float32Array(c),
-    vals: new Float32Array(o)
+    rows: new Float32Array(f),
+    cols: new Float32Array(l),
+    vals: new Float32Array(c)
   };
 }
 const Z = `// UMAP SGD compute shader — processes one graph edge per GPU thread.
@@ -150,13 +145,20 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     dist_sq += diff * diff;
   }
 
-  let grad_coeff_attr = -2.0 * params.a * params.b * pow(dist_sq, params.b - 1.0)
-                        / (params.a * pow(dist_sq, params.b) + 1.0);
+  let pow_b = pow(dist_sq, params.b);
+  // Guard dist_sq == 0: b-1 is negative so pow(0, b-1) = +Inf.
+  // Mirror CPU: use pow_b / dist_sq only when dist_sq > 0, else 0.
+  let grad_coeff_attr = select(
+    -2.0 * params.a * params.b * (pow_b / dist_sq) / (params.a * pow_b + 1.0),
+    0.0,
+    dist_sq == 0.0
+  );
 
   for (var d = 0u; d < nc; d++) {
     let diff = embedding[i * nc + d] - embedding[j * nc + d];
     let grad = clip(grad_coeff_attr * diff, -4.0, 4.0);
     embedding[i * nc + d] += params.alpha * grad;
+    embedding[j * nc + d] -= params.alpha * grad;
   }
 
   epoch_of_next_sample[edge_idx] += epochs_per_sample[edge_idx];
@@ -197,13 +199,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 `;
 class W {
   constructor() {
-    E(this, "device");
-    E(this, "pipeline");
+    G(this, "device");
+    G(this, "pipeline");
   }
   async init() {
-    const e = await navigator.gpu.requestAdapter();
-    if (!e) throw new Error("WebGPU not supported");
-    this.device = await e.requestDevice(), this.pipeline = this.device.createComputePipeline({
+    const n = await navigator.gpu.requestAdapter();
+    if (!n) throw new Error("WebGPU not supported");
+    this.device = await n.requestDevice(), this.pipeline = this.device.createComputePipeline({
       layout: "auto",
       compute: {
         module: this.device.createShaderModule({ code: Z }),
@@ -224,218 +226,218 @@ class W {
    * @param params      - UMAP curve parameters and repulsion settings
    * @returns Optimized embedding as Float32Array
    */
-  async optimize(e, i, a, s, p, f, h, c, o) {
-    const { device: n } = this, r = i.length, l = this.makeBuffer(
-      e,
+  async optimize(n, r, s, t, p, f, l, c, a) {
+    const { device: o } = this, i = r.length, h = this.makeBuffer(
+      n,
       GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
-    ), d = this.makeBuffer(i, GPUBufferUsage.STORAGE), u = this.makeBuffer(a, GPUBufferUsage.STORAGE), g = this.makeBuffer(s, GPUBufferUsage.STORAGE), _ = new Float32Array(r).fill(0), w = this.makeBuffer(_, GPUBufferUsage.STORAGE), y = new Float32Array(r);
-    for (let m = 0; m < r; m++)
-      y[m] = s[m] / c.negativeSampleRate;
-    const b = this.makeBuffer(y, GPUBufferUsage.STORAGE), M = new Uint32Array(r);
-    for (let m = 0; m < r; m++)
-      M[m] = Math.random() * 4294967295 | 0;
-    const x = this.makeBuffer(M, GPUBufferUsage.STORAGE), B = n.createBuffer({
+    ), d = this.makeBuffer(r, GPUBufferUsage.STORAGE), u = this.makeBuffer(s, GPUBufferUsage.STORAGE), g = this.makeBuffer(t, GPUBufferUsage.STORAGE), m = new Float32Array(i).fill(0), w = this.makeBuffer(m, GPUBufferUsage.STORAGE), b = new Float32Array(i);
+    for (let _ = 0; _ < i; _++)
+      b[_] = t[_] / c.negativeSampleRate;
+    const y = this.makeBuffer(b, GPUBufferUsage.STORAGE), v = new Uint32Array(i);
+    for (let _ = 0; _ < i; _++)
+      v[_] = Math.random() * 4294967295 | 0;
+    const x = this.makeBuffer(v, GPUBufferUsage.STORAGE), B = o.createBuffer({
       size: 40,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
     });
-    for (let m = 0; m < h; m++) {
-      const F = 1 - m / h, A = new ArrayBuffer(40), v = new Uint32Array(A), N = new Float32Array(A);
-      v[0] = r, v[1] = p, v[2] = f, v[3] = m, v[4] = h, N[5] = F, N[6] = c.a, N[7] = c.b, N[8] = c.gamma, v[9] = c.negativeSampleRate, n.queue.writeBuffer(B, 0, A);
-      const S = n.createBindGroup({
+    for (let _ = 0; _ < l; _++) {
+      const F = 1 - _ / l, A = new ArrayBuffer(40), M = new Uint32Array(A), N = new Float32Array(A);
+      M[0] = i, M[1] = p, M[2] = f, M[3] = _, M[4] = l, N[5] = F, N[6] = c.a, N[7] = c.b, N[8] = c.gamma, M[9] = c.negativeSampleRate, o.queue.writeBuffer(B, 0, A);
+      const S = o.createBindGroup({
         layout: this.pipeline.getBindGroupLayout(0),
         entries: [
           { binding: 0, resource: { buffer: g } },
           { binding: 1, resource: { buffer: d } },
           { binding: 2, resource: { buffer: u } },
-          { binding: 3, resource: { buffer: l } },
+          { binding: 3, resource: { buffer: h } },
           { binding: 4, resource: { buffer: w } },
-          { binding: 5, resource: { buffer: b } },
+          { binding: 5, resource: { buffer: y } },
           { binding: 6, resource: { buffer: B } },
           { binding: 7, resource: { buffer: x } }
         ]
-      }), k = n.createCommandEncoder(), U = k.beginComputePass();
-      U.setPipeline(this.pipeline), U.setBindGroup(0, S), U.dispatchWorkgroups(Math.ceil(r / 256)), U.end(), n.queue.submit([k.finish()]), m % 10 === 0 && (await n.queue.onSubmittedWorkDone(), o == null || o(m, h));
+      }), O = o.createCommandEncoder(), U = O.beginComputePass();
+      U.setPipeline(this.pipeline), U.setBindGroup(0, S), U.dispatchWorkgroups(Math.ceil(i / 256)), U.end(), o.queue.submit([O.finish()]), _ % 10 === 0 && (await o.queue.onSubmittedWorkDone(), a == null || a(_, l));
     }
-    const G = n.createBuffer({
-      size: e.byteLength,
+    const E = o.createBuffer({
+      size: n.byteLength,
       usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
-    }), R = n.createCommandEncoder();
-    R.copyBufferToBuffer(l, 0, G, 0, e.byteLength), n.queue.submit([R.finish()]), await G.mapAsync(GPUMapMode.READ);
-    const O = new Float32Array(G.getMappedRange().slice(0));
-    return G.unmap(), l.destroy(), d.destroy(), u.destroy(), g.destroy(), w.destroy(), b.destroy(), x.destroy(), B.destroy(), G.destroy(), O;
+    }), R = o.createCommandEncoder();
+    R.copyBufferToBuffer(h, 0, E, 0, n.byteLength), o.queue.submit([R.finish()]), await E.mapAsync(GPUMapMode.READ);
+    const k = new Float32Array(E.getMappedRange().slice(0));
+    return E.unmap(), h.destroy(), d.destroy(), u.destroy(), g.destroy(), w.destroy(), y.destroy(), x.destroy(), B.destroy(), E.destroy(), k;
   }
-  makeBuffer(e, i) {
-    const a = this.device.createBuffer({
-      size: e.byteLength,
-      usage: i,
+  makeBuffer(n, r) {
+    const s = this.device.createBuffer({
+      size: n.byteLength,
+      usage: r,
       mappedAtCreation: !0
     });
-    return e instanceof Float32Array ? new Float32Array(a.getMappedRange()).set(e) : new Uint32Array(a.getMappedRange()).set(e), a.unmap(), a;
+    return n instanceof Float32Array ? new Float32Array(s.getMappedRange()).set(n) : new Uint32Array(s.getMappedRange()).set(n), s.unmap(), s;
   }
 }
-function P(t) {
-  return Math.max(-4, Math.min(4, t));
+function q(e) {
+  return Math.max(-4, Math.min(4, e));
 }
-function q(t, e, i, a, s, p, f, h) {
-  const { a: c, b: o, gamma: n = 1, negativeSampleRate: r = 5 } = f, l = e.rows.length, d = new Uint32Array(e.rows), u = new Uint32Array(e.cols), g = new Float32Array(l).fill(0), _ = new Float32Array(l);
-  for (let w = 0; w < l; w++)
-    _[w] = i[w] / r;
+function z(e, n, r, s, t, p, f, l) {
+  const { a: c, b: a, gamma: o = 1, negativeSampleRate: i = 5 } = f, h = n.rows.length, d = new Uint32Array(n.rows), u = new Uint32Array(n.cols), g = new Float32Array(h).fill(0), m = new Float32Array(h);
+  for (let w = 0; w < h; w++)
+    m[w] = r[w] / i;
   for (let w = 0; w < p; w++) {
-    h == null || h(w, p);
-    const y = 1 - w / p;
-    for (let b = 0; b < l; b++) {
-      if (g[b] > w) continue;
-      const M = d[b], x = u[b];
+    l == null || l(w, p);
+    const b = 1 - w / p;
+    for (let y = 0; y < h; y++) {
+      if (g[y] > w) continue;
+      const v = d[y], x = u[y];
       let B = 0;
-      for (let m = 0; m < s; m++) {
-        const F = t[M * s + m] - t[x * s + m];
+      for (let _ = 0; _ < t; _++) {
+        const F = e[v * t + _] - e[x * t + _];
         B += F * F;
       }
-      const G = Math.pow(B, o), R = -2 * c * o * (B > 0 ? G / B : 0) / (c * G + 1);
-      for (let m = 0; m < s; m++) {
-        const F = t[M * s + m] - t[x * s + m], A = P(R * F);
-        t[M * s + m] += y * A;
+      const E = Math.pow(B, a), R = -2 * c * a * (B > 0 ? E / B : 0) / (c * E + 1);
+      for (let _ = 0; _ < t; _++) {
+        const F = e[v * t + _] - e[x * t + _], A = q(R * F);
+        e[v * t + _] += b * A, e[x * t + _] -= b * A;
       }
-      g[b] += i[b];
-      const O = _[b] > 0 ? Math.floor(i[b] / _[b]) : 0;
-      for (let m = 0; m < O; m++) {
-        const F = Math.floor(Math.random() * a);
-        if (F === M) continue;
+      g[y] += r[y];
+      const k = m[y] > 0 ? Math.floor(r[y] / m[y]) : 0;
+      for (let _ = 0; _ < k; _++) {
+        const F = Math.floor(Math.random() * s);
+        if (F === v) continue;
         let A = 0;
-        for (let S = 0; S < s; S++) {
-          const k = t[M * s + S] - t[F * s + S];
-          A += k * k;
+        for (let S = 0; S < t; S++) {
+          const O = e[v * t + S] - e[F * t + S];
+          A += O * O;
         }
-        const v = Math.pow(A, o), N = 2 * n * o / ((1e-3 + A) * (c * v + 1));
-        for (let S = 0; S < s; S++) {
-          const k = t[M * s + S] - t[F * s + S], U = P(N * k);
-          t[M * s + S] += y * U;
+        const M = Math.pow(A, a), N = 2 * o * a / ((1e-3 + A) * (c * M + 1));
+        for (let S = 0; S < t; S++) {
+          const O = e[v * t + S] - e[F * t + S], U = q(N * O);
+          e[v * t + S] += b * U;
         }
       }
-      _[b] += i[b] / r;
+      m[y] += r[y] / i;
     }
   }
-  return t;
+  return e;
 }
-function $(t, e, i, a, s, p, f, h, c, o) {
-  const { a: n, b: r, gamma: l = 1, negativeSampleRate: d = 5 } = c, u = i.rows.length, g = new Uint32Array(i.rows), _ = new Uint32Array(i.cols), w = new Float32Array(u).fill(0), y = new Float32Array(u);
-  for (let b = 0; b < u; b++)
-    y[b] = a[b] / d;
-  for (let b = 0; b < h; b++) {
-    const M = 1 - b / h;
+function $(e, n, r, s, t, p, f, l, c, a) {
+  const { a: o, b: i, gamma: h = 1, negativeSampleRate: d = 5 } = c, u = r.rows.length, g = new Uint32Array(r.rows), m = new Uint32Array(r.cols), w = new Float32Array(u).fill(0), b = new Float32Array(u);
+  for (let y = 0; y < u; y++)
+    b[y] = s[y] / d;
+  for (let y = 0; y < l; y++) {
+    const v = 1 - y / l;
     for (let x = 0; x < u; x++) {
-      if (w[x] > b) continue;
-      const B = g[x], G = _[x];
+      if (w[x] > y) continue;
+      const B = g[x], E = m[x];
       let R = 0;
       for (let A = 0; A < f; A++) {
-        const v = t[B * f + A] - e[G * f + A];
-        R += v * v;
+        const M = e[B * f + A] - n[E * f + A];
+        R += M * M;
       }
-      const O = Math.pow(R, r), m = -2 * n * r * (R > 0 ? O / R : 0) / (n * O + 1);
+      const k = Math.pow(R, i), _ = -2 * o * i * (R > 0 ? k / R : 0) / (o * k + 1);
       for (let A = 0; A < f; A++) {
-        const v = t[B * f + A] - e[G * f + A];
-        t[B * f + A] += M * P(m * v);
+        const M = e[B * f + A] - n[E * f + A];
+        e[B * f + A] += v * q(_ * M);
       }
-      w[x] += a[x];
-      const F = y[x] > 0 ? Math.floor(a[x] / y[x]) : 0;
+      w[x] += s[x];
+      const F = b[x] > 0 ? Math.floor(s[x] / b[x]) : 0;
       for (let A = 0; A < F; A++) {
-        const v = Math.floor(Math.random() * p);
-        if (v === G) continue;
+        const M = Math.floor(Math.random() * p);
+        if (M === E) continue;
         let N = 0;
         for (let U = 0; U < f; U++) {
-          const z = t[B * f + U] - e[v * f + U];
-          N += z * z;
+          const P = e[B * f + U] - n[M * f + U];
+          N += P * P;
         }
-        const S = Math.pow(N, r), k = 2 * l * r / ((1e-3 + N) * (n * S + 1));
+        const S = Math.pow(N, i), O = 2 * h * i / ((1e-3 + N) * (o * S + 1));
         for (let U = 0; U < f; U++) {
-          const z = t[B * f + U] - e[v * f + U];
-          t[B * f + U] += M * P(k * z);
+          const P = e[B * f + U] - n[M * f + U];
+          e[B * f + U] += v * q(O * P);
         }
       }
-      y[x] += a[x] / d;
+      b[x] += s[x] / d;
     }
   }
-  return t;
+  return e;
 }
-function K() {
+function j() {
   return typeof navigator < "u" && !!navigator.gpu;
 }
-async function ae(t, e = {}, i) {
+async function ae(e, n = {}, r) {
   const {
-    nComponents: a = 2,
-    nNeighbors: s = 15,
+    nComponents: s = 2,
+    nNeighbors: t = 15,
     minDist: p = 0.1,
     spread: f = 1,
-    hnsw: h = {}
-  } = e, c = e.nEpochs ?? (t.length > 1e4 ? 200 : 500);
+    hnsw: l = {}
+  } = n, c = n.nEpochs ?? (e.length > 1e4 ? 200 : 500);
   console.time("knn");
-  const { indices: o, distances: n } = await Y(t, s, {
-    M: h.M ?? 16,
-    efConstruction: h.efConstruction ?? 200,
-    efSearch: h.efSearch ?? 50
+  const { indices: a, distances: o } = await Y(e, t, {
+    M: l.M ?? 16,
+    efConstruction: l.efConstruction ?? 200,
+    efSearch: l.efSearch ?? 50
   });
   console.timeEnd("knn"), console.time("fuzzy-set");
-  const r = D(o, n, s);
+  const i = D(a, o, t);
   console.timeEnd("fuzzy-set");
-  const { a: l, b: d } = j(p, f), u = I(r.vals, c), g = t.length, _ = new Float32Array(g * a);
-  for (let y = 0; y < _.length; y++)
-    _[y] = Math.random() * 20 - 10;
+  const { a: h, b: d } = K(p, f), u = I(i.vals, c), g = e.length, m = new Float32Array(g * s);
+  for (let b = 0; b < m.length; b++)
+    m[b] = Math.random() * 20 - 10;
   console.time("sgd");
   let w;
-  if (K())
+  if (j())
     try {
-      const y = new W();
-      await y.init(), w = await y.optimize(
-        _,
-        new Uint32Array(r.rows),
-        new Uint32Array(r.cols),
+      const b = new W();
+      await b.init(), w = await b.optimize(
+        m,
+        new Uint32Array(i.rows),
+        new Uint32Array(i.cols),
         u,
         g,
-        a,
+        s,
         c,
-        { a: l, b: d, gamma: 1, negativeSampleRate: 5 },
-        i
+        { a: h, b: d, gamma: 1, negativeSampleRate: 5 },
+        r
       );
-    } catch (y) {
-      console.warn("WebGPU SGD failed, falling back to CPU:", y), w = q(_, r, u, g, a, c, { a: l, b: d }, i);
+    } catch (b) {
+      console.warn("WebGPU SGD failed, falling back to CPU:", b), w = z(m, i, u, g, s, c, { a: h, b: d }, r);
     }
   else
-    w = q(_, r, u, g, a, c, { a: l, b: d }, i);
+    w = z(m, i, u, g, s, c, { a: h, b: d }, r);
   return console.timeEnd("sgd"), w;
 }
-function j(t, e) {
-  if (Math.abs(e - 1) < 1e-6 && Math.abs(t - 0.1) < 1e-6)
+function K(e, n) {
+  if (Math.abs(n - 1) < 1e-6 && Math.abs(e - 0.1) < 1e-6)
     return { a: 1.9292, b: 0.7915 };
-  if (Math.abs(e - 1) < 1e-6 && Math.abs(t - 0) < 1e-6)
+  if (Math.abs(n - 1) < 1e-6 && Math.abs(e - 0) < 1e-6)
     return { a: 1.8956, b: 0.8006 };
-  if (Math.abs(e - 1) < 1e-6 && Math.abs(t - 0.5) < 1e-6)
+  if (Math.abs(n - 1) < 1e-6 && Math.abs(e - 0.5) < 1e-6)
     return { a: 1.5769, b: 0.8951 };
-  const i = ee(t, e);
-  return { a: te(t, e, i), b: i };
+  const r = ee(e, n);
+  return { a: ne(e, n, r), b: r };
 }
-function ee(t, e) {
-  return 1 / (e * 1.2);
+function ee(e, n) {
+  return 1 / (n * 1.2);
 }
-function te(t, e, i) {
-  return t < 1e-6 ? 1.8956 : (1 / (1 + 1e-3) - 1) / -Math.pow(t, 2 * i);
+function ne(e, n, r) {
+  return e < 1e-6 ? 1.8956 : (1 / (1 + 1e-3) - 1) / -Math.pow(e, 2 * r);
 }
 class re {
-  constructor(e = {}) {
-    E(this, "_nComponents");
-    E(this, "_nNeighbors");
-    E(this, "_minDist");
-    E(this, "_spread");
-    E(this, "_nEpochs");
-    E(this, "_hnswOpts");
-    E(this, "_a");
-    E(this, "_b");
+  constructor(n = {}) {
+    G(this, "_nComponents");
+    G(this, "_nNeighbors");
+    G(this, "_minDist");
+    G(this, "_spread");
+    G(this, "_nEpochs");
+    G(this, "_hnswOpts");
+    G(this, "_a");
+    G(this, "_b");
     /** The low-dimensional embedding produced by the last fit() call. */
-    E(this, "embedding", null);
-    E(this, "_hnswIndex", null);
-    E(this, "_nTrain", 0);
-    this._nComponents = e.nComponents ?? 2, this._nNeighbors = e.nNeighbors ?? 15, this._minDist = e.minDist ?? 0.1, this._spread = e.spread ?? 1, this._nEpochs = e.nEpochs, this._hnswOpts = e.hnsw ?? {};
-    const { a: i, b: a } = j(this._minDist, this._spread);
-    this._a = i, this._b = a;
+    G(this, "embedding", null);
+    G(this, "_hnswIndex", null);
+    G(this, "_nTrain", 0);
+    this._nComponents = n.nComponents ?? 2, this._nNeighbors = n.nNeighbors ?? 15, this._minDist = n.minDist ?? 0.1, this._spread = n.spread ?? 1, this._nEpochs = n.nEpochs, this._hnswOpts = n.hnsw ?? {};
+    const { a: r, b: s } = K(this._minDist, this._spread);
+    this._a = r, this._b = s;
   }
   /**
    * Train UMAP on `vectors`.
@@ -443,45 +445,45 @@ class re {
    * index so that transform() can project new points later.
    * Returns `this` for chaining.
    */
-  async fit(e, i) {
-    const a = e.length, s = this._nEpochs ?? (a > 1e4 ? 200 : 500), { M: p = 16, efConstruction: f = 200, efSearch: h = 50 } = this._hnswOpts;
+  async fit(n, r) {
+    const s = n.length, t = this._nEpochs ?? (s > 1e4 ? 200 : 500), { M: p = 16, efConstruction: f = 200, efSearch: l = 50 } = this._hnswOpts;
     console.time("knn");
-    const { knn: c, index: o } = await J(e, this._nNeighbors, {
+    const { knn: c, index: a } = await Q(n, this._nNeighbors, {
       M: p,
       efConstruction: f,
-      efSearch: h
+      efSearch: l
     });
-    this._hnswIndex = o, this._nTrain = a, console.timeEnd("knn"), console.time("fuzzy-set");
-    const n = D(c.indices, c.distances, this._nNeighbors);
+    this._hnswIndex = a, this._nTrain = s, console.timeEnd("knn"), console.time("fuzzy-set");
+    const o = D(c.indices, c.distances, this._nNeighbors);
     console.timeEnd("fuzzy-set");
-    const r = I(n.vals, s), l = new Float32Array(a * this._nComponents);
-    for (let d = 0; d < l.length; d++)
-      l[d] = Math.random() * 20 - 10;
-    if (console.time("sgd"), K())
+    const i = I(o.vals, t), h = new Float32Array(s * this._nComponents);
+    for (let d = 0; d < h.length; d++)
+      h[d] = Math.random() * 20 - 10;
+    if (console.time("sgd"), j())
       try {
         const d = new W();
         await d.init(), this.embedding = await d.optimize(
-          l,
-          new Uint32Array(n.rows),
-          new Uint32Array(n.cols),
-          r,
-          a,
-          this._nComponents,
+          h,
+          new Uint32Array(o.rows),
+          new Uint32Array(o.cols),
+          i,
           s,
+          this._nComponents,
+          t,
           { a: this._a, b: this._b, gamma: 1, negativeSampleRate: 5 },
-          i
+          r
         );
       } catch (d) {
-        console.warn("WebGPU SGD failed, falling back to CPU:", d), this.embedding = q(l, n, r, a, this._nComponents, s, {
+        console.warn("WebGPU SGD failed, falling back to CPU:", d), this.embedding = z(h, o, i, s, this._nComponents, t, {
           a: this._a,
           b: this._b
-        }, i);
+        }, r);
       }
     else
-      this.embedding = q(l, n, r, a, this._nComponents, s, {
+      this.embedding = z(h, o, i, s, this._nComponents, t, {
         a: this._a,
         b: this._b
-      }, i);
+      }, r);
     return console.timeEnd("sgd"), this;
   }
   /**
@@ -495,35 +497,35 @@ class re {
    *   returned embedding to [0, 1].  The stored training embedding is never
    *   mutated.  Defaults to `false`.
    */
-  async transform(e, i = !1) {
+  async transform(n, r = !1) {
     if (!this._hnswIndex || !this.embedding)
       throw new Error("UMAP.transform() must be called after fit()");
-    const a = e.length, s = this._nEpochs ?? (this._nTrain > 1e4 ? 200 : 500), p = Math.max(100, Math.floor(s / 4)), f = this._hnswIndex.searchKnn(e, this._nNeighbors), h = Q(f.indices, f.distances, this._nNeighbors), c = new Uint32Array(h.rows), o = new Uint32Array(h.cols), n = new Float32Array(a), r = new Float32Array(a * this._nComponents);
+    const s = n.length, t = this._nEpochs ?? (this._nTrain > 1e4 ? 200 : 500), p = Math.max(100, Math.floor(t / 4)), f = this._hnswIndex.searchKnn(n, this._nNeighbors), l = J(f.indices, f.distances, this._nNeighbors), c = new Uint32Array(l.rows), a = new Uint32Array(l.cols), o = new Float32Array(s), i = new Float32Array(s * this._nComponents);
     for (let u = 0; u < c.length; u++) {
-      const g = c[u], _ = o[u], w = h.vals[u];
-      n[g] += w;
-      for (let y = 0; y < this._nComponents; y++)
-        r[g * this._nComponents + y] += w * this.embedding[_ * this._nComponents + y];
+      const g = c[u], m = a[u], w = l.vals[u];
+      o[g] += w;
+      for (let b = 0; b < this._nComponents; b++)
+        i[g * this._nComponents + b] += w * this.embedding[m * this._nComponents + b];
     }
-    for (let u = 0; u < a; u++)
-      if (n[u] > 0)
+    for (let u = 0; u < s; u++)
+      if (o[u] > 0)
         for (let g = 0; g < this._nComponents; g++)
-          r[u * this._nComponents + g] /= n[u];
+          i[u * this._nComponents + g] /= o[u];
       else
         for (let g = 0; g < this._nComponents; g++)
-          r[u * this._nComponents + g] = Math.random() * 20 - 10;
-    const l = I(h.vals, p), d = $(
-      r,
+          i[u * this._nComponents + g] = Math.random() * 20 - 10;
+    const h = I(l.vals, p), d = $(
+      i,
       this.embedding,
-      h,
       l,
-      a,
+      h,
+      s,
       this._nTrain,
       this._nComponents,
       p,
       { a: this._a, b: this._b }
     );
-    return i ? T(d, a, this._nComponents) : d;
+    return r ? T(d, s, this._nComponents) : d;
   }
   /**
    * Convenience method equivalent to `fit(vectors)` followed by
@@ -534,37 +536,37 @@ class re {
    *   returned embedding to [0, 1].  `this.embedding` is never mutated.
    *   Defaults to `false`.
    */
-  async fit_transform(e, i, a = !1) {
-    return await this.fit(e, i), a ? T(this.embedding, e.length, this._nComponents) : this.embedding;
+  async fit_transform(n, r, s = !1) {
+    return await this.fit(n, r), s ? T(this.embedding, n.length, this._nComponents) : this.embedding;
   }
 }
-function T(t, e, i) {
-  const a = new Float32Array(t.length);
-  for (let s = 0; s < i; s++) {
+function T(e, n, r) {
+  const s = new Float32Array(e.length);
+  for (let t = 0; t < r; t++) {
     let p = 1 / 0, f = -1 / 0;
-    for (let c = 0; c < e; c++) {
-      const o = t[c * i + s];
-      o < p && (p = o), o > f && (f = o);
+    for (let c = 0; c < n; c++) {
+      const a = e[c * r + t];
+      a < p && (p = a), a > f && (f = a);
     }
-    const h = f - p;
-    for (let c = 0; c < e; c++)
-      a[c * i + s] = h > 0 ? (t[c * i + s] - p) / h : 0;
+    const l = f - p;
+    for (let c = 0; c < n; c++)
+      s[c * r + t] = l > 0 ? (e[c * r + t] - p) / l : 0;
   }
-  return a;
+  return s;
 }
-function I(t, e) {
-  let i = -1 / 0;
-  for (let s = 0; s < t.length; s++)
-    t[s] > i && (i = t[s]);
-  const a = new Float32Array(t.length);
-  for (let s = 0; s < t.length; s++) {
-    const p = t[s] / i;
-    a[s] = p > 0 ? e / p : -1;
+function I(e, n) {
+  let r = -1 / 0;
+  for (let t = 0; t < e.length; t++)
+    e[t] > r && (r = e[t]);
+  const s = new Float32Array(e.length);
+  for (let t = 0; t < e.length; t++) {
+    const p = e[t] / r;
+    s[t] = p > 0 ? n / p : -1;
   }
-  return a;
+  return s;
 }
 export {
   re as UMAP,
   ae as fit,
-  K as isWebGPUAvailable
+  j as isWebGPUAvailable
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "umap-gpu",
-  "version": "0.2.9",
+  "version": "0.2.11",
   "description": "UMAP with HNSW kNN and WebGPU-accelerated SGD",
   "type": "module",
   "main": "dist/index.js",