umap-gpu 0.1.0 → 0.2.6

package/README.md

@@ -1,87 +1,104 @@
 # umap-gpu
 
-UMAP dimensionality reduction with HNSW k-nearest-neighbor search and WebGPU-accelerated SGD optimization, with a transparent CPU fallback.
+UMAP dimensionality reduction with WebGPU-accelerated SGD and HNSW approximate nearest neighbors.
 
-## What it does
+Embed millions of high-dimensional vectors into 2D in seconds — not minutes.
 
-Takes a set of high-dimensional vectors and returns a low-dimensional embedding (default: 2D) suitable for visualization or downstream tasks.
+## Why GPU?
 
-The pipeline runs in three stages:
+The bottleneck in UMAP is the SGD optimization loop: thousands of epochs, millions of edge updates per epoch. On CPU this is sequential. On GPU, all edges run in parallel across thousands of shader cores — expect a significant speedup on large datasets, scaling with both the number of points and the number of epochs.
 
-1. **k-NN** — approximate nearest neighbors via [hnswlib-wasm](https://github.com/yoshoku/hnswlib-wasm) (O(n log n))
-2. **Fuzzy simplicial set** — builds a weighted graph from the k-NN graph using smooth distances
-3. **SGD** — optimizes the embedding using attraction/repulsion forces:
-   - **WebGPU** compute shader when available (Chrome 113+, Edge 113+)
-   - **CPU** fallback otherwise — identical output, just slower
+The k-NN stage uses [hnswlib-wasm](https://github.com/yoshoku/hnswlib-wasm) (O(n log n)) so it stays fast regardless.
+A transparent CPU fallback guarantees identical output everywhere WebGPU isn't available.
 
 ## Install
 
 ```bash
+# npm
 npm install umap-gpu
-```
 
-> Requires a browser or runtime with WebGPU support for GPU acceleration. The CPU fallback works anywhere.
+# Bun
+bun add umap-gpu
+
+# pnpm
+pnpm add umap-gpu
+```
 
-## Usage
+## Quick start
 
 ```ts
 import { fit } from 'umap-gpu';
 
 const vectors = [
-  [1.0, 0.0, 0.3],
-  [0.9, 0.1, 0.4],
-  [0.0, 1.0, 0.8],
+  [0.1, 0.4, 0.9, ...],  // high-dimensional points
+  [0.2, 0.3, 0.8, ...],
   // ...
 ];
 
 const embedding = await fit(vectors);
-// Float32Array of length n * nComponents (default: n * 2)
-// embedding[i*2], embedding[i*2 + 1] → 2D coordinates of point i
+// Float32Array — embedding[i*2], embedding[i*2+1] are the 2D coords of point i
+```
+
+## Train once, project many times
+
+Use the `UMAP` class to embed a training set and later project new points into the same space without retraining.
+
+```ts
+import { UMAP } from 'umap-gpu';
+
+const umap = new UMAP({ nNeighbors: 15, minDist: 0.1 });
+
+// Train
+await umap.fit(trainVectors);
+console.log(umap.embedding); // Float32Array [nTrain × 2]
+
+// Project new points (training embedding stays fixed)
+const projected = await umap.transform(newVectors);
+// Float32Array [nNew × 2]
 ```
 
-### Options
+## Options
 
 ```ts
-const embedding = await fit(vectors, {
-  nComponents: 2,      // output dimensions (default: 2)
-  nNeighbors:  15,     // k-NN graph degree (default: 15)
-  nEpochs:     500,    // SGD iterations (default: 500 for <10k points, 200 otherwise)
-  minDist:     0.1,    // minimum distance between points in the embedding (default: 0.1)
-  spread:      1.0,    // scale of the embedding (default: 1.0)
+const umap = new UMAP({
+  nComponents: 2,      // output dimensions          (default: 2)
+  nNeighbors:  15,     // k-NN graph degree          (default: 15)
+  nEpochs:     500,    // SGD iterations             (default: auto — 500 for <10k points, 200 otherwise)
+  minDist:     0.1,    // min distance in embedding  (default: 0.1)
+  spread:      1.0,    // scale of the embedding     (default: 1.0)
   hnsw: {
-    M:               16,  // HNSW graph connectivity (default: 16)
-    efConstruction: 200,  // build-time search width (default: 200)
-    efSearch:        50,  // query-time search width (default: 50)
+    M:               16,  // graph connectivity        (default: 16)
+    efConstruction: 200,  // build-time search width  (default: 200)
+    efSearch:        50,  // query-time search width  (default: 50)
   },
 });
+
+// Same options work with the functional API
+const embedding = await fit(vectors, { nNeighbors: 15, minDist: 0.05 });
 ```
 
-### Checking GPU availability
+## Check GPU availability
 
 ```ts
 import { isWebGPUAvailable } from 'umap-gpu';
 
-if (isWebGPUAvailable()) {
-  console.log('Will use WebGPU-accelerated SGD');
-} else {
-  console.log('Will fall back to CPU SGD');
-}
-```
-
-## Build
-
-```bash
-npm run build   # compiles TypeScript to dist/
-npm test        # runs the unit test suite (Vitest)
+console.log(isWebGPUAvailable()); // true → GPU path, false → CPU fallback
 ```
 
 ## Browser support
 
-| Feature | Requirement |
+| Feature | Supported in |
 |---------|-------------|
 | WebGPU SGD | Chrome 113+, Edge 113+, Safari 18+ |
-| CPU fallback | Any modern browser / Node.js |
-| HNSW (WASM) | Any environment with WebAssembly support |
+| CPU fallback | Any modern browser / Node.js / Bun |
+| HNSW (WASM) | Any environment with WebAssembly |
+
+## Development
+
+```bash
+npm test        # Vitest unit tests
+npm run build   # TypeScript → dist/
+```
 
 ## License
 

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

@@ -10,3 +10,20 @@ export interface CPUSgdParams {
  * 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;
+/**
+ * CPU SGD for UMAP.transform(): optimizes only the new-point embeddings.
+ * The training embedding is read-only; attraction pulls new points toward
+ * their training neighbors, and repulsion pushes them away from random
+ * training points.
+ *
+ * @param embeddingNew   - New-point embeddings to optimize [nNew × nComponents]
+ * @param embeddingTrain - Fixed training embeddings [nTrain × nComponents]
+ * @param graph          - Bipartite graph: rows=new-point indices, cols=training-point indices
+ * @param epochsPerSample - Per-edge epoch sampling schedule
+ * @param nNew           - Number of new points
+ * @param nTrain         - Number of training points
+ * @param nComponents    - Embedding dimensionality
+ * @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;

package/dist/fuzzy-set.d.ts

@@ -10,3 +10,16 @@ export interface FuzzyGraph {
  * (sigmas, rhos) and symmetrizes with the fuzzy set union operation.
  */
 export declare function computeFuzzySimplicialSet(knnIndices: number[][], knnDistances: number[][], nNeighbors: number, setOpMixRatio?: number): FuzzyGraph;
+/**
+ * Compute the fuzzy weight graph between new (query) points and training points.
+ * Used by UMAP.transform() to project unseen data into an existing embedding.
+ *
+ * Unlike computeFuzzySimplicialSet, this produces a bipartite graph
+ * (new points → training points) with no symmetrization.
+ *
+ * @param knnIndices   - For each new point, the indices of its training neighbors
+ * @param knnDistances - For each new point, the distances to those neighbors
+ * @param nNeighbors   - Number of neighbors used
+ * @returns FuzzyGraph where rows are new-point indices, cols are training-point indices
+ */
+export declare function computeTransformFuzzyWeights(knnIndices: number[][], knnDistances: number[][], nNeighbors: number): FuzzyGraph;

package/dist/hnsw-knn.d.ts

@@ -7,9 +7,24 @@ export interface HNSWOptions {
     efConstruction?: number;
     efSearch?: number;
 }
+/**
+ * A built HNSW index that can be queried to find nearest neighbors in the
+ * training data for new (unseen) points — used by UMAP.transform().
+ */
+export interface HNSWSearchableIndex {
+    searchKnn(queryVectors: number[][], nNeighbors: number): KNNResult;
+}
 /**
  * Compute k-nearest neighbors using HNSW (Hierarchical Navigable Small World)
  * via hnswlib-wasm, replacing the O(n^2) brute-force search in umap-js with
  * an O(n log n) approximate nearest neighbor search.
  */
 export declare function computeKNN(vectors: number[][], nNeighbors: number, opts?: HNSWOptions): Promise<KNNResult>;
+/**
+ * Like computeKNN, but also returns the built HNSW index so it can be reused
+ * later to project new points (used by UMAP.transform()).
+ */
+export declare function computeKNNWithIndex(vectors: number[][], nNeighbors: number, opts?: HNSWOptions): Promise<{
+    knn: KNNResult;
+    index: HNSWSearchableIndex;
+}>;

package/dist/index.d.ts

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

package/dist/index.js

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

package/dist/umap.d.ts

@@ -36,6 +36,60 @@ export declare function findAB(minDist: number, spread: number): {
     a: number;
     b: number;
 };
+/**
+ * Stateful UMAP model that supports separate fit / transform / fit_transform.
+ *
+ * Usage:
+ * ```ts
+ * const umap = new UMAP({ nNeighbors: 15, nComponents: 2 });
+ *
+ * // Train on high-dimensional data:
+ * await umap.fit(trainVectors);
+ * console.log(umap.embedding); // Float32Array [nTrain * nComponents]
+ *
+ * // Project new points into the same space:
+ * const newEmbedding = await umap.transform(testVectors);
+ *
+ * // Or do both at once:
+ * const embedding = await umap.fit_transform(vectors);
+ * ```
+ */
+export declare class UMAP {
+    private readonly _nComponents;
+    private readonly _nNeighbors;
+    private readonly _minDist;
+    private readonly _spread;
+    private readonly _nEpochs;
+    private readonly _hnswOpts;
+    private readonly _a;
+    private readonly _b;
+    /** The low-dimensional embedding produced by the last fit() call. */
+    embedding: Float32Array | null;
+    private _hnswIndex;
+    private _nTrain;
+    constructor(opts?: UMAPOptions);
+    /**
+     * Train UMAP on `vectors`.
+     * Stores the resulting embedding in `this.embedding` and retains the HNSW
+     * index so that transform() can project new points later.
+     * Returns `this` for chaining.
+     */
+    fit(vectors: number[][]): 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].
+     */
+    transform(vectors: number[][]): 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.
+     */
+    fit_transform(vectors: number[][]): Promise<Float32Array>;
+}
 /**
  * Compute per-edge epoch sampling periods based on edge weights.
  * Higher-weight edges are sampled more frequently.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "umap-gpu",
-  "version": "0.1.0",
+  "version": "0.2.6",
   "description": "UMAP with HNSW kNN and WebGPU-accelerated SGD",
   "type": "module",
   "main": "dist/index.js",
@@ -8,6 +8,10 @@
   "files": [
     "dist"
   ],
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/Achuttarsing/umap-gpu"
+  },
   "scripts": {
     "build": "vite build && tsc",
     "dev": "vite",