umap-gpu 0.2.6 → 0.2.9

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): 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): 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): 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): Float32Array;

package/dist/gpu/sgd.d.ts

@@ -25,6 +25,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): 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): Promise<Float32Array>;
     private makeBuffer;
 }

package/dist/index.d.ts

@@ -1,5 +1,5 @@
 export { fit, UMAP } from './umap';
-export type { UMAPOptions } from './umap';
+export type { UMAPOptions, ProgressCallback } from './umap';
 export type { KNNResult, HNSWOptions, HNSWSearchableIndex } from './hnsw-knn';
 export type { FuzzyGraph } from './fuzzy-set';
 export { isWebGPUAvailable } from './gpu/device';

package/dist/index.js

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

package/dist/umap.d.ts

@@ -16,6 +16,14 @@ export interface UMAPOptions {
         efSearch?: number;
     };
 }
+/**
+ * Called after each completed SGD epoch (or every 10 epochs on the GPU path,
+ * piggybacking on the existing GPU synchronisation point to avoid extra stalls).
+ *
+ * @param epoch   - Zero-based index of the epoch that just finished.
+ * @param nEpochs - Total number of epochs.
+ */
+export type ProgressCallback = (epoch: number, nEpochs: number) => void;
 /**
  * Fit UMAP to the given high-dimensional vectors and return a low-dimensional embedding.
  *
@@ -24,7 +32,7 @@ export interface UMAPOptions {
  * 2. Fuzzy simplicial set construction (graph weights)
  * 3. SGD optimization (WebGPU accelerated, with CPU fallback)
  */
-export declare function fit(vectors: number[][], opts?: UMAPOptions): Promise<Float32Array>;
+export declare function fit(vectors: number[][], opts?: UMAPOptions, onProgress?: ProgressCallback): Promise<Float32Array>;
 /**
  * Compute the a, b parameters for the UMAP curve 1/(1 + a*d^(2b)).
  *
@@ -74,21 +82,29 @@ export declare class UMAP {
      * index so that transform() can project new points later.
      * Returns `this` for chaining.
      */
-    fit(vectors: number[][]): Promise<this>;
+    fit(vectors: number[][], onProgress?: ProgressCallback): Promise<this>;
     /**
      * Project new (unseen) `vectors` into the embedding space learned by fit().
      * Must be called after fit().
      *
      * The training embedding is kept fixed; only the new-point positions are
      * optimised. Returns a Float32Array of shape [vectors.length × nComponents].
+     *
+     * @param normalize - When `true`, min-max normalise each dimension of the
+     *   returned embedding to [0, 1].  The stored training embedding is never
+     *   mutated.  Defaults to `false`.
      */
-    transform(vectors: number[][]): Promise<Float32Array>;
+    transform(vectors: number[][], normalize?: boolean): Promise<Float32Array>;
     /**
      * Convenience method equivalent to `fit(vectors)` followed by
      * `transform(vectors)` — but more efficient because the training embedding
      * is returned directly without a second optimization pass.
+     *
+     * @param normalize - When `true`, min-max normalise each dimension of the
+     *   returned embedding to [0, 1].  `this.embedding` is never mutated.
+     *   Defaults to `false`.
      */
-    fit_transform(vectors: number[][]): Promise<Float32Array>;
+    fit_transform(vectors: number[][], onProgress?: ProgressCallback, normalize?: boolean): Promise<Float32Array>;
 }
 /**
  * Compute per-edge epoch sampling periods based on edge weights.

package/package.json

@@ -1,30 +1,60 @@
 {
   "name": "umap-gpu",
-  "version": "0.2.6",
+  "version": "0.2.9",
   "description": "UMAP with HNSW kNN and WebGPU-accelerated SGD",
   "type": "module",
   "main": "dist/index.js",
-  "types": "dist/index.d.ts",
+  "types": "dist/umap-gpu.d.ts",
+  "sideEffects": false,
+  "exports": {
+    ".": {
+      "import": {
+        "types": "./dist/umap-gpu.d.ts",
+        "default": "./dist/index.js"
+      }
+    }
+  },
   "files": [
     "dist"
   ],
+  "engines": {
+    "node": ">=18"
+  },
+  "keywords": [
+    "umap",
+    "dimensionality-reduction",
+    "webgpu",
+    "hnsw",
+    "machine-learning",
+    "gpu"
+  ],
   "repository": {
     "type": "git",
     "url": "https://github.com/Achuttarsing/umap-gpu"
   },
+  "bugs": {
+    "url": "https://github.com/Achuttarsing/umap-gpu/issues"
+  },
+  "homepage": "https://achuttarsing.github.io/umap-gpu",
   "scripts": {
     "build": "vite build && tsc",
     "dev": "vite",
     "test": "vitest run",
-    "prepublishOnly": "npm test && npm run build"
+    "prepublishOnly": "bun test && bun run build",
+    "docs:dev": "vitepress dev docs",
+    "docs:build": "vitepress build docs",
+    "docs:generate": "bun run build && bunx api-extractor run && bun run docs:build"
   },
   "dependencies": {
     "hnswlib-wasm": "^0.8.2"
   },
   "devDependencies": {
+    "@microsoft/api-extractor": "^7.57.6",
     "@webgpu/types": "^0.1.40",
     "typescript": "^5.4.0",
     "vite": "^5.0.0",
+    "vitepress": "^1.6.4",
+    "vitepress-plugin-llms": "^1.11.0",
     "vitest": "^4.0.18"
   },
   "license": "MIT"