umap-gpu 0.1.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Adrien Chuttarsing
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,88 @@
+# umap-gpu
+
+UMAP dimensionality reduction with HNSW k-nearest-neighbor search and WebGPU-accelerated SGD optimization, with a transparent CPU fallback.
+
+## What it does
+
+Takes a set of high-dimensional vectors and returns a low-dimensional embedding (default: 2D) suitable for visualization or downstream tasks.
+
+The pipeline runs in three stages:
+
+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
+
+## Install
+
+```bash
+npm install umap-gpu
+```
+
+> Requires a browser or runtime with WebGPU support for GPU acceleration. The CPU fallback works anywhere.
+
+## Usage
+
+```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],
+  // ...
+];
+
+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
+```
+
+### 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)
+  hnsw: {
+    M:               16,  // HNSW graph connectivity (default: 16)
+    efConstruction: 200,  // build-time search width (default: 200)
+    efSearch:        50,  // query-time search width (default: 50)
+  },
+});
+```
+
+### Checking 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)
+```
+
+## Browser support
+
+| Feature | Requirement |
+|---------|-------------|
+| WebGPU SGD | Chrome 113+, Edge 113+, Safari 18+ |
+| CPU fallback | Any modern browser / Node.js |
+| HNSW (WASM) | Any environment with WebAssembly support |
+
+## License
+
+MIT

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

@@ -0,0 +1,12 @@
+import type { FuzzyGraph } from '../fuzzy-set';
+export interface CPUSgdParams {
+    a: number;
+    b: number;
+    gamma?: number;
+    negativeSampleRate?: number;
+}
+/**
+ * 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;

package/dist/fuzzy-set.d.ts

@@ -0,0 +1,12 @@
+export interface FuzzyGraph {
+    rows: Float32Array;
+    cols: Float32Array;
+    vals: Float32Array;
+    nVertices: number;
+}
+/**
+ * Compute the fuzzy simplicial set from kNN results.
+ * This builds the high-dimensional graph weights using smooth kNN distances
+ * (sigmas, rhos) and symmetrizes with the fuzzy set union operation.
+ */
+export declare function computeFuzzySimplicialSet(knnIndices: number[][], knnDistances: number[][], nNeighbors: number, setOpMixRatio?: number): FuzzyGraph;

package/dist/gpu/device.d.ts

@@ -0,0 +1,12 @@
+/**
+ * WebGPU device management — handles adapter/device acquisition and
+ * provides a single shared device instance.
+ */
+/**
+ * Request and cache a WebGPU device. Returns null if WebGPU is not available.
+ */
+export declare function getGPUDevice(): Promise<GPUDevice | null>;
+/**
+ * Check whether WebGPU is available in the current environment.
+ */
+export declare function isWebGPUAvailable(): boolean;

package/dist/gpu/sgd.d.ts

