umap-gpu 0.2.14 → 0.2.16

package/README.md

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

package/dist/fallback/cpu-sgd.d.ts

@@ -9,7 +9,7 @@ export interface CPUSgdParams {
  * CPU fallback SGD optimizer for environments without WebGPU.
  * Mirrors the GPU shader logic: per-edge attraction + negative-sample repulsion.
  */
-export declare function cpuSgd(embedding: Float32Array, graph: FuzzyGraph, epochsPerSample: Float32Array, nVertices: number, nComponents: number, nEpochs: number, params: CPUSgdParams, onProgress?: (epoch: number, nEpochs: number) => void): Float32Array;
+export declare function cpuSgd(embedding: Float32Array, graph: FuzzyGraph, epochsPerSample: Float32Array, nVertices: number, nComponents: number, nEpochs: number, params: CPUSgdParams, onProgress?: (epoch: number, nEpochs: number) => void, rng?: () => number): Float32Array;
 /**
  * CPU SGD for UMAP.transform(): optimizes only the new-point embeddings.
  * The training embedding is read-only; attraction pulls new points toward
@@ -26,4 +26,4 @@ export declare function cpuSgd(embedding: Float32Array, graph: FuzzyGraph, epoch
  * @param nEpochs        - Number of optimization epochs
  * @param params         - UMAP curve parameters
  */
-export declare function cpuSgdTransform(embeddingNew: Float32Array, embeddingTrain: Float32Array, graph: FuzzyGraph, epochsPerSample: Float32Array, nNew: number, nTrain: number, nComponents: number, nEpochs: number, params: CPUSgdParams, onProgress?: (epoch: number, nEpochs: number) => void): Float32Array;
+export declare function cpuSgdTransform(embeddingNew: Float32Array, embeddingTrain: Float32Array, graph: FuzzyGraph, epochsPerSample: Float32Array, nNew: number, nTrain: number, nComponents: number, nEpochs: number, params: CPUSgdParams, onProgress?: (epoch: number, nEpochs: number) => void, rng?: () => number): Float32Array;

package/dist/gpu/sgd.d.ts

@@ -37,6 +37,6 @@ export declare class GPUSgd {
      * @param params      - UMAP curve parameters and repulsion settings
      * @returns Optimized embedding as Float32Array
      */
-    optimize(embedding: Float32Array, head: Uint32Array, tail: Uint32Array, epochsPerSample: Float32Array, nVertices: number, nComponents: number, nEpochs: number, params: SGDParams, onProgress?: (epoch: number, nEpochs: number) => void): Promise<Float32Array>;
+    optimize(embedding: Float32Array, head: Uint32Array, tail: Uint32Array, epochsPerSample: Float32Array, nVertices: number, nComponents: number, nEpochs: number, params: SGDParams, onProgress?: (epoch: number, nEpochs: number) => void, rng?: () => number): Promise<Float32Array>;
     private makeBuffer;
 }

package/dist/index.js

@@ -1,93 +1,83 @@
-var te = Object.defineProperty;
-var ne = (e, t, f) => t in e ? te(e, t, { enumerable: !0, configurable: !0, writable: !0, value: f }) : e[t] = f;
-var C = (e, t, f) => ne(e, typeof t != "symbol" ? t + "" : t, f);
-import { loadHnswlib as H } from "hnswlib-wasm";
-async function se(e, t, f = {}) {
-  const { M: a = 16, efConstruction: s = 200, efSearch: u = 50 } = f, c = await H(), d = e[0].length, i = e.length, n = new c.HierarchicalNSW("l2", d, "");
-  n.initIndex(i, a, s, 200), n.setEfSearch(Math.max(u, t)), n.addItems(e, !1);
-  const r = [], o = [];
-  for (let h = 0; h < i; h++) {
-    const l = n.searchKnn(e[h], t + 1, void 0), p = l.neighbors.map((g, _) => ({ idx: g, dist: l.distances[_] })).filter(({ idx: g }) => g !== h).slice(0, t);
-    r.push(p.map(({ idx: g }) => g)), o.push(p.map(({ dist: g }) => Math.sqrt(g)));
+var X = Object.defineProperty;
+var Z = (t, e, s) => e in t ? X(t, e, { enumerable: !0, configurable: !0, writable: !0, value: s }) : t[e] = s;
+var G = (t, e, s) => Z(t, typeof e != "symbol" ? e + "" : e, s);
+import { loadHnswlib as $ } from "hnswlib-wasm";
+async function ee(t, e, s = {}) {
+  const { M: r = 16, efConstruction: n = 200, efSearch: u = 50 } = s, o = await $(), f = t[0].length, c = t.length, a = new o.HierarchicalNSW("l2", f, "");
+  a.initIndex(c, r, n, 200), a.setEfSearch(Math.max(u, e)), a.addItems(t, !1);
+  const d = [], i = [];
+  for (let l = 0; l < c; l++) {
+    const g = a.searchKnn(t[l], e + 1, void 0), _ = g.neighbors.map((w, x) => ({ idx: w, dist: g.distances[x] })).filter(({ idx: w }) => w !== l).slice(0, e);
+    d.push(_.map(({ idx: w }) => w)), i.push(_.map(({ dist: w }) => Math.sqrt(w)));
   }
-  return { indices: r, distances: o };
-}
-async function ae(e, t, f = {}) {
-  const { M: a = 16, efConstruction: s = 200, efSearch: u = 50 } = f, c = await H(), d = e[0].length, i = e.length, n = new c.HierarchicalNSW("l2", d, "");
-  n.initIndex(i, a, s, 200), n.setEfSearch(Math.max(u, t)), n.addItems(e, !1);
-  const r = [], o = [];
-  for (let l = 0; l < i; l++) {
-    const p = n.searchKnn(e[l], t + 1, void 0), g = p.neighbors.map((_, m) => ({ idx: _, dist: p.distances[m] })).filter(({ idx: _ }) => _ !== l).slice(0, t);
-    r.push(g.map(({ idx: _ }) => _)), o.push(g.map(({ dist: _ }) => Math.sqrt(_)));
-  }
-  return { knn: { indices: r, distances: o }, index: {
-    searchKnn(l, p) {
-      const g = [], _ = [];
-      for (const m of l) {
-        const y = n.searchKnn(m, p, void 0), b = y.neighbors.map((x, v) => ({ idx: x, dist: y.distances[v] })).sort((x, v) => x.dist - v.dist).slice(0, p);
-        g.push(b.map(({ idx: x }) => x)), _.push(b.map(({ dist: x }) => Math.sqrt(x)));
+  return { knn: { indices: d, distances: i }, index: {
+    searchKnn(l, g) {
+      const _ = [], w = [];
+      for (const x of l) {
+        const m = a.searchKnn(x, g, void 0), F = m.neighbors.map((p, M) => ({ idx: p, dist: m.distances[M] })).sort((p, M) => p.dist - M.dist).slice(0, g);
+        _.push(F.map(({ idx: p }) => p)), w.push(F.map(({ dist: p }) => Math.sqrt(p)));
       }
-      return { indices: g, distances: _ };
+      return { indices: _, distances: w };
     }
   } };
 }
-function V(e, t, f, a = 1) {
-  const s = e.length, { sigmas: u, rhos: c } = Y(t, f), d = [], i = [], n = [];
-  for (let o = 0; o < s; o++)
-    for (let h = 0; h < e[o].length; h++) {
-      const l = t[o][h], p = l <= c[o] ? 1 : Math.exp(-((l - c[o]) / u[o]));
-      d.push(o), i.push(e[o][h]), n.push(p);
+function te(t, e, s, r = 1) {
+  const n = t.length, { sigmas: u, rhos: o } = H(e, s), f = [], c = [], a = [];
+  for (let i = 0; i < n; i++)
+    for (let h = 0; h < t[i].length; h++) {
+      const l = e[i][h], g = l <= o[i] ? 1 : Math.exp(-((l - o[i]) / u[i]));
+      f.push(i), c.push(t[i][h]), a.push(g);
     }
-  return { ...oe(d, i, n, s, a), nVertices: s };
+  return { ...se(f, c, a, n, r), nVertices: n };
 }
-function re(e, t, f) {
-  const a = e.length, { sigmas: s, rhos: u } = Y(t, f), c = [], d = [], i = [];
-  for (let n = 0; n < a; n++)
-    for (let r = 0; r < e[n].length; r++) {
-      const o = t[n][r], h = o <= u[n] ? 1 : Math.exp(-((o - u[n]) / s[n]));
-      c.push(n), d.push(e[n][r]), i.push(h);
+function ne(t, e, s) {
+  const r = t.length, { sigmas: n, rhos: u } = H(e, s), o = [], f = [], c = [];
+  for (let a = 0; a < r; a++)
+    for (let d = 0; d < t[a].length; d++) {
+      const i = e[a][d], h = i <= u[a] ? 1 : Math.exp(-((i - u[a]) / n[a]));
+      o.push(a), f.push(t[a][d]), c.push(h);
     }
   return {
-    rows: new Uint32Array(c),
-    cols: new Uint32Array(d),
-    vals: new Float32Array(i),
-    nVertices: a
+    rows: new Uint32Array(o),
+    cols: new Uint32Array(f),
+    vals: new Float32Array(c),
+    nVertices: r
   };
 }
-function Y(e, t) {
-  const a = e.length, s = new Float32Array(a), u = new Float32Array(a);
-  for (let c = 0; c < a; c++) {
-    const d = e[c];
-    u[c] = d.find((h) => h > 0) ?? 0;
-    let i = 0, n = 1 / 0, r = 1;
-    const o = Math.log2(t);
+function H(t, e) {
+  const r = t.length, n = new Float32Array(r), u = new Float32Array(r);
+  for (let o = 0; o < r; o++) {
+    const f = t[o];
+    u[o] = f.find((h) => h > 0) ?? 0;
+    let c = 0, a = 1 / 0, d = 1;
+    const i = Math.log2(e);
     for (let h = 0; h < 64; h++) {
       let l = 0;
-      for (let p = 0; p < d.length; p++)
-        l += Math.exp(-Math.max(0, d[p] - u[c]) / r);
-      if (Math.abs(l - o) < 1e-5) break;
-      l > o ? (n = r, r = (i + n) / 2) : (i = r, r = n === 1 / 0 ? r * 2 : (i + n) / 2);
+      for (let g = 0; g < f.length; g++)
+        l += Math.exp(-Math.max(0, f[g] - u[o]) / d);
+      if (Math.abs(l - i) < 1e-5) break;
+      l > i ? (a = d, d = (c + a) / 2) : (c = d, d = a === 1 / 0 ? d * 2 : (c + a) / 2);
     }
-    s[c] = r;
+    n[o] = d;
   }
-  return { sigmas: s, rhos: u };
+  return { sigmas: n, rhos: u };
 }
-function oe(e, t, f, a, s) {
+function se(t, e, s, r, n) {
   const u = /* @__PURE__ */ new Map();
-  for (let n = 0; n < e.length; n++)
-    u.set(e[n] * a + t[n], f[n]);
-  const c = [], d = [], i = [];
-  for (const [n, r] of u) {
-    const o = Math.floor(n / a), h = n % a, l = u.get(h * a + o) ?? 0, p = r + l - r * l, g = r * l;
-    c.push(o), d.push(h), i.push(s * p + (1 - s) * g);
+  for (let a = 0; a < t.length; a++)
+    u.set(t[a] * r + e[a], s[a]);
+  const o = [], f = [], c = [];
+  for (const [a, d] of u) {
+    const i = Math.floor(a / r), h = a % r, l = u.get(h * r + i) ?? 0, g = d + l - d * l, _ = d * l;
+    o.push(i), f.push(h), c.push(n * g + (1 - n) * _);
   }
   return {
-    rows: new Uint32Array(c),
-    cols: new Uint32Array(d),
-    vals: new Float32Array(i)
+    rows: new Uint32Array(o),
+    cols: new Uint32Array(f),
+    vals: new Float32Array(c)
   };
 }
-const ie = `// UMAP SGD compute shader — processes one graph edge per GPU thread.
+const ae = `// UMAP SGD compute shader — processes one graph edge per GPU thread.
 // Computes attraction and repulsion forces and accumulates them atomically
 // into a forces buffer. A separate apply-forces pass then updates embeddings,
 // eliminating write-write races on shared vertex positions.
@@ -181,7 +171,10 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     epoch_of_next_negative_sample[edge_idx] += f32(n_neg) * epochs_per_neg;
   }
 
-  var rng = xorshift(rng_seeds[edge_idx] + params.current_epoch * 6364136223u);
+  // 2654435761u is the 32-bit golden-ratio hash constant (0x9E3779B1),
+  // which fits in u32 unlike the 64-bit LCG value 6364136223 used originally.
+  // Bug 14 fix: the original constant exceeded u32 range and failed WGSL validation.
+  var rng = xorshift(rng_seeds[edge_idx] + params.current_epoch * 2654435761u);
 
   for (var s = 0u; s < n_neg; s++) {
     rng = xorshift(rng);
@@ -204,7 +197,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     }
   }
 }
-`, ce = `// Apply-forces shader — second pass of the two-pass GPU SGD.
+`, re = `// Apply-forces shader — second pass of the two-pass GPU SGD.
 //
 // After the SGD pass has atomically accumulated all gradients into the forces
 // buffer, this shader applies each element's accumulated force to the
@@ -233,40 +226,40 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 `;
 let z = null;
-async function Q() {
+async function Y() {
   if (z) return z;
   if (typeof navigator > "u" || !navigator.gpu)
     return null;
-  const e = await navigator.gpu.requestAdapter();
-  return e ? (z = await e.requestDevice(), z.lost.then(() => {
+  const t = await navigator.gpu.requestAdapter();
+  return t ? (z = await t.requestDevice(), z.lost.then(() => {
     z = null;
   }), z) : null;
 }
-function J() {
+function ie() {
   return typeof navigator < "u" && !!navigator.gpu;
 }
-async function he() {
-  return await Q() !== null;
+async function pe() {
+  return await Y() !== null;
 }
-class X {
+class oe {
   constructor() {
-    C(this, "device");
-    C(this, "sgdPipeline");
-    C(this, "applyForcesPipeline");
+    G(this, "device");
+    G(this, "sgdPipeline");
+    G(this, "applyForcesPipeline");
   }
   async init() {
-    const t = await Q();
-    if (!t) throw new Error("WebGPU not supported");
-    this.device = t, this.sgdPipeline = this.device.createComputePipeline({
+    const e = await Y();
+    if (!e) throw new Error("WebGPU not supported");
+    this.device = e, this.sgdPipeline = this.device.createComputePipeline({
       layout: "auto",
       compute: {
-        module: this.device.createShaderModule({ code: ie }),
+        module: this.device.createShaderModule({ code: ae }),
         entryPoint: "main"
       }
     }), this.applyForcesPipeline = this.device.createComputePipeline({
       layout: "auto",
       compute: {
-        module: this.device.createShaderModule({ code: ce }),
+        module: this.device.createShaderModule({ code: re }),
         entryPoint: "main"
       }
     });
@@ -284,254 +277,225 @@ class X {
    * @param params      - UMAP curve parameters and repulsion settings
    * @returns Optimized embedding as Float32Array
    */
-  async optimize(t, f, a, s, u, c, d, i, n) {
-    const { device: r } = this, o = f.length, h = u * c, l = this.makeBuffer(
-      t,
+  async optimize(e, s, r, n, u, o, f, c, a, d = Math.random) {
+    const { device: i } = this, h = s.length, l = u * o, g = this.makeBuffer(
+      e,
       GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
-    ), p = this.makeBuffer(f, GPUBufferUsage.STORAGE), g = this.makeBuffer(a, GPUBufferUsage.STORAGE), _ = this.makeBuffer(s, GPUBufferUsage.STORAGE), m = new Float32Array(s), y = this.makeBuffer(m, GPUBufferUsage.STORAGE), b = new Float32Array(o);
-    for (let A = 0; A < o; A++)
-      b[A] = s[A] / i.negativeSampleRate;
-    const x = this.makeBuffer(b, GPUBufferUsage.STORAGE), v = new Uint32Array(o);
-    for (let A = 0; A < o; A++)
-      v[A] = Math.random() * 4294967295 | 0;
-    const P = this.makeBuffer(v, GPUBufferUsage.STORAGE), G = r.createBuffer({
-      size: h * 4,
+    ), _ = this.makeBuffer(s, GPUBufferUsage.STORAGE), w = this.makeBuffer(r, GPUBufferUsage.STORAGE), x = this.makeBuffer(n, GPUBufferUsage.STORAGE), m = new Float32Array(n), F = this.makeBuffer(m, GPUBufferUsage.STORAGE), p = new Float32Array(h);
+    for (let y = 0; y < h; y++)
+      p[y] = n[y] / c.negativeSampleRate;
+    const M = this.makeBuffer(p, GPUBufferUsage.STORAGE), B = new Uint32Array(h);
+    for (let y = 0; y < h; y++)
+      B[y] = d() * 4294967295 | 0;
+    const S = this.makeBuffer(B, GPUBufferUsage.STORAGE), v = i.createBuffer({
+      size: l * 4,
       usage: GPUBufferUsage.STORAGE,
       mappedAtCreation: !0
     });
-    new Int32Array(G.getMappedRange()).fill(0), G.unmap();
-    const U = r.createBuffer({
+    new Int32Array(v.getMappedRange()).fill(0), v.unmap();
+    const R = i.createBuffer({
       size: 40,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
-    }), F = r.createBuffer({
+    }), O = i.createBuffer({
       size: 16,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
-    }), N = r.createBindGroup({
+    }), C = i.createBindGroup({
       layout: this.sgdPipeline.getBindGroupLayout(0),
       entries: [
-        { binding: 0, resource: { buffer: _ } },
-        { binding: 1, resource: { buffer: p } },
-        { binding: 2, resource: { buffer: g } },
-        { binding: 3, resource: { buffer: l } },
-        { binding: 4, resource: { buffer: y } },
-        { binding: 5, resource: { buffer: x } },
-        { binding: 6, resource: { buffer: U } },
-        { binding: 7, resource: { buffer: P } },
-        { binding: 8, resource: { buffer: G } }
+        { binding: 0, resource: { buffer: x } },
+        { binding: 1, resource: { buffer: _ } },
+        { binding: 2, resource: { buffer: w } },
+        { binding: 3, resource: { buffer: g } },
+        { binding: 4, resource: { buffer: F } },
+        { binding: 5, resource: { buffer: M } },
+        { binding: 6, resource: { buffer: R } },
+        { binding: 7, resource: { buffer: S } },
+        { binding: 8, resource: { buffer: v } }
       ]
-    }), M = r.createBindGroup({
+    }), U = i.createBindGroup({
       layout: this.applyForcesPipeline.getBindGroupLayout(0),
       entries: [
-        { binding: 0, resource: { buffer: l } },
-        { binding: 1, resource: { buffer: G } },
-        { binding: 2, resource: { buffer: F } }
+        { binding: 0, resource: { buffer: g } },
+        { binding: 1, resource: { buffer: v } },
+        { binding: 2, resource: { buffer: O } }
       ]
     });
-    for (let A = 0; A < d; A++) {
-      const B = 1 - A / d, O = new ArrayBuffer(40), S = new Uint32Array(O), k = new Float32Array(O);
-      S[0] = o, S[1] = u, S[2] = c, S[3] = A, S[4] = d, k[5] = B, k[6] = i.a, k[7] = i.b, k[8] = i.gamma, S[9] = i.negativeSampleRate, r.queue.writeBuffer(U, 0, O);
-      const D = new ArrayBuffer(16), $ = new Uint32Array(D), ee = new Float32Array(D);
-      $[0] = h, ee[1] = B, r.queue.writeBuffer(F, 0, D);
-      const T = r.createCommandEncoder(), q = T.beginComputePass();
-      q.setPipeline(this.sgdPipeline), q.setBindGroup(0, N), q.dispatchWorkgroups(Math.ceil(o / 256)), q.end();
-      const I = T.beginComputePass();
-      I.setPipeline(this.applyForcesPipeline), I.setBindGroup(0, M), I.dispatchWorkgroups(Math.ceil(h / 256)), I.end(), r.queue.submit([T.finish()]), A % 10 === 0 && (await r.queue.onSubmittedWorkDone(), n == null || n(A, d));
+    for (let y = 0; y < f; y++) {
+      const N = 1 - y / f, k = new ArrayBuffer(40), P = new Uint32Array(k), q = new Float32Array(k);
+      P[0] = h, P[1] = u, P[2] = o, P[3] = y, P[4] = f, q[5] = N, q[6] = c.a, q[7] = c.b, q[8] = c.gamma, P[9] = c.negativeSampleRate, i.queue.writeBuffer(R, 0, k);
+      const I = new ArrayBuffer(16), Q = new Uint32Array(I), J = new Float32Array(I);
+      Q[0] = l, J[1] = N, i.queue.writeBuffer(O, 0, I);
+      const D = i.createCommandEncoder(), L = D.beginComputePass();
+      L.setPipeline(this.sgdPipeline), L.setBindGroup(0, C), L.dispatchWorkgroups(Math.ceil(h / 256)), L.end();
+      const j = D.beginComputePass();
+      j.setPipeline(this.applyForcesPipeline), j.setBindGroup(0, U), j.dispatchWorkgroups(Math.ceil(l / 256)), j.end(), i.queue.submit([D.finish()]), y % 10 === 0 && (await i.queue.onSubmittedWorkDone(), a == null || a(y, f));
     }
-    const E = r.createBuffer({
-      size: t.byteLength,
+    const A = i.createBuffer({
+      size: e.byteLength,
       usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
-    }), w = r.createCommandEncoder();
-    w.copyBufferToBuffer(l, 0, E, 0, t.byteLength), r.queue.submit([w.finish()]), await E.mapAsync(GPUMapMode.READ);
-    const R = new Float32Array(E.getMappedRange().slice(0));
-    return E.unmap(), l.destroy(), p.destroy(), g.destroy(), _.destroy(), y.destroy(), x.destroy(), P.destroy(), G.destroy(), U.destroy(), F.destroy(), E.destroy(), R;
+    }), b = i.createCommandEncoder();
+    b.copyBufferToBuffer(g, 0, A, 0, e.byteLength), i.queue.submit([b.finish()]), await A.mapAsync(GPUMapMode.READ);
+    const E = new Float32Array(A.getMappedRange().slice(0));
+    return A.unmap(), g.destroy(), _.destroy(), w.destroy(), x.destroy(), F.destroy(), M.destroy(), S.destroy(), v.destroy(), R.destroy(), O.destroy(), A.destroy(), E;
   }
-  makeBuffer(t, f) {
-    const a = this.device.createBuffer({
-      size: t.byteLength,
-      usage: f,
+  makeBuffer(e, s) {
+    const r = this.device.createBuffer({
+      size: e.byteLength,
+      usage: s,
       mappedAtCreation: !0
     });
-    return t instanceof Float32Array ? new Float32Array(a.getMappedRange()).set(t) : new Uint32Array(a.getMappedRange()).set(t), a.unmap(), a;
+    return e instanceof Float32Array ? new Float32Array(r.getMappedRange()).set(e) : new Uint32Array(r.getMappedRange()).set(e), r.unmap(), r;
   }
 }
-function j(e) {
-  return Math.max(-4, Math.min(4, e));
+function T(t) {
+  return Math.max(-4, Math.min(4, t));
 }
-function L(e, t, f, a, s, u, c, d) {
-  const { a: i, b: n, gamma: r = 1, negativeSampleRate: o = 5 } = c, h = t.rows.length, l = new Uint32Array(t.rows), p = new Uint32Array(t.cols), g = new Float32Array(f), _ = new Float32Array(h);
-  for (let m = 0; m < h; m++)
-    _[m] = f[m] / o;
+function W(t, e, s, r, n, u, o, f, c = Math.random) {
+  const { a, b: d, gamma: i = 1, negativeSampleRate: h = 5 } = o, l = e.rows.length, g = new Uint32Array(e.rows), _ = new Uint32Array(e.cols), w = new Float32Array(s), x = new Float32Array(l);
+  for (let m = 0; m < l; m++)
+    x[m] = s[m] / h;
   for (let m = 0; m < u; m++) {
-    d == null || d(m, u);
-    const y = 1 - m / u;
-    for (let b = 0; b < h; b++) {
-      if (g[b] > m) continue;
-      const x = l[b], v = p[b];
-      let P = 0;
-      for (let M = 0; M < s; M++) {
-        const E = e[x * s + M] - e[v * s + M];
-        P += E * E;
+    f == null || f(m, u);
+    const F = 1 - m / u;
+    for (let p = 0; p < l; p++) {
+      if (w[p] > m) continue;
+      const M = g[p], B = _[p];
+      let S = 0;
+      for (let U = 0; U < n; U++) {
+        const A = t[M * n + U] - t[B * n + U];
+        S += A * A;
       }
-      const G = Math.pow(P, n), U = -2 * i * n * (P > 0 ? G / P : 0) / (i * G + 1);
-      for (let M = 0; M < s; M++) {
-        const E = e[x * s + M] - e[v * s + M], w = j(U * E);
-        e[x * s + M] += y * w, e[v * s + M] -= y * w;
+      const v = Math.pow(S, d), R = -2 * a * d * (S > 0 ? v / S : 0) / (a * v + 1);
+      for (let U = 0; U < n; U++) {
+        const A = t[M * n + U] - t[B * n + U], b = T(R * A);
+        t[M * n + U] += F * b, t[B * n + U] -= F * b;
       }
-      g[b] += f[b];
-      const F = f[b] / o, N = Math.max(0, Math.floor(
-        (m - _[b]) / F
+      w[p] += s[p];
+      const O = s[p] / h, C = Math.max(0, Math.floor(
+        (m - x[p]) / O
       ));
-      _[b] += N * F;
-      for (let M = 0; M < N; M++) {
-        const E = Math.floor(Math.random() * a);
-        if (E === x) continue;
-        let w = 0;
-        for (let B = 0; B < s; B++) {
-          const O = e[x * s + B] - e[E * s + B];
-          w += O * O;
+      x[p] += C * O;
+      for (let U = 0; U < C; U++) {
+        const A = Math.floor(c() * r);
+        if (A === M) continue;
+        let b = 0;
+        for (let N = 0; N < n; N++) {
+          const k = t[M * n + N] - t[A * n + N];
+          b += k * k;
         }
-        const R = Math.pow(w, n), A = 2 * r * n / ((1e-3 + w) * (i * R + 1));
-        for (let B = 0; B < s; B++) {
-          const O = e[x * s + B] - e[E * s + B], S = j(A * O);
-          e[x * s + B] += y * S;
+        const E = Math.pow(b, d), y = 2 * i * d / ((1e-3 + b) * (a * E + 1));
+        for (let N = 0; N < n; N++) {
+          const k = t[M * n + N] - t[A * n + N], P = T(y * k);
+          t[M * n + N] += F * P;
         }
       }
     }
   }
-  return e;
+  return t;
 }
-function fe(e, t, f, a, s, u, c, d, i, n) {
-  const { a: r, b: o, gamma: h = 1, negativeSampleRate: l = 5 } = i, p = f.rows.length, g = new Uint32Array(f.rows), _ = new Uint32Array(f.cols), m = new Float32Array(a), y = new Float32Array(p);
-  for (let b = 0; b < p; b++)
-    y[b] = a[b] / l;
-  for (let b = 0; b < d; b++) {
-    const x = 1 - b / d;
-    for (let v = 0; v < p; v++) {
-      if (m[v] > b) continue;
-      const P = g[v], G = _[v];
-      let U = 0;
-      for (let w = 0; w < c; w++) {
-        const R = e[P * c + w] - t[G * c + w];
-        U += R * R;
+function ce(t, e, s, r, n, u, o, f, c, a, d = Math.random) {
+  const { a: i, b: h, gamma: l = 1, negativeSampleRate: g = 5 } = c, _ = s.rows.length, w = new Uint32Array(s.rows), x = new Uint32Array(s.cols), m = new Float32Array(r), F = new Float32Array(_);
+  for (let p = 0; p < _; p++)
+    F[p] = r[p] / g;
+  for (let p = 0; p < f; p++) {
+    const M = 1 - p / f;
+    for (let B = 0; B < _; B++) {
+      if (m[B] > p) continue;
+      const S = w[B], v = x[B];
+      let R = 0;
+      for (let b = 0; b < o; b++) {
+        const E = t[S * o + b] - e[v * o + b];
+        R += E * E;
       }
-      const F = Math.pow(U, o), N = -2 * r * o * (U > 0 ? F / U : 0) / (r * F + 1);
-      for (let w = 0; w < c; w++) {
-        const R = e[P * c + w] - t[G * c + w];
-        e[P * c + w] += x * j(N * R);
+      const O = Math.pow(R, h), C = -2 * i * h * (R > 0 ? O / R : 0) / (i * O + 1);
+      for (let b = 0; b < o; b++) {
+        const E = t[S * o + b] - e[v * o + b];
+        t[S * o + b] += M * T(C * E);
       }
-      m[v] += a[v];
-      const M = a[v] / l, E = Math.max(0, Math.floor(
-        (b - y[v]) / M
+      m[B] += r[B];
+      const U = r[B] / g, A = Math.max(0, Math.floor(
+        (p - F[B]) / U
       ));
-      y[v] += E * M;
-      for (let w = 0; w < E; w++) {
-        const R = Math.floor(Math.random() * u);
-        if (R === G) continue;
-        let A = 0;
-        for (let S = 0; S < c; S++) {
-          const k = e[P * c + S] - t[R * c + S];
-          A += k * k;
+      F[B] += A * U;
+      for (let b = 0; b < A; b++) {
+        const E = Math.floor(d() * u);
+        if (E === v) continue;
+        let y = 0;
+        for (let P = 0; P < o; P++) {
+          const q = t[S * o + P] - e[E * o + P];
+          y += q * q;
         }
-        const B = Math.pow(A, o), O = 2 * h * o / ((1e-3 + A) * (r * B + 1));
-        for (let S = 0; S < c; S++) {
-          const k = e[P * c + S] - t[R * c + S];
-          e[P * c + S] += x * j(O * k);
+        const N = Math.pow(y, h), k = 2 * l * h / ((1e-3 + y) * (i * N + 1));
+        for (let P = 0; P < o; P++) {
+          const q = t[S * o + P] - e[E * o + P];
+          t[S * o + P] += M * T(k * q);
         }
       }
     }
   }
-  return e;
+  return t;
 }
-async function pe(e, t = {}, f) {
-  const {
-    nComponents: a = 2,
-    nNeighbors: s = 15,
-    minDist: u = 0.1,
-    spread: c = 1,
-    hnsw: d = {}
-  } = t, i = t.nEpochs ?? (e.length > 1e4 ? 200 : 500);
-  console.time("knn");
-  const { indices: n, distances: r } = await se(e, s, {
-    M: d.M ?? 16,
-    efConstruction: d.efConstruction ?? 200,
-    efSearch: d.efSearch ?? 50
-  });
-  console.timeEnd("knn"), console.time("fuzzy-set");
-  const o = V(n, r, s);
-  console.timeEnd("fuzzy-set");
-  const { a: h, b: l } = Z(u, c), p = W(o.vals), g = e.length, _ = new Float32Array(g * a);
-  for (let y = 0; y < _.length; y++)
-    _[y] = Math.random() * 20 - 10;
-  console.time("sgd");
-  let m;
-  if (J())
-    try {
-      const y = new X();
-      await y.init(), m = await y.optimize(
-        _,
-        new Uint32Array(o.rows),
-        new Uint32Array(o.cols),
-        p,
-        g,
-        a,
-        i,
-        { a: h, b: l, gamma: 1, negativeSampleRate: 5 },
-        f
-      );
-    } catch (y) {
-      console.warn("WebGPU SGD failed, falling back to CPU:", y), m = L(_, o, p, g, a, i, { a: h, b: l }, f);
-    }
-  else
-    m = L(_, o, p, g, a, i, { a: h, b: l }, f);
-  return console.timeEnd("sgd"), m;
+function fe(t) {
+  if (t === void 0) return Math.random;
+  let e = t | 0;
+  return () => {
+    e = e + 1831565813 | 0;
+    let s = Math.imul(e ^ e >>> 15, 1 | e);
+    return s = s + Math.imul(s ^ s >>> 7, 61 | s) ^ s, ((s ^ s >>> 14) >>> 0) / 4294967296;
+  };
 }
-function Z(e, t) {
-  return Math.abs(t - 1) < 1e-6 && Math.abs(e - 0.1) < 1e-6 ? { a: 1.9292, b: 0.7915 } : Math.abs(t - 1) < 1e-6 && Math.abs(e - 0) < 1e-6 ? { a: 1.8956, b: 0.8006 } : Math.abs(t - 1) < 1e-6 && Math.abs(e - 0.5) < 1e-6 ? { a: 1.5769, b: 0.8951 } : de(e, t);
+async function _e(t, e = {}, s) {
+  return new ue(e).fit_transform(t, s);
 }
-function de(e, t) {
-  const a = [], s = [];
-  for (let i = 0; i < 299; i++) {
-    const n = (i + 1) / 299 * t * 3;
-    a.push(n), s.push(n < e ? 1 : Math.exp(-(n - e) / t));
+function de(t, e) {
+  return Math.abs(e - 1) < 1e-6 && Math.abs(t - 0.1) < 1e-6 ? { a: 1.9292, b: 0.7915 } : Math.abs(e - 1) < 1e-6 && Math.abs(t - 0) < 1e-6 ? { a: 1.8956, b: 0.8006 } : Math.abs(e - 1) < 1e-6 && Math.abs(t - 0.5) < 1e-6 ? { a: 1.5769, b: 0.8951 } : le(t, e);
+}
+function le(t, e) {
+  const r = [], n = [];
+  for (let c = 0; c < 299; c++) {
+    const a = (c + 1) / 299 * e * 3;
+    r.push(a), n.push(a < t ? 1 : Math.exp(-(a - t) / e));
   }
-  let u = 1, c = 1, d = 1e-3;
-  for (let i = 0; i < 500; i++) {
-    let n = 0, r = 0, o = 0, h = 0, l = 0, p = 0;
-    for (let U = 0; U < 299; U++) {
-      const F = a[U], N = Math.pow(F, 2 * c), M = 1 + u * N, w = 1 / M - s[U];
-      p += w * w;
-      const R = M * M, A = -N / R, B = F > 0 ? -2 * Math.log(F) * u * N / R : 0;
-      n += A * w, r += B * w, o += A * A, h += B * B, l += A * B;
+  let u = 1, o = 1, f = 1e-3;
+  for (let c = 0; c < 500; c++) {
+    let a = 0, d = 0, i = 0, h = 0, l = 0, g = 0;
+    for (let v = 0; v < 299; v++) {
+      const R = r[v], O = Math.pow(R, 2 * o), C = 1 + u * O, A = 1 / C - n[v];
+      g += A * A;
+      const b = C * C, E = -O / b, y = R > 0 ? -2 * Math.log(R) * u * O / b : 0;
+      a += E * A, d += y * A, i += E * E, h += y * y, l += E * y;
     }
-    const g = o + d, _ = h + d, m = l, y = g * _ - m * m;
-    if (Math.abs(y) < 1e-20) break;
-    const b = -(_ * n - m * r) / y, x = -(g * r - m * n) / y, v = Math.max(1e-4, u + b), P = Math.max(1e-4, c + x);
-    let G = 0;
-    for (let U = 0; U < 299; U++) {
-      const F = Math.pow(a[U], 2 * P), N = 1 / (1 + v * F) - s[U];
-      G += N * N;
+    const _ = i + f, w = h + f, x = l, m = _ * w - x * x;
+    if (Math.abs(m) < 1e-20) break;
+    const F = -(w * a - x * d) / m, p = -(_ * d - x * a) / m, M = Math.max(1e-4, u + F), B = Math.max(1e-4, o + p);
+    let S = 0;
+    for (let v = 0; v < 299; v++) {
+      const R = Math.pow(r[v], 2 * B), O = 1 / (1 + M * R) - n[v];
+      S += O * O;
     }
-    if (G < p ? (u = v, c = P, d = Math.max(1e-10, d / 10)) : d = Math.min(1e10, d * 10), Math.abs(b) < 1e-8 && Math.abs(x) < 1e-8) break;
+    if (S < g ? (u = M, o = B, f = Math.max(1e-10, f / 10)) : f = Math.min(1e10, f * 10), Math.abs(F) < 1e-8 && Math.abs(p) < 1e-8) break;
   }
-  return { a: u, b: c };
+  return { a: u, b: o };
 }
-class ge {
-  constructor(t = {}) {
-    C(this, "_nComponents");
-    C(this, "_nNeighbors");
-    C(this, "_minDist");
-    C(this, "_spread");
-    C(this, "_nEpochs");
-    C(this, "_hnswOpts");
-    C(this, "_a");
-    C(this, "_b");
+class ue {
+  constructor(e = {}) {
+    G(this, "_nComponents");
+    G(this, "_nNeighbors");
+    G(this, "_minDist");
+    G(this, "_spread");
+    G(this, "_nEpochs");
+    G(this, "_hnswOpts");
+    G(this, "_a");
+    G(this, "_b");
+    G(this, "_debug");
+    G(this, "_rng");
     /** The low-dimensional embedding produced by the last fit() call. */
-    C(this, "embedding", null);
-    C(this, "_hnswIndex", null);
-    C(this, "_nTrain", 0);
-    this._nComponents = t.nComponents ?? 2, this._nNeighbors = t.nNeighbors ?? 15, this._minDist = t.minDist ?? 0.1, this._spread = t.spread ?? 1, this._nEpochs = t.nEpochs, this._hnswOpts = t.hnsw ?? {};
-    const { a: f, b: a } = Z(this._minDist, this._spread);
-    this._a = f, this._b = a;
+    G(this, "embedding", null);
+    G(this, "_hnswIndex", null);
+    G(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 ?? {}, this._debug = e.debug ?? !1, this._rng = fe(e.seed);
+    const { a: s, b: r } = de(this._minDist, this._spread);
+    this._a = s, this._b = r;
   }
   /**
    * Train UMAP on `vectors`.
@@ -539,46 +503,47 @@ class ge {
    * index so that transform() can project new points later.
    * Returns `this` for chaining.
    */
-  async fit(t, f) {
-    const a = t.length, s = this._nEpochs ?? (a > 1e4 ? 200 : 500), { M: u = 16, efConstruction: c = 200, efSearch: d = 50 } = this._hnswOpts;
-    console.time("knn");
-    const { knn: i, index: n } = await ae(t, this._nNeighbors, {
+  async fit(e, s) {
+    const r = e.length, n = this._nEpochs ?? (r > 1e4 ? 200 : 500), { M: u = 16, efConstruction: o = 200, efSearch: f = 50 } = this._hnswOpts;
+    this._debug && console.time("knn");
+    const { knn: c, index: a } = await ee(e, this._nNeighbors, {
       M: u,
-      efConstruction: c,
-      efSearch: d
+      efConstruction: o,
+      efSearch: f
     });
-    this._hnswIndex = n, this._nTrain = a, console.timeEnd("knn"), console.time("fuzzy-set");
-    const r = V(i.indices, i.distances, this._nNeighbors);
-    console.timeEnd("fuzzy-set");
-    const o = W(r.vals), h = new Float32Array(a * this._nComponents);
+    this._hnswIndex = a, this._nTrain = r, this._debug && console.timeEnd("knn"), this._debug && console.time("fuzzy-set");
+    const d = te(c.indices, c.distances, this._nNeighbors);
+    this._debug && console.timeEnd("fuzzy-set");
+    const i = V(d.vals), h = new Float32Array(r * this._nComponents);
     for (let l = 0; l < h.length; l++)
-      h[l] = Math.random() * 20 - 10;
-    if (console.time("sgd"), J())
+      h[l] = this._rng() * 20 - 10;
+    if (this._debug && console.time("sgd"), ie())
       try {
-        const l = new X();
+        const l = new oe();
         await l.init(), this.embedding = await l.optimize(
           h,
-          new Uint32Array(r.rows),
-          new Uint32Array(r.cols),
-          o,
-          a,
+          new Uint32Array(d.rows),
+          new Uint32Array(d.cols),
+          i,
+          r,
           this._nComponents,
-          s,
+          n,
           { a: this._a, b: this._b, gamma: 1, negativeSampleRate: 5 },
-          f
+          s,
+          this._rng
         );
       } catch (l) {
-        console.warn("WebGPU SGD failed, falling back to CPU:", l), this.embedding = L(h, r, o, a, this._nComponents, s, {
+        console.warn("WebGPU SGD failed, falling back to CPU:", l), this.embedding = W(h, d, i, r, this._nComponents, n, {
           a: this._a,
           b: this._b
-        }, f);
+        }, s, this._rng);
       }
     else
-      this.embedding = L(h, r, o, a, this._nComponents, s, {
+      this.embedding = W(h, d, i, r, this._nComponents, n, {
         a: this._a,
         b: this._b
-      }, f);
-    return console.timeEnd("sgd"), this;
+      }, s, this._rng);
+    return this._debug && console.timeEnd("sgd"), this;
   }
   /**
    * Project new (unseen) `vectors` into the embedding space learned by fit().
@@ -591,35 +556,37 @@ class ge {
    *   returned embedding to [0, 1].  The stored training embedding is never
    *   mutated.  Defaults to `false`.
    */
-  async transform(t, f = !1) {
+  async transform(e, s = !1) {
     if (!this._hnswIndex || !this.embedding)
       throw new Error("UMAP.transform() must be called after fit()");
-    const a = t.length, s = this._nEpochs ?? (this._nTrain > 1e4 ? 200 : 500), u = Math.max(100, Math.floor(s / 4)), c = this._hnswIndex.searchKnn(t, this._nNeighbors), d = re(c.indices, c.distances, this._nNeighbors), i = new Uint32Array(d.rows), n = new Uint32Array(d.cols), r = new Float32Array(a), o = new Float32Array(a * this._nComponents);
-    for (let p = 0; p < i.length; p++) {
-      const g = i[p], _ = n[p], m = d.vals[p];
-      r[g] += m;
-      for (let y = 0; y < this._nComponents; y++)
-        o[g * this._nComponents + y] += m * this.embedding[_ * this._nComponents + y];
+    const r = e.length, n = this._nEpochs ?? (this._nTrain > 1e4 ? 200 : 500), u = Math.max(100, Math.floor(n / 4)), o = this._hnswIndex.searchKnn(e, this._nNeighbors), f = ne(o.indices, o.distances, this._nNeighbors), c = new Uint32Array(f.rows), a = new Uint32Array(f.cols), d = new Float32Array(r), i = new Float32Array(r * this._nComponents);
+    for (let g = 0; g < c.length; g++) {
+      const _ = c[g], w = a[g], x = f.vals[g];
+      d[_] += x;
+      for (let m = 0; m < this._nComponents; m++)
+        i[_ * this._nComponents + m] += x * this.embedding[w * this._nComponents + m];
     }
-    for (let p = 0; p < a; p++)
-      if (r[p] > 0)
-        for (let g = 0; g < this._nComponents; g++)
-          o[p * this._nComponents + g] /= r[p];
+    for (let g = 0; g < r; g++)
+      if (d[g] > 0)
+        for (let _ = 0; _ < this._nComponents; _++)
+          i[g * this._nComponents + _] /= d[g];
       else
-        for (let g = 0; g < this._nComponents; g++)
-          o[p * this._nComponents + g] = Math.random() * 20 - 10;
-    const h = W(d.vals), l = fe(
-      o,
+        for (let _ = 0; _ < this._nComponents; _++)
+          i[g * this._nComponents + _] = this._rng() * 20 - 10;
+    const h = V(f.vals), l = ce(
+      i,
       this.embedding,
-      d,
+      f,
       h,
-      a,
+      r,
       this._nTrain,
       this._nComponents,
       u,
-      { a: this._a, b: this._b }
+      { a: this._a, b: this._b },
+      void 0,
+      this._rng
     );
-    return f ? K(l, a, this._nComponents) : l;
+    return s ? K(l, r, this._nComponents) : l;
   }
   /**
    * Convenience method equivalent to `fit(vectors)` followed by
@@ -630,38 +597,38 @@ class ge {
    *   returned embedding to [0, 1].  `this.embedding` is never mutated.
    *   Defaults to `false`.
    */
-  async fit_transform(t, f, a = !1) {
-    return await this.fit(t, f), a ? K(this.embedding, t.length, this._nComponents) : this.embedding;
+  async fit_transform(e, s, r = !1) {
+    return await this.fit(e, s), r ? K(this.embedding, e.length, this._nComponents) : this.embedding;
   }
 }
-function K(e, t, f) {
-  const a = new Float32Array(e.length);
-  for (let s = 0; s < f; s++) {
-    let u = 1 / 0, c = -1 / 0;
-    for (let i = 0; i < t; i++) {
-      const n = e[i * f + s];
-      n < u && (u = n), n > c && (c = n);
+function K(t, e, s) {
+  const r = new Float32Array(t.length);
+  for (let n = 0; n < s; n++) {
+    let u = 1 / 0, o = -1 / 0;
+    for (let c = 0; c < e; c++) {
+      const a = t[c * s + n];
+      a < u && (u = a), a > o && (o = a);
     }
-    const d = c - u;
-    for (let i = 0; i < t; i++)
-      a[i * f + s] = d > 0 ? (e[i * f + s] - u) / d : 0;
+    const f = o - u;
+    for (let c = 0; c < e; c++)
+      r[c * s + n] = f > 0 ? (t[c * s + n] - u) / f : 0;
   }
-  return a;
+  return r;
 }
-function W(e, t) {
-  let f = -1 / 0;
-  for (let s = 0; s < e.length; s++)
-    e[s] > f && (f = e[s]);
-  const a = new Float32Array(e.length);
-  for (let s = 0; s < e.length; s++) {
-    const u = e[s] / f;
-    a[s] = u > 0 ? 1 / u : -1;
+function V(t, e) {
+  let s = -1 / 0;
+  for (let n = 0; n < t.length; n++)
+    t[n] > s && (s = t[n]);
+  const r = new Float32Array(t.length);
+  for (let n = 0; n < t.length; n++) {
+    const u = t[n] / s;
+    r[n] = u > 0 ? 1 / u : -1;
   }
-  return a;
+  return r;
 }
 export {
-  ge as UMAP,
-  he as checkWebGPUAvailable,
-  pe as fit,
-  J as isWebGPUAvailable
+  ue as UMAP,
+  pe as checkWebGPUAvailable,
+  _e as fit,
+  ie as isWebGPUAvailable
 };

package/dist/rng.d.ts

@@ -0,0 +1,2 @@
+/** Returns a seeded PRNG if seed is provided, otherwise wraps Math.random. */
+export declare function makeRng(seed?: number): () => number;

package/dist/umap.d.ts

@@ -15,6 +15,10 @@ export interface UMAPOptions {
         efConstruction?: number;
         efSearch?: number;
     };
+    /** Enable timing instrumentation via console.time/timeEnd (default: false) */
+    debug?: boolean;
+    /** Random seed for reproducible results. When omitted, Math.random() is used. */
+    seed?: number;
 }
 /**
  * Called after each completed SGD epoch (or every 10 epochs on the GPU path,
@@ -71,6 +75,8 @@ export declare class UMAP {
     private readonly _hnswOpts;
     private readonly _a;
     private readonly _b;
+    private readonly _debug;
+    private readonly _rng;
     /** The low-dimensional embedding produced by the last fit() call. */
     embedding: Float32Array | null;
     private _hnswIndex;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "umap-gpu",
-  "version": "0.2.14",
+  "version": "0.2.16",
   "description": "UMAP with HNSW kNN and WebGPU-accelerated SGD",
   "type": "module",
   "main": "dist/index.js",
@@ -40,6 +40,7 @@
     "build": "vite build && tsc",
     "dev": "vite",
     "test": "vitest run",
+    "bench": "npx vitest run --config vitest.bench.config.ts",
     "prepublishOnly": "bun test && bun run build",
     "docs:dev": "vitepress dev docs",
     "docs:build": "vitepress build docs",
@@ -55,7 +56,8 @@
     "vite": "^5.0.0",
     "vitepress": "^1.6.4",
     "vitepress-plugin-llms": "^1.11.0",
-    "vitest": "^4.0.18"
+    "vitest": "^4.0.18",
+    "webgpu": "^0.3.8"
   },
   "license": "MIT"
 }