umap-gpu 0.2.10 → 0.2.13

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,93 +1,93 @@
-var H = Object.defineProperty;
-var V = (e, t, r) => t in e ? H(e, t, { enumerable: !0, configurable: !0, writable: !0, value: r }) : e[t] = r;
-var G = (e, t, r) => V(e, typeof t != "symbol" ? t + "" : t, r);
-import { loadHnswlib as C } from "hnswlib-wasm";
-async function Y(e, t, r = {}) {
-  const { M: s = 16, efConstruction: n = 200, efSearch: p = 50 } = r, f = await C(), d = e[0].length, c = e.length, a = new f.HierarchicalNSW("l2", d, "");
-  a.initIndex(c, s, n, 200), a.setEfSearch(Math.max(p, t)), a.addItems(e, !1);
+var V = Object.defineProperty;
+var Y = (e, t, r) => t in e ? V(e, t, { enumerable: !0, configurable: !0, writable: !0, value: r }) : e[t] = r;
+var G = (e, t, r) => Y(e, typeof t != "symbol" ? t + "" : t, r);
+import { loadHnswlib as D } from "hnswlib-wasm";
+async function Q(e, t, r = {}) {
+  const { M: s = 16, efConstruction: n = 200, efSearch: g = 50 } = r, c = await D(), d = e[0].length, f = e.length, a = new c.HierarchicalNSW("l2", d, "");
+  a.initIndex(f, s, n, 200), a.setEfSearch(Math.max(g, t)), a.addItems(e, !1);
   const o = [], i = [];
-  for (let h = 0; h < c; h++) {
-    const l = a.searchKnn(e[h], t + 1, void 0), u = l.neighbors.map((g, m) => ({ idx: g, dist: l.distances[m] })).filter(({ idx: g }) => g !== h).slice(0, t);
-    o.push(u.map(({ idx: g }) => g)), i.push(u.map(({ dist: g }) => g));
+  for (let h = 0; h < f; h++) {
+    const l = a.searchKnn(e[h], t + 1, void 0), u = l.neighbors.map((p, m) => ({ idx: p, dist: l.distances[m] })).filter(({ idx: p }) => p !== h).slice(0, t);
+    o.push(u.map(({ idx: p }) => p)), i.push(u.map(({ dist: p }) => p));
   }
   return { indices: o, distances: i };
 }
-async function Q(e, t, r = {}) {
-  const { M: s = 16, efConstruction: n = 200, efSearch: p = 50 } = r, f = await C(), d = e[0].length, c = e.length, a = new f.HierarchicalNSW("l2", d, "");
-  a.initIndex(c, s, n, 200), a.setEfSearch(Math.max(p, t)), a.addItems(e, !1);
+async function J(e, t, r = {}) {
+  const { M: s = 16, efConstruction: n = 200, efSearch: g = 50 } = r, c = await D(), d = e[0].length, f = e.length, a = new c.HierarchicalNSW("l2", d, "");
+  a.initIndex(f, s, n, 200), a.setEfSearch(Math.max(g, t)), a.addItems(e, !1);
   const o = [], i = [];
-  for (let l = 0; l < c; l++) {
-    const u = a.searchKnn(e[l], t + 1, void 0), g = u.neighbors.map((m, w) => ({ idx: m, dist: u.distances[w] })).filter(({ idx: m }) => m !== l).slice(0, t);
-    o.push(g.map(({ idx: m }) => m)), i.push(g.map(({ dist: m }) => m));
+  for (let l = 0; l < f; l++) {
+    const u = a.searchKnn(e[l], t + 1, void 0), p = u.neighbors.map((m, _) => ({ idx: m, dist: u.distances[_] })).filter(({ idx: m }) => m !== l).slice(0, t);
+    o.push(p.map(({ idx: m }) => m)), i.push(p.map(({ dist: m }) => m));
   }
   return { knn: { indices: o, distances: i }, index: {
     searchKnn(l, u) {
-      const g = [], m = [];
-      for (const w of l) {
-        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));
+      const p = [], m = [];
+      for (const _ of l) {
+        const w = a.searchKnn(_, u, void 0), b = w.neighbors.map((A, x) => ({ idx: A, dist: w.distances[x] })).sort((A, x) => A.dist - x.dist).slice(0, u);
+        p.push(b.map(({ idx: A }) => A)), m.push(b.map(({ dist: A }) => A));
       }
-      return { indices: g, distances: m };
+      return { indices: p, distances: m };
     }
   } };
 }
-function D(e, t, r, s = 1) {
-  const n = e.length, { sigmas: p, rhos: f } = L(t, r), d = [], c = [], a = [];
+function L(e, t, r, s = 1) {
+  const n = e.length, { sigmas: g, rhos: c } = W(t, r), d = [], f = [], a = [];
   for (let i = 0; i < n; i++)
     for (let h = 0; h < e[i].length; h++) {
-      const l = t[i][h], u = l <= f[i] ? 1 : Math.exp(-((l - f[i]) / p[i]));
-      d.push(i), c.push(e[i][h]), a.push(u);
+      const l = t[i][h], u = l <= c[i] ? 1 : Math.exp(-((l - c[i]) / g[i]));
+      d.push(i), f.push(e[i][h]), a.push(u);
     }
-  return { ...X(d, c, a, n, s), nVertices: n };
+  return { ...Z(d, f, a, n, s), nVertices: n };
 }
-function J(e, t, r) {
-  const s = e.length, { sigmas: n, rhos: p } = L(t, r), f = [], d = [], c = [];
+function X(e, t, r) {
+  const s = e.length, { sigmas: n, rhos: g } = W(t, r), c = [], d = [], f = [];
   for (let a = 0; a < s; a++)
     for (let o = 0; o < e[a].length; o++) {
-      const i = t[a][o], h = i <= p[a] ? 1 : Math.exp(-((i - p[a]) / n[a]));
-      f.push(a), d.push(e[a][o]), c.push(h);
+      const i = t[a][o], h = i <= g[a] ? 1 : Math.exp(-((i - g[a]) / n[a]));
+      c.push(a), d.push(e[a][o]), f.push(h);
     }
   return {
-    rows: new Float32Array(f),
+    rows: new Float32Array(c),
     cols: new Float32Array(d),
-    vals: new Float32Array(c),
+    vals: new Float32Array(f),
     nVertices: s
   };
 }
-function L(e, t) {
-  const s = e.length, n = new Float32Array(s), p = new Float32Array(s);
-  for (let f = 0; f < s; f++) {
-    const d = e[f];
-    p[f] = d.find((h) => h > 0) ?? 0;
-    let c = 0, a = 1 / 0, o = 1;
+function W(e, t) {
+  const s = e.length, n = new Float32Array(s), g = new Float32Array(s);
+  for (let c = 0; c < s; c++) {
+    const d = e[c];
+    g[c] = d.find((h) => h > 0) ?? 0;
+    let f = 0, a = 1 / 0, o = 1;
     const i = Math.log2(t);
     for (let h = 0; h < 64; h++) {
       let l = 0;
-      for (let u = 1; u < d.length; u++)
-        l += Math.exp(-Math.max(0, d[u] - p[f]) / o);
+      for (let u = 0; u < d.length; u++)
+        l += Math.exp(-Math.max(0, d[u] - g[c]) / o);
       if (Math.abs(l - i) < 1e-5) break;
-      l > i ? (a = o, o = (c + a) / 2) : (c = o, o = a === 1 / 0 ? o * 2 : (c + a) / 2);
+      l > i ? (a = o, o = (f + a) / 2) : (f = o, o = a === 1 / 0 ? o * 2 : (f + a) / 2);
     }
-    n[f] = o;
+    n[c] = o;
   }
-  return { sigmas: n, rhos: p };
+  return { sigmas: n, rhos: g };
 }
-function X(e, t, r, s, n) {
-  const p = /* @__PURE__ */ new Map();
+function Z(e, t, r, s, n) {
+  const g = /* @__PURE__ */ new Map();
   for (let a = 0; a < e.length; a++)
-    p.set(e[a] * s + t[a], r[a]);
-  const f = [], d = [], c = [];
-  for (const [a, o] of p) {
-    const i = Math.floor(a / s), h = a % s, l = p.get(h * s + i) ?? 0, u = o + l - o * l, g = o * l;
-    f.push(i), d.push(h), c.push(n * u + (1 - n) * g);
+    g.set(e[a] * s + t[a], r[a]);
+  const c = [], d = [], f = [];
+  for (const [a, o] of g) {
+    const i = Math.floor(a / s), h = a % s, l = g.get(h * s + i) ?? 0, u = o + l - o * l, p = o * l;
+    c.push(i), d.push(h), f.push(n * u + (1 - n) * p);
   }
   return {
-    rows: new Float32Array(f),
+    rows: new Float32Array(c),
     cols: new Float32Array(d),
-    vals: new Float32Array(c)
+    vals: new Float32Array(f)
   };
 }
-const Z = `// UMAP SGD compute shader — processes one graph edge per GPU thread.
+const $ = `// UMAP SGD compute shader — processes one graph edge per GPU thread.
 // Applies attraction forces between connected nodes and repulsion forces
 // against negative samples.
 
@@ -158,6 +158,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     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];
@@ -196,7 +197,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     epochs_per_sample[edge_idx] / f32(params.negative_sample_rate);
 }
 `;
-class W {
+class j {
   constructor() {
     G(this, "device");
     G(this, "pipeline");
@@ -207,7 +208,7 @@ class W {
     this.device = await t.requestDevice(), this.pipeline = this.device.createComputePipeline({
       layout: "auto",
       compute: {
-        module: this.device.createShaderModule({ code: Z }),
+        module: this.device.createShaderModule({ code: $ }),
         entryPoint: "main"
       }
     });
@@ -225,45 +226,45 @@ class W {
    * @param params      - UMAP curve parameters and repulsion settings
    * @returns Optimized embedding as Float32Array
    */
-  async optimize(t, r, s, n, p, f, d, c, a) {
+  async optimize(t, r, s, n, g, c, d, f, a) {
     const { device: o } = this, i = r.length, h = this.makeBuffer(
       t,
       GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
-    ), l = this.makeBuffer(r, GPUBufferUsage.STORAGE), u = this.makeBuffer(s, GPUBufferUsage.STORAGE), g = this.makeBuffer(n, GPUBufferUsage.STORAGE), m = new Float32Array(i).fill(0), w = this.makeBuffer(m, GPUBufferUsage.STORAGE), b = new Float32Array(i);
-    for (let _ = 0; _ < i; _++)
-      b[_] = n[_] / 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({
+    ), l = this.makeBuffer(r, GPUBufferUsage.STORAGE), u = this.makeBuffer(s, GPUBufferUsage.STORAGE), p = this.makeBuffer(n, GPUBufferUsage.STORAGE), m = new Float32Array(i).fill(0), _ = this.makeBuffer(m, GPUBufferUsage.STORAGE), w = new Float32Array(i);
+    for (let M = 0; M < i; M++)
+      w[M] = n[M] / f.negativeSampleRate;
+    const b = this.makeBuffer(w, GPUBufferUsage.STORAGE), A = new Uint32Array(i);
+    for (let M = 0; M < i; M++)
+      A[M] = Math.random() * 4294967295 | 0;
+    const x = this.makeBuffer(A, GPUBufferUsage.STORAGE), U = o.createBuffer({
       size: 40,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
     });
-    for (let _ = 0; _ < d; _++) {
-      const F = 1 - _ / d, 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] = d, 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({
+    for (let M = 0; M < d; M++) {
+      const v = 1 - M / d, S = new ArrayBuffer(40), y = new Uint32Array(S), B = new Float32Array(S);
+      y[0] = i, y[1] = g, y[2] = c, y[3] = M, y[4] = d, B[5] = v, B[6] = f.a, B[7] = f.b, B[8] = f.gamma, y[9] = f.negativeSampleRate, o.queue.writeBuffer(U, 0, S);
+      const k = o.createBindGroup({
         layout: this.pipeline.getBindGroupLayout(0),
         entries: [
-          { binding: 0, resource: { buffer: g } },
+          { binding: 0, resource: { buffer: p } },
           { binding: 1, resource: { buffer: l } },
           { binding: 2, resource: { buffer: u } },
           { binding: 3, resource: { buffer: h } },
-          { binding: 4, resource: { buffer: w } },
-          { binding: 5, resource: { buffer: y } },
-          { binding: 6, resource: { buffer: B } },
+          { binding: 4, resource: { buffer: _ } },
+          { binding: 5, resource: { buffer: b } },
+          { binding: 6, resource: { buffer: U } },
           { binding: 7, resource: { buffer: x } }
         ]
-      }), 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(_, d));
+      }), E = o.createCommandEncoder(), P = E.beginComputePass();
+      P.setPipeline(this.pipeline), P.setBindGroup(0, k), P.dispatchWorkgroups(Math.ceil(i / 256)), P.end(), o.queue.submit([E.finish()]), M % 10 === 0 && (await o.queue.onSubmittedWorkDone(), a == null || a(M, d));
     }
-    const E = o.createBuffer({
+    const N = o.createBuffer({
       size: t.byteLength,
       usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
     }), R = o.createCommandEncoder();
-    R.copyBufferToBuffer(h, 0, E, 0, t.byteLength), o.queue.submit([R.finish()]), await E.mapAsync(GPUMapMode.READ);
-    const k = new Float32Array(E.getMappedRange().slice(0));
-    return E.unmap(), h.destroy(), l.destroy(), u.destroy(), g.destroy(), w.destroy(), y.destroy(), x.destroy(), B.destroy(), E.destroy(), k;
+    R.copyBufferToBuffer(h, 0, N, 0, t.byteLength), o.queue.submit([R.finish()]), await N.mapAsync(GPUMapMode.READ);
+    const O = new Float32Array(N.getMappedRange().slice(0));
+    return N.unmap(), h.destroy(), l.destroy(), u.destroy(), p.destroy(), _.destroy(), b.destroy(), x.destroy(), U.destroy(), N.destroy(), O;
   }
   makeBuffer(t, r) {
     const s = this.device.createBuffer({
@@ -274,86 +275,90 @@ class W {
     return t instanceof Float32Array ? new Float32Array(s.getMappedRange()).set(t) : new Uint32Array(s.getMappedRange()).set(t), s.unmap(), s;
   }
 }
-function q(e) {
+function z(e) {
   return Math.max(-4, Math.min(4, e));
 }
-function z(e, t, r, s, n, p, f, d) {
-  const { a: c, b: a, gamma: o = 1, negativeSampleRate: i = 5 } = f, h = t.rows.length, l = new Uint32Array(t.rows), u = new Uint32Array(t.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++) {
-    d == null || d(w, p);
-    const b = 1 - w / p;
-    for (let y = 0; y < h; y++) {
-      if (g[y] > w) continue;
-      const v = l[y], x = u[y];
-      let B = 0;
-      for (let _ = 0; _ < n; _++) {
-        const F = e[v * n + _] - e[x * n + _];
-        B += F * F;
+function I(e, t, r, s, n, g, c, d) {
+  const { a: f, b: a, gamma: o = 1, negativeSampleRate: i = 5 } = c, h = t.rows.length, l = new Uint32Array(t.rows), u = new Uint32Array(t.cols), p = new Float32Array(h).fill(0), m = new Float32Array(h);
+  for (let _ = 0; _ < h; _++)
+    m[_] = r[_] / i;
+  for (let _ = 0; _ < g; _++) {
+    d == null || d(_, g);
+    const w = 1 - _ / g;
+    for (let b = 0; b < h; b++) {
+      if (p[b] > _) continue;
+      const A = l[b], x = u[b];
+      let U = 0;
+      for (let v = 0; v < n; v++) {
+        const S = e[A * n + v] - e[x * n + v];
+        U += S * S;
       }
-      const E = Math.pow(B, a), R = -2 * c * a * (B > 0 ? E / B : 0) / (c * E + 1);
-      for (let _ = 0; _ < n; _++) {
-        const F = e[v * n + _] - e[x * n + _], A = q(R * F);
-        e[v * n + _] += b * A, e[x * n + _] -= b * A;
+      const N = Math.pow(U, a), R = -2 * f * a * (U > 0 ? N / U : 0) / (f * N + 1);
+      for (let v = 0; v < n; v++) {
+        const S = e[A * n + v] - e[x * n + v], y = z(R * S);
+        e[A * n + v] += w * y, e[x * n + v] -= w * y;
       }
-      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 < n; S++) {
-          const O = e[v * n + S] - e[F * n + S];
-          A += O * O;
+      p[b] += r[b];
+      const O = r[b] / i, M = Math.max(0, Math.floor(
+        (_ - m[b]) / O
+      ));
+      m[b] += M * O;
+      for (let v = 0; v < M; v++) {
+        const S = Math.floor(Math.random() * s);
+        if (S === A) continue;
+        let y = 0;
+        for (let E = 0; E < n; E++) {
+          const P = e[A * n + E] - e[S * n + E];
+          y += P * P;
         }
-        const M = Math.pow(A, a), N = 2 * o * a / ((1e-3 + A) * (c * M + 1));
-        for (let S = 0; S < n; S++) {
-          const O = e[v * n + S] - e[F * n + S], U = q(N * O);
-          e[v * n + S] += b * U;
+        const B = Math.pow(y, a), k = 2 * o * a / ((1e-3 + y) * (f * B + 1));
+        for (let E = 0; E < n; E++) {
+          const P = e[A * n + E] - e[S * n + E], F = z(k * P);
+          e[A * n + E] += w * F;
         }
       }
-      m[y] += r[y] / i;
     }
   }
   return e;
 }
-function $(e, t, r, s, n, p, f, d, c, a) {
-  const { a: o, b: i, gamma: h = 1, negativeSampleRate: l = 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] / l;
-  for (let y = 0; y < d; y++) {
-    const v = 1 - y / d;
+function ee(e, t, r, s, n, g, c, d, f, a) {
+  const { a: o, b: i, gamma: h = 1, negativeSampleRate: l = 5 } = f, u = r.rows.length, p = new Uint32Array(r.rows), m = new Uint32Array(r.cols), _ = new Float32Array(u).fill(0), w = new Float32Array(u);
+  for (let b = 0; b < u; b++)
+    w[b] = s[b] / l;
+  for (let b = 0; b < d; b++) {
+    const A = 1 - b / d;
     for (let x = 0; x < u; x++) {
-      if (w[x] > y) continue;
-      const B = g[x], E = m[x];
+      if (_[x] > b) continue;
+      const U = p[x], N = m[x];
       let R = 0;
-      for (let A = 0; A < f; A++) {
-        const M = e[B * f + A] - t[E * f + A];
-        R += M * M;
+      for (let y = 0; y < c; y++) {
+        const B = e[U * c + y] - t[N * c + y];
+        R += B * B;
       }
-      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 M = e[B * f + A] - t[E * f + A];
-        e[B * f + A] += v * q(_ * M);
+      const O = Math.pow(R, i), M = -2 * o * i * (R > 0 ? O / R : 0) / (o * O + 1);
+      for (let y = 0; y < c; y++) {
+        const B = e[U * c + y] - t[N * c + y];
+        e[U * c + y] += A * z(M * B);
       }
-      w[x] += s[x];
-      const F = b[x] > 0 ? Math.floor(s[x] / b[x]) : 0;
-      for (let A = 0; A < F; A++) {
-        const M = Math.floor(Math.random() * p);
-        if (M === E) continue;
-        let N = 0;
-        for (let U = 0; U < f; U++) {
-          const P = e[B * f + U] - t[M * f + U];
-          N += P * P;
+      _[x] += s[x];
+      const v = s[x] / l, S = Math.max(0, Math.floor(
+        (b - w[x]) / v
+      ));
+      w[x] += S * v;
+      for (let y = 0; y < S; y++) {
+        const B = Math.floor(Math.random() * g);
+        if (B === N) continue;
+        let k = 0;
+        for (let F = 0; F < c; F++) {
+          const q = e[U * c + F] - t[B * c + F];
+          k += q * q;
         }
-        const S = Math.pow(N, i), O = 2 * h * i / ((1e-3 + N) * (o * S + 1));
-        for (let U = 0; U < f; U++) {
-          const P = e[B * f + U] - t[M * f + U];
-          e[B * f + U] += v * q(O * P);
+        const E = Math.pow(k, i), P = 2 * h * i / ((1e-3 + k) * (o * E + 1));
+        for (let F = 0; F < c; F++) {
+          const q = e[U * c + F] - t[B * c + F];
+          e[U * c + F] += A * z(P * q);
         }
       }
-      b[x] += s[x] / l;
     }
   }
   return e;
@@ -361,66 +366,66 @@ function $(e, t, r, s, n, p, f, d, c, a) {
 function K() {
   return typeof navigator < "u" && !!navigator.gpu;
 }
-async function ae(e, t = {}, r) {
+async function re(e, t = {}, r) {
   const {
     nComponents: s = 2,
     nNeighbors: n = 15,
-    minDist: p = 0.1,
-    spread: f = 1,
+    minDist: g = 0.1,
+    spread: c = 1,
     hnsw: d = {}
-  } = t, c = t.nEpochs ?? (e.length > 1e4 ? 200 : 500);
+  } = t, f = t.nEpochs ?? (e.length > 1e4 ? 200 : 500);
   console.time("knn");
-  const { indices: a, distances: o } = await Y(e, n, {
+  const { indices: a, distances: o } = await Q(e, n, {
     M: d.M ?? 16,
     efConstruction: d.efConstruction ?? 200,
     efSearch: d.efSearch ?? 50
   });
   console.timeEnd("knn"), console.time("fuzzy-set");
-  const i = D(a, o, n);
+  const i = L(a, o, n);
   console.timeEnd("fuzzy-set");
-  const { a: h, b: l } = j(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;
+  const { a: h, b: l } = H(g, c), u = T(i.vals), p = e.length, m = new Float32Array(p * s);
+  for (let w = 0; w < m.length; w++)
+    m[w] = Math.random() * 20 - 10;
   console.time("sgd");
-  let w;
+  let _;
   if (K())
     try {
-      const b = new W();
-      await b.init(), w = await b.optimize(
+      const w = new j();
+      await w.init(), _ = await w.optimize(
         m,
         new Uint32Array(i.rows),
         new Uint32Array(i.cols),
         u,
-        g,
+        p,
         s,
-        c,
+        f,
         { a: h, b: l, gamma: 1, negativeSampleRate: 5 },
         r
       );
-    } catch (b) {
-      console.warn("WebGPU SGD failed, falling back to CPU:", b), w = z(m, i, u, g, s, c, { a: h, b: l }, r);
+    } catch (w) {
+      console.warn("WebGPU SGD failed, falling back to CPU:", w), _ = I(m, i, u, p, s, f, { a: h, b: l }, r);
     }
   else
-    w = z(m, i, u, g, s, c, { a: h, b: l }, r);
-  return console.timeEnd("sgd"), w;
+    _ = I(m, i, u, p, s, f, { a: h, b: l }, r);
+  return console.timeEnd("sgd"), _;
 }
-function j(e, t) {
+function H(e, t) {
   if (Math.abs(t - 1) < 1e-6 && Math.abs(e - 0.1) < 1e-6)
     return { a: 1.9292, b: 0.7915 };
   if (Math.abs(t - 1) < 1e-6 && Math.abs(e - 0) < 1e-6)
     return { a: 1.8956, b: 0.8006 };
   if (Math.abs(t - 1) < 1e-6 && Math.abs(e - 0.5) < 1e-6)
     return { a: 1.5769, b: 0.8951 };
-  const r = ee(e, t);
-  return { a: te(e, t, r), b: r };
+  const r = te(e, t);
+  return { a: ne(e, t, r), b: r };
 }
-function ee(e, t) {
+function te(e, t) {
   return 1 / (t * 1.2);
 }
-function te(e, t, r) {
+function ne(e, t, r) {
   return e < 1e-6 ? 1.8956 : (1 / (1 + 1e-3) - 1) / -Math.pow(e, 2 * r);
 }
-class re {
+class ie {
   constructor(t = {}) {
     G(this, "_nComponents");
     G(this, "_nNeighbors");
@@ -435,7 +440,7 @@ class re {
     G(this, "_hnswIndex", null);
     G(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: r, b: s } = j(this._minDist, this._spread);
+    const { a: r, b: s } = H(this._minDist, this._spread);
     this._a = r, this._b = s;
   }
   /**
@@ -445,22 +450,22 @@ class re {
    * Returns `this` for chaining.
    */
   async fit(t, r) {
-    const s = t.length, n = this._nEpochs ?? (s > 1e4 ? 200 : 500), { M: p = 16, efConstruction: f = 200, efSearch: d = 50 } = this._hnswOpts;
+    const s = t.length, n = this._nEpochs ?? (s > 1e4 ? 200 : 500), { M: g = 16, efConstruction: c = 200, efSearch: d = 50 } = this._hnswOpts;
     console.time("knn");
-    const { knn: c, index: a } = await Q(t, this._nNeighbors, {
-      M: p,
-      efConstruction: f,
+    const { knn: f, index: a } = await J(t, this._nNeighbors, {
+      M: g,
+      efConstruction: c,
       efSearch: d
     });
     this._hnswIndex = a, this._nTrain = s, console.timeEnd("knn"), console.time("fuzzy-set");
-    const o = D(c.indices, c.distances, this._nNeighbors);
+    const o = L(f.indices, f.distances, this._nNeighbors);
     console.timeEnd("fuzzy-set");
-    const i = I(o.vals, n), h = new Float32Array(s * this._nComponents);
+    const i = T(o.vals), h = new Float32Array(s * this._nComponents);
     for (let l = 0; l < h.length; l++)
       h[l] = Math.random() * 20 - 10;
     if (console.time("sgd"), K())
       try {
-        const l = new W();
+        const l = new j();
         await l.init(), this.embedding = await l.optimize(
           h,
           new Uint32Array(o.rows),
@@ -473,13 +478,13 @@ class re {
           r
         );
       } catch (l) {
-        console.warn("WebGPU SGD failed, falling back to CPU:", l), this.embedding = z(h, o, i, s, this._nComponents, n, {
+        console.warn("WebGPU SGD failed, falling back to CPU:", l), this.embedding = I(h, o, i, s, this._nComponents, n, {
           a: this._a,
           b: this._b
         }, r);
       }
     else
-      this.embedding = z(h, o, i, s, this._nComponents, n, {
+      this.embedding = I(h, o, i, s, this._nComponents, n, {
         a: this._a,
         b: this._b
       }, r);
@@ -499,21 +504,21 @@ class re {
   async transform(t, r = !1) {
     if (!this._hnswIndex || !this.embedding)
       throw new Error("UMAP.transform() must be called after fit()");
-    const s = t.length, n = this._nEpochs ?? (this._nTrain > 1e4 ? 200 : 500), p = Math.max(100, Math.floor(n / 4)), f = this._hnswIndex.searchKnn(t, this._nNeighbors), d = J(f.indices, f.distances, this._nNeighbors), c = new Uint32Array(d.rows), a = new Uint32Array(d.cols), o = new Float32Array(s), i = new Float32Array(s * this._nComponents);
-    for (let u = 0; u < c.length; u++) {
-      const g = c[u], m = a[u], w = d.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];
+    const s = t.length, n = this._nEpochs ?? (this._nTrain > 1e4 ? 200 : 500), g = Math.max(100, Math.floor(n / 4)), c = this._hnswIndex.searchKnn(t, this._nNeighbors), d = X(c.indices, c.distances, this._nNeighbors), f = new Uint32Array(d.rows), a = new Uint32Array(d.cols), o = new Float32Array(s), i = new Float32Array(s * this._nComponents);
+    for (let u = 0; u < f.length; u++) {
+      const p = f[u], m = a[u], _ = d.vals[u];
+      o[p] += _;
+      for (let w = 0; w < this._nComponents; w++)
+        i[p * this._nComponents + w] += _ * this.embedding[m * this._nComponents + w];
     }
     for (let u = 0; u < s; u++)
       if (o[u] > 0)
-        for (let g = 0; g < this._nComponents; g++)
-          i[u * this._nComponents + g] /= o[u];
+        for (let p = 0; p < this._nComponents; p++)
+          i[u * this._nComponents + p] /= o[u];
       else
-        for (let g = 0; g < this._nComponents; g++)
-          i[u * this._nComponents + g] = Math.random() * 20 - 10;
-    const h = I(d.vals, p), l = $(
+        for (let p = 0; p < this._nComponents; p++)
+          i[u * this._nComponents + p] = Math.random() * 20 - 10;
+    const h = T(d.vals), l = ee(
       i,
       this.embedding,
       d,
@@ -521,10 +526,10 @@ class re {
       s,
       this._nTrain,
       this._nComponents,
-      p,
+      g,
       { a: this._a, b: this._b }
     );
-    return r ? T(l, s, this._nComponents) : l;
+    return r ? C(l, s, this._nComponents) : l;
   }
   /**
    * Convenience method equivalent to `fit(vectors)` followed by
@@ -536,36 +541,36 @@ class re {
    *   Defaults to `false`.
    */
   async fit_transform(t, r, s = !1) {
-    return await this.fit(t, r), s ? T(this.embedding, t.length, this._nComponents) : this.embedding;
+    return await this.fit(t, r), s ? C(this.embedding, t.length, this._nComponents) : this.embedding;
   }
 }
-function T(e, t, r) {
+function C(e, t, r) {
   const s = new Float32Array(e.length);
   for (let n = 0; n < r; n++) {
-    let p = 1 / 0, f = -1 / 0;
-    for (let c = 0; c < t; c++) {
-      const a = e[c * r + n];
-      a < p && (p = a), a > f && (f = a);
+    let g = 1 / 0, c = -1 / 0;
+    for (let f = 0; f < t; f++) {
+      const a = e[f * r + n];
+      a < g && (g = a), a > c && (c = a);
     }
-    const d = f - p;
-    for (let c = 0; c < t; c++)
-      s[c * r + n] = d > 0 ? (e[c * r + n] - p) / d : 0;
+    const d = c - g;
+    for (let f = 0; f < t; f++)
+      s[f * r + n] = d > 0 ? (e[f * r + n] - g) / d : 0;
   }
   return s;
 }
-function I(e, t) {
+function T(e, t) {
   let r = -1 / 0;
   for (let n = 0; n < e.length; n++)
     e[n] > r && (r = e[n]);
   const s = new Float32Array(e.length);
   for (let n = 0; n < e.length; n++) {
-    const p = e[n] / r;
-    s[n] = p > 0 ? t / p : -1;
+    const g = e[n] / r;
+    s[n] = g > 0 ? 1 / g : -1;
   }
   return s;
 }
 export {
-  re as UMAP,
-  ae as fit,
+  ie as UMAP,
+  re as fit,
   K as isWebGPUAvailable
 };

package/package.json

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