@@ -0,0 +1,30 @@
+export interface SGDParams {
+    a: number;
+    b: number;
+    gamma: number;
+    negativeSampleRate: number;
+}
+/**
+ * GPU-accelerated SGD optimizer for UMAP embedding.
+ * Each GPU thread processes one graph edge, applying attraction and repulsion forces.
+ */
+export declare class GPUSgd {
+    private device;
+    private pipeline;
+    init(): Promise<void>;
+    /**
+     * Run SGD optimization on the GPU.
+     *
+     * @param embedding   - Initial embedding positions [nVertices * nComponents]
+     * @param head        - Edge source node indices
+     * @param tail        - Edge target node indices
+     * @param epochsPerSample - Per-edge epoch sampling period
+     * @param nVertices   - Number of data points
+     * @param nComponents - Embedding dimensionality (typically 2)
+     * @param nEpochs     - Total number of optimization epochs
+     * @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>;
+    private makeBuffer;
+}

package/dist/hnsw-knn.d.ts

@@ -0,0 +1,15 @@
+export interface KNNResult {
+    indices: number[][];
+    distances: number[][];
+}
+export interface HNSWOptions {
+    M?: number;
+    efConstruction?: number;
+    efSearch?: number;
+}
+/**
+ * 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>;

package/dist/index.d.ts

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

package/dist/index.js

@@ -0,0 +1,361 @@
+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: _ }) => _));
+  }
+  return { indices: r, distances: s };
+}
+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);
+    }
+  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);
+    }
+    t[p] = r;
+  }
+  return { sigmas: t, rhos: h };
+}
+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);
+  };
+  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
+    );
+  }
+  return {
+    rows: new Float32Array(u),
+    cols: new Float32Array(c),
+    vals: new Float32Array(o)
+  };
+}
+const W = `// UMAP SGD compute shader — processes one graph edge per GPU thread.
+// Applies attraction forces between connected nodes and repulsion forces
+// against negative samples.
+
+@group(0) @binding(0) var<storage, read>       epochs_per_sample : array<f32>;
+@group(0) @binding(1) var<storage, read>       head              : array<u32>;  // edge source
+@group(0) @binding(2) var<storage, read>       tail              : array<u32>;  // edge target
+@group(0) @binding(3) var<storage, read_write> embedding         : array<f32>;  // [n * nComponents]
+@group(0) @binding(4) var<storage, read_write> epoch_of_next_sample : array<f32>;
+@group(0) @binding(5) var<storage, read_write> epoch_of_next_negative_sample : array<f32>;
+@group(0) @binding(6) var<uniform>             params            : Params;
+@group(0) @binding(7) var<storage, read>       rng_seeds         : array<u32>;  // per-edge seed
+
+struct Params {
+  n_edges        : u32,
+  n_vertices     : u32,
+  n_components   : u32,
+  current_epoch  : u32,
+  n_epochs       : u32,
+  alpha          : f32,   // learning rate
+  a              : f32,
+  b              : f32,
+  gamma          : f32,   // repulsion strength
+  negative_sample_rate : u32,
+}
+
+fn clip(v: f32, lo: f32, hi: f32) -> f32 {
+  return max(lo, min(hi, v));
+}
+
+// Simple xorshift RNG per thread
+fn xorshift(seed: u32) -> u32 {
+  var s = seed;
+  s ^= s << 13u;
+  s ^= s >> 17u;
+  s ^= s << 5u;
+  return s;
+}
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let edge_idx = gid.x;
+  if (edge_idx >= params.n_edges) { return; }
+
+  // Only process this edge if its epoch has come
+  if (epoch_of_next_sample[edge_idx] > f32(params.current_epoch)) { return; }
+
+  let i = head[edge_idx];
+  let j = tail[edge_idx];
+  let nc = params.n_components;
+
+  // --- Attraction ---
+  var dist_sq : f32 = 0.0;
+  for (var d = 0u; d < nc; d++) {
+    let diff = embedding[i * nc + d] - embedding[j * nc + d];
+    dist_sq += diff * diff;
+  }
+
+  let grad_coeff_attr = -2.0 * params.a * params.b * pow(dist_sq, params.b - 1.0)
+                        / (params.a * pow(dist_sq, params.b) + 1.0);
+
+  for (var d = 0u; d < nc; d++) {
+    let diff = embedding[i * nc + d] - embedding[j * nc + d];
+    let grad = clip(grad_coeff_attr * diff, -4.0, 4.0);
+    embedding[i * nc + d] += params.alpha * grad;
+  }
+
+  epoch_of_next_sample[edge_idx] += epochs_per_sample[edge_idx];
+
+  // --- Repulsion (negative samples) ---
+  let eps = epochs_per_sample[edge_idx];
+  let neg_eps = epoch_of_next_negative_sample[edge_idx];
+  var n_neg = 0u;
+  if (neg_eps > 0.0) {
+    n_neg = u32(eps / neg_eps);
+  }
+  var rng = xorshift(rng_seeds[edge_idx] + params.current_epoch * 6364136223u);
+
+  for (var s = 0u; s < n_neg; s++) {
+    rng = xorshift(rng);
+    let k = rng % params.n_vertices;
+    if (k == i) { continue; }
+
+    var neg_dist_sq : f32 = 0.0;
+    for (var d = 0u; d < nc; d++) {
+      let diff = embedding[i * nc + d] - embedding[k * nc + d];
+      neg_dist_sq += diff * diff;
+    }
+
+    let grad_coeff_rep = 2.0 * params.gamma * params.b
+                         / ((0.001 + neg_dist_sq) * (params.a * pow(neg_dist_sq, params.b) + 1.0));
+
+    for (var d = 0u; d < nc; d++) {
+      let diff = embedding[i * nc + d] - embedding[k * nc + d];
+      let grad = clip(grad_coeff_rep * diff, -4.0, 4.0);
+      embedding[i * nc + d] += params.alpha * grad;
+    }
+  }
+
+  epoch_of_next_negative_sample[edge_idx] +=
+    epochs_per_sample[edge_idx] / f32(params.negative_sample_rate);
+}
+`;
+class j {
+  constructor() {
+    O(this, "device");
+    O(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({
+      layout: "auto",
+      compute: {
+        module: this.device.createShaderModule({ code: W }),
+        entryPoint: "main"
+      }
+    });
+  }
+  /**
+   * Run SGD optimization on the GPU.
+   *
+   * @param embedding   - Initial embedding positions [nVertices * nComponents]
+   * @param head        - Edge source node indices
+   * @param tail        - Edge target node indices
+   * @param epochsPerSample - Per-edge epoch sampling period
+   * @param nVertices   - Number of data points
+   * @param nComponents - Embedding dimensionality (typically 2)
+   * @param nEpochs     - Total number of optimization epochs
+   * @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,
+      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({
+      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({
+        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 } }
+        ]
+      }), 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();
+    }
+    const x = o.createBuffer({
+      size: n.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;
+  }
+  makeBuffer(n, a) {
+    const d = this.device.createBuffer({
+      size: n.byteLength,
+      usage: a,
+      mappedAtCreation: !0
+    });
+    return n instanceof Float32Array ? new Float32Array(d.getMappedRange()).set(n) : new Uint32Array(d.getMappedRange()).set(n), d.unmap(), d;
+  }
+}
+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];
+      let x = 0;
+      for (let i = 0; i < t; i++) {
+        const A = e[v * t + i] - e[M * t + i];
+        x += A * A;
+      }
+      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;
+      }
+      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;
+        }
+        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;
+        }
+      }
+      _[b] += a[b] / r;
+    }
+  }
+  return e;
+}
+function I() {
+  return typeof navigator < "u" && !!navigator.gpu;
+}
+async function Q(e, n = {}) {
+  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);
+  console.time("knn");
+  const { indices: c, distances: o } = await T(e, d, {
+    M: p.M ?? 16,
+    efConstruction: p.efConstruction ?? 200,
+    efSearch: p.efSearch ?? 50
+  });
+  console.timeEnd("knn"), console.time("fuzzy-set");
+  const r = L(c, o, d);
+  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;
+  console.time("sgd");
+  let w;
+  if (I())
+    try {
+      const g = new j();
+      await g.init(), w = await g.optimize(
+        _,
+        new Uint32Array(r.rows),
+        new Uint32Array(r.cols),
+        l,
+        m,
+        a,
+        u,
+        { a: s, b: f, 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 });
+    }
+  else
+    w = F(_, r, l, m, a, u, { a: s, b: f });
+  return console.timeEnd("sgd"), w;
+}
+function K(e, n) {
+  if (Math.abs(n - 1) < 1e-6 && Math.abs(e - 0.1) < 1e-6)
+    return { a: 1.9292, b: 0.7915 };
+  if (Math.abs(n - 1) < 1e-6 && Math.abs(e - 0) < 1e-6)
+    return { a: 1.8956, b: 0.8006 };
+  if (Math.abs(n - 1) < 1e-6 && Math.abs(e - 0.5) < 1e-6)
+    return { a: 1.5769, b: 0.8951 };
+  const a = H(e, n);
+  return { a: V(e, n, a), b: a };
+}
+function H(e, n) {
+  return 1 / (n * 1.2);
+}
+function V(e, n, a) {
+  return e < 1e-6 ? 1.8956 : (1 / (1 + 1e-3) - 1) / -Math.pow(e, 2 * a);
+}
+function Y(e, n) {
+  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;
+  }
+  return d;
+}
+export {
+  Q as fit,
+  I as isWebGPUAvailable
+};

package/dist/umap.d.ts

@@ -0,0 +1,43 @@
+export interface UMAPOptions {
+    /** Embedding dimensionality (default: 2) */
+    nComponents?: number;
+    /** Number of nearest neighbors (default: 15) */
+    nNeighbors?: number;
+    /** Number of optimization epochs (default: auto based on dataset size) */
+    nEpochs?: number;
+    /** Minimum distance in the embedding (default: 0.1) */
+    minDist?: number;
+    /** Spread of the embedding (default: 1.0) */
+    spread?: number;
+    /** HNSW index parameters */
+    hnsw?: {
+        M?: number;
+        efConstruction?: number;
+        efSearch?: number;
+    };
+}
+/**
+ * Fit UMAP to the given high-dimensional vectors and return a low-dimensional embedding.
+ *
+ * Pipeline:
+ * 1. HNSW k-nearest neighbor search (O(n log n) via hnswlib-wasm)
+ * 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>;
+/**
+ * Compute the a, b parameters for the UMAP curve 1/(1 + a*d^(2b)).
+ *
+ * For arbitrary minDist/spread values, a proper implementation would use
+ * Levenberg-Marquardt curve fitting. Here we provide pre-fitted constants
+ * for common parameter combinations plus an approximation for others.
+ */
+export declare function findAB(minDist: number, spread: number): {
+    a: number;
+    b: number;
+};
+/**
+ * Compute per-edge epoch sampling periods based on edge weights.
+ * Higher-weight edges are sampled more frequently.
+ */
+export declare function computeEpochsPerSample(weights: Float32Array, nEpochs: number): Float32Array;

package/package.json

@@ -0,0 +1,27 @@
+{
+  "name": "umap-gpu",
+  "version": "0.1.0",
+  "description": "UMAP with HNSW kNN and WebGPU-accelerated SGD",
+  "type": "module",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "files": [
+    "dist"
+  ],
+  "scripts": {
+    "build": "vite build && tsc",
+    "dev": "vite",
+    "test": "vitest run",
+    "prepublishOnly": "npm test && npm run build"
+  },
+  "dependencies": {
+    "hnswlib-wasm": "^0.8.2"
+  },
+  "devDependencies": {
+    "@webgpu/types": "^0.1.40",
+    "typescript": "^5.4.0",
+    "vite": "^5.0.0",
+    "vitest": "^4.0.18"
+  },
+  "license": "MIT"
+}