umap-web 0.2.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/CHANGELOG.md +98 -0
- package/LICENSE +29 -0
- package/NOTICE +19 -0
- package/README.md +205 -0
- package/dist/chunk-Bp6m_JJh.js +13 -0
- package/dist/core-BPpH5_PW.js +2816 -0
- package/dist/core-BPpH5_PW.js.map +1 -0
- package/dist/core-IHdw-3yY.d.ts +373 -0
- package/dist/core.d.ts +3 -0
- package/dist/core.js +4 -0
- package/dist/index.d.ts +261 -0
- package/dist/index.js +2991 -0
- package/dist/index.js.map +1 -0
- package/dist/init-DLiEtGpx.js +802 -0
- package/dist/init-DLiEtGpx.js.map +1 -0
- package/dist/sparse-BxSeVhO4.d.ts +150 -0
- package/dist/webgpu.d.ts +111 -0
- package/dist/webgpu.js +955 -0
- package/dist/webgpu.js.map +1 -0
- package/dist/worker.d.ts +1 -0
- package/dist/worker.js +1095 -0
- package/dist/worker.js.map +1 -0
- package/package.json +76 -0
package/dist/index.js
ADDED
|
@@ -0,0 +1,2991 @@
|
|
|
1
|
+
import { A as validateMetricOptionSizes, B as csrCopy, C as resolveInternalMetric, D as isNamedMetric, E as canonicalMetricName, F as smoothKnnDist, G as csrMaxAbsDiff, H as csrForEach, I as defaultNEpochs, J as csrTranspose, K as csrNnz, L as makeEpochsPerSample, M as fuzzySimplicialSet, N as MIN_K_DIST_SCALE, O as metrics_exports, P as SMOOTH_K_TOLERANCE, R as pruneGraphForEpochs, S as knnFromPrecomputed, T as NAMED_METRICS, U as csrFromCoo, V as csrEliminateZeros, W as csrMapValues, Y as findABParams, _ as floorMod, a as SMALL_DATA_CUTOFF, b as exactKnnDense, c as prepareSearchIndex, d as umapNnDescentParams, f as checkedFlaggedHeapPush, g as Pcg32, h as makeHeap, i as validateCsr, j as computeMembershipStrengths, k as resolveMetric, l as searchKnn, m as deheapSort, n as raceTimeout, o as nearestNeighborsDense, p as checkedHeapPush, q as csrRowNormalizeMax, r as toMatrix, s as nearestNeighborsSparse, t as isCsrInput, u as nnDescent, v as tauRand, w as ANGULAR_METRICS, x as exactKnnSparse, y as tauRandInt, z as csrBinaryOp } from "./core-BPpH5_PW.js";
|
|
2
|
+
import { _ as silentLogger, a as degreeInvSqrt, c as spectralLayout, d as symmetricEigen, f as UmapBackendError, g as consoleLogger, h as UmapValidationError, i as connectedComponents, l as pcaScores, m as UmapError, n as initializeEmbedding, o as lobpcgSmallest, p as UmapConvergenceWarning, r as noisyScaleCoords, t as finalInitRescale, u as mgsOrthonormalize, v as throwIfAborted } from "./init-DLiEtGpx.js";
|
|
3
|
+
|
|
4
|
+
//#region src/layout/sgd.ts
|
|
5
|
+
function clip(val) {
|
|
6
|
+
if (val > 4) return 4;
|
|
7
|
+
if (val < -4) return -4;
|
|
8
|
+
return val;
|
|
9
|
+
}
|
|
10
|
+
/** Per-vertex tau streams: base rng_state + bit-pattern of the vertex's first coord. */
|
|
11
|
+
function perVertexRngStates(n, rngState, embedding, dim) {
|
|
12
|
+
const states = new Int32Array(n * 3);
|
|
13
|
+
const buf = /* @__PURE__ */ new DataView(/* @__PURE__ */ new ArrayBuffer(8));
|
|
14
|
+
for (let j = 0; j < n; j++) {
|
|
15
|
+
buf.setFloat64(0, embedding[j * dim], true);
|
|
16
|
+
const lo = buf.getInt32(0, true);
|
|
17
|
+
const hi = buf.getInt32(4, true);
|
|
18
|
+
states[j * 3] = rngState[0] + lo | 0;
|
|
19
|
+
states[j * 3 + 1] = rngState[1] + hi | 0;
|
|
20
|
+
states[j * 3 + 2] = rngState[2] + (lo ^ hi) | 0;
|
|
21
|
+
}
|
|
22
|
+
return states;
|
|
23
|
+
}
|
|
24
|
+
/**
|
|
25
|
+
* optimize_layout_euclidean. Mutates headEmbedding (and tailEmbedding when
|
|
26
|
+
* moveOther and they alias). Returns the final embedding.
|
|
27
|
+
*/
|
|
28
|
+
function optimizeLayoutEuclidean(headEmbedding, tailEmbedding, dim, head, tail, nEpochs, nVertices, epochsPerSample, a, b, rngState, { gamma = 1, initialAlpha = 1, negativeSampleRate = 5, moveOther = true, densmap = null, onEpoch, signal, snapshotEvery = 0 } = {}) {
|
|
29
|
+
const nEdges = epochsPerSample.length;
|
|
30
|
+
let alpha = initialAlpha;
|
|
31
|
+
const epochsPerNegativeSample = new Float64Array(nEdges);
|
|
32
|
+
for (let i = 0; i < nEdges; i++) epochsPerNegativeSample[i] = epochsPerSample[i] / negativeSampleRate;
|
|
33
|
+
const epochOfNextNegativeSample = epochsPerNegativeSample.slice();
|
|
34
|
+
const epochOfNextSample = epochsPerSample.slice();
|
|
35
|
+
const vertexRng = perVertexRngStates(headEmbedding.length / dim, rngState, headEmbedding, dim);
|
|
36
|
+
let densReSum = null;
|
|
37
|
+
let densPhiSum = null;
|
|
38
|
+
let densMuTot = 0;
|
|
39
|
+
if (densmap) {
|
|
40
|
+
densReSum = new Float32Array(nVertices);
|
|
41
|
+
densPhiSum = new Float32Array(nVertices);
|
|
42
|
+
for (let i = 0; i < densmap.muSum.length; i++) densMuTot += densmap.muSum[i];
|
|
43
|
+
densMuTot /= 2;
|
|
44
|
+
}
|
|
45
|
+
const rngRow = new Int32Array(3);
|
|
46
|
+
for (let n = 0; n < nEpochs; n++) {
|
|
47
|
+
throwIfAborted(signal);
|
|
48
|
+
let densmapFlag = false;
|
|
49
|
+
let densReStd = 0;
|
|
50
|
+
let densReMean = 0;
|
|
51
|
+
let densReCov = 0;
|
|
52
|
+
if (densmap && densmap.lambda > 0 && (n + 1) / nEpochs > 1 - densmap.frac) {
|
|
53
|
+
densmapFlag = true;
|
|
54
|
+
densReSum.fill(0);
|
|
55
|
+
densPhiSum.fill(0);
|
|
56
|
+
for (let i = 0; i < nEdges; i++) {
|
|
57
|
+
const j = head[i];
|
|
58
|
+
const k = tail[i];
|
|
59
|
+
let d2 = 0;
|
|
60
|
+
for (let d = 0; d < dim; d++) {
|
|
61
|
+
const diff = headEmbedding[j * dim + d] - tailEmbedding[k * dim + d];
|
|
62
|
+
d2 += diff * diff;
|
|
63
|
+
}
|
|
64
|
+
const phi = 1 / (1 + a * d2 ** b);
|
|
65
|
+
densReSum[j] = densReSum[j] + phi * d2;
|
|
66
|
+
densReSum[k] = densReSum[k] + phi * d2;
|
|
67
|
+
densPhiSum[j] = densPhiSum[j] + phi;
|
|
68
|
+
densPhiSum[k] = densPhiSum[k] + phi;
|
|
69
|
+
}
|
|
70
|
+
let mean = 0;
|
|
71
|
+
for (let i = 0; i < nVertices; i++) {
|
|
72
|
+
densReSum[i] = Math.log(1e-8 + densReSum[i] / densPhiSum[i]);
|
|
73
|
+
mean += densReSum[i];
|
|
74
|
+
}
|
|
75
|
+
mean /= nVertices;
|
|
76
|
+
let variance = 0;
|
|
77
|
+
let cov = 0;
|
|
78
|
+
for (let i = 0; i < nVertices; i++) {
|
|
79
|
+
variance += (densReSum[i] - mean) ** 2;
|
|
80
|
+
cov += densReSum[i] * densmap.R[i];
|
|
81
|
+
}
|
|
82
|
+
variance /= nVertices;
|
|
83
|
+
densReStd = Math.sqrt(variance + densmap.varShift);
|
|
84
|
+
densReMean = mean;
|
|
85
|
+
densReCov = cov / (nVertices - 1);
|
|
86
|
+
}
|
|
87
|
+
if (dim === 2 && !densmap) {
|
|
88
|
+
const fround = Math.fround;
|
|
89
|
+
const bm1 = b - 1;
|
|
90
|
+
for (let i = 0; i < nEdges; i++) {
|
|
91
|
+
if (epochOfNextSample[i] > n) continue;
|
|
92
|
+
const j = head[i];
|
|
93
|
+
let k = tail[i];
|
|
94
|
+
const jBase = j * 2;
|
|
95
|
+
let kBase = k * 2;
|
|
96
|
+
let cj0 = headEmbedding[jBase];
|
|
97
|
+
let cj1 = headEmbedding[jBase + 1];
|
|
98
|
+
let dx = cj0 - tailEmbedding[kBase];
|
|
99
|
+
let dy = cj1 - tailEmbedding[kBase + 1];
|
|
100
|
+
const d2 = dx * dx + dy * dy;
|
|
101
|
+
if (d2 > 0) {
|
|
102
|
+
const gc = -2 * a * b * d2 ** bm1 / (a * d2 ** b + 1);
|
|
103
|
+
const g0 = clip(gc * dx);
|
|
104
|
+
const g1 = clip(gc * dy);
|
|
105
|
+
cj0 = fround(cj0 + g0 * alpha);
|
|
106
|
+
cj1 = fround(cj1 + g1 * alpha);
|
|
107
|
+
if (moveOther) {
|
|
108
|
+
tailEmbedding[kBase] = tailEmbedding[kBase] - g0 * alpha;
|
|
109
|
+
tailEmbedding[kBase + 1] = tailEmbedding[kBase + 1] - g1 * alpha;
|
|
110
|
+
}
|
|
111
|
+
}
|
|
112
|
+
epochOfNextSample[i] = epochOfNextSample[i] + epochsPerSample[i];
|
|
113
|
+
const nNegSamples = Math.trunc((n - epochOfNextNegativeSample[i]) / epochsPerNegativeSample[i]);
|
|
114
|
+
let s0 = vertexRng[j * 3] | 0;
|
|
115
|
+
let s1 = vertexRng[j * 3 + 1] | 0;
|
|
116
|
+
let s2 = vertexRng[j * 3 + 2] | 0;
|
|
117
|
+
for (let s = 0; s < nNegSamples; s++) {
|
|
118
|
+
const u0 = s0 >>> 0;
|
|
119
|
+
const u1 = s1 >>> 0;
|
|
120
|
+
const u2 = s2 >>> 0;
|
|
121
|
+
s0 = (u0 & 4294967294) << 12 ^ (u0 << 13 >>> 0 ^ u0) >>> 19 | 0;
|
|
122
|
+
s1 = (u1 & 4294967288) << 4 ^ (u1 << 2 >>> 0 ^ u1) >>> 25 | 0;
|
|
123
|
+
s2 = (u2 & 4294967280) << 17 ^ (u2 << 3 >>> 0 ^ u2) >>> 11 | 0;
|
|
124
|
+
const rm = (s0 ^ s1 ^ s2 | 0) % nVertices;
|
|
125
|
+
k = rm + (rm < 0 ? nVertices : 0);
|
|
126
|
+
if (j === k) continue;
|
|
127
|
+
kBase = k * 2;
|
|
128
|
+
dx = cj0 - tailEmbedding[kBase];
|
|
129
|
+
dy = cj1 - tailEmbedding[kBase + 1];
|
|
130
|
+
const dn = dx * dx + dy * dy;
|
|
131
|
+
if (dn > 0) {
|
|
132
|
+
const rc = 2 * gamma * b / ((.001 + dn) * (a * dn ** b + 1));
|
|
133
|
+
cj0 = fround(cj0 + clip(rc * dx) * alpha);
|
|
134
|
+
cj1 = fround(cj1 + clip(rc * dy) * alpha);
|
|
135
|
+
}
|
|
136
|
+
}
|
|
137
|
+
headEmbedding[jBase] = cj0;
|
|
138
|
+
headEmbedding[jBase + 1] = cj1;
|
|
139
|
+
vertexRng[j * 3] = s0;
|
|
140
|
+
vertexRng[j * 3 + 1] = s1;
|
|
141
|
+
vertexRng[j * 3 + 2] = s2;
|
|
142
|
+
epochOfNextNegativeSample[i] = epochOfNextNegativeSample[i] + nNegSamples * epochsPerNegativeSample[i];
|
|
143
|
+
}
|
|
144
|
+
alpha = initialAlpha * (1 - n / nEpochs);
|
|
145
|
+
if (onEpoch) {
|
|
146
|
+
const info = {
|
|
147
|
+
epoch: n,
|
|
148
|
+
nEpochs,
|
|
149
|
+
phase: "layout"
|
|
150
|
+
};
|
|
151
|
+
if (snapshotEvery > 0 && (n % snapshotEvery === 0 || n === nEpochs - 1)) info.embedding = headEmbedding;
|
|
152
|
+
onEpoch(info);
|
|
153
|
+
}
|
|
154
|
+
continue;
|
|
155
|
+
}
|
|
156
|
+
for (let i = 0; i < nEdges; i++) {
|
|
157
|
+
if (epochOfNextSample[i] > n) continue;
|
|
158
|
+
const j = head[i];
|
|
159
|
+
let k = tail[i];
|
|
160
|
+
const jBase = j * dim;
|
|
161
|
+
let kBase = k * dim;
|
|
162
|
+
let distSquared = 0;
|
|
163
|
+
for (let d = 0; d < dim; d++) {
|
|
164
|
+
const diff = headEmbedding[jBase + d] - tailEmbedding[kBase + d];
|
|
165
|
+
distSquared += diff * diff;
|
|
166
|
+
}
|
|
167
|
+
let gradCorCoeff = 0;
|
|
168
|
+
if (densmapFlag) {
|
|
169
|
+
const phi = 1 / (1 + a * distSquared ** b);
|
|
170
|
+
const dphiTerm = a * b * distSquared ** (b - 1) / (1 + a * distSquared ** b);
|
|
171
|
+
const qjk = phi / densPhiSum[k];
|
|
172
|
+
const qkj = phi / densPhiSum[j];
|
|
173
|
+
const drk = qjk * ((1 - b * (1 - phi)) / Math.exp(densReSum[k]) + dphiTerm);
|
|
174
|
+
const drj = qkj * ((1 - b * (1 - phi)) / Math.exp(densReSum[j]) + dphiTerm);
|
|
175
|
+
const reStdSq = densReStd * densReStd;
|
|
176
|
+
const weightK = densmap.R[k] - densReCov * (densReSum[k] - densReMean) / reStdSq;
|
|
177
|
+
const weightJ = densmap.R[j] - densReCov * (densReSum[j] - densReMean) / reStdSq;
|
|
178
|
+
gradCorCoeff = densmap.lambda * densMuTot / (densmap.mu[i] * densReStd) * (weightK * drk + weightJ * drj) / nVertices;
|
|
179
|
+
}
|
|
180
|
+
let gradCoeff;
|
|
181
|
+
if (distSquared > 0) gradCoeff = -2 * a * b * distSquared ** (b - 1) / (a * distSquared ** b + 1);
|
|
182
|
+
else gradCoeff = 0;
|
|
183
|
+
for (let d = 0; d < dim; d++) {
|
|
184
|
+
const diff = headEmbedding[jBase + d] - tailEmbedding[kBase + d];
|
|
185
|
+
let gradD = clip(gradCoeff * diff);
|
|
186
|
+
if (densmapFlag) gradD += clip(2 * gradCorCoeff * diff);
|
|
187
|
+
headEmbedding[jBase + d] = headEmbedding[jBase + d] + gradD * alpha;
|
|
188
|
+
if (moveOther) tailEmbedding[kBase + d] = tailEmbedding[kBase + d] - gradD * alpha;
|
|
189
|
+
}
|
|
190
|
+
epochOfNextSample[i] = epochOfNextSample[i] + epochsPerSample[i];
|
|
191
|
+
const nNegSamples = Math.trunc((n - epochOfNextNegativeSample[i]) / epochsPerNegativeSample[i]);
|
|
192
|
+
rngRow[0] = vertexRng[j * 3];
|
|
193
|
+
rngRow[1] = vertexRng[j * 3 + 1];
|
|
194
|
+
rngRow[2] = vertexRng[j * 3 + 2];
|
|
195
|
+
for (let p = 0; p < nNegSamples; p++) {
|
|
196
|
+
const r = tauRandInt(rngRow);
|
|
197
|
+
k = r % nVertices + (r % nVertices < 0 ? nVertices : 0);
|
|
198
|
+
kBase = k * dim;
|
|
199
|
+
distSquared = 0;
|
|
200
|
+
for (let d = 0; d < dim; d++) {
|
|
201
|
+
const diff = headEmbedding[jBase + d] - tailEmbedding[kBase + d];
|
|
202
|
+
distSquared += diff * diff;
|
|
203
|
+
}
|
|
204
|
+
let repCoeff;
|
|
205
|
+
if (distSquared > 0) repCoeff = 2 * gamma * b / ((.001 + distSquared) * (a * distSquared ** b + 1));
|
|
206
|
+
else if (j === k) continue;
|
|
207
|
+
else repCoeff = 0;
|
|
208
|
+
for (let d = 0; d < dim; d++) {
|
|
209
|
+
let gradD;
|
|
210
|
+
if (repCoeff > 0) {
|
|
211
|
+
const diff = headEmbedding[jBase + d] - tailEmbedding[kBase + d];
|
|
212
|
+
gradD = clip(repCoeff * diff);
|
|
213
|
+
} else gradD = 0;
|
|
214
|
+
headEmbedding[jBase + d] = headEmbedding[jBase + d] + gradD * alpha;
|
|
215
|
+
}
|
|
216
|
+
}
|
|
217
|
+
vertexRng[j * 3] = rngRow[0];
|
|
218
|
+
vertexRng[j * 3 + 1] = rngRow[1];
|
|
219
|
+
vertexRng[j * 3 + 2] = rngRow[2];
|
|
220
|
+
epochOfNextNegativeSample[i] = epochOfNextNegativeSample[i] + nNegSamples * epochsPerNegativeSample[i];
|
|
221
|
+
}
|
|
222
|
+
alpha = initialAlpha * (1 - n / nEpochs);
|
|
223
|
+
if (onEpoch) {
|
|
224
|
+
const info = {
|
|
225
|
+
epoch: n,
|
|
226
|
+
nEpochs,
|
|
227
|
+
phase: "layout"
|
|
228
|
+
};
|
|
229
|
+
if (snapshotEvery > 0 && (n % snapshotEvery === 0 || n === nEpochs - 1)) info.embedding = headEmbedding;
|
|
230
|
+
onEpoch(info);
|
|
231
|
+
}
|
|
232
|
+
}
|
|
233
|
+
return headEmbedding;
|
|
234
|
+
}
|
|
235
|
+
/**
|
|
236
|
+
* optimize_layout_inverse (layouts.py:663-732; notes §13): optimizes DATA-space
|
|
237
|
+
* positions of inverse-transformed points against the raw training data.
|
|
238
|
+
* Negative samples use the SHARED rng state (not per-vertex).
|
|
239
|
+
*/
|
|
240
|
+
function optimizeLayoutInverse(headEmbedding, tailEmbedding, dim, head, tail, weight, sigmas, rhos, nEpochs, nVertices, epochsPerSample, rngState, outputMetric, { gamma = 1, initialAlpha = 1, negativeSampleRate = 5, signal } = {}) {
|
|
241
|
+
const nEdges = epochsPerSample.length;
|
|
242
|
+
let alpha = initialAlpha;
|
|
243
|
+
const epochsPerNegativeSample = new Float64Array(nEdges);
|
|
244
|
+
for (let i = 0; i < nEdges; i++) epochsPerNegativeSample[i] = epochsPerSample[i] / negativeSampleRate;
|
|
245
|
+
const epochOfNextNegativeSample = epochsPerNegativeSample.slice();
|
|
246
|
+
const epochOfNextSample = epochsPerSample.slice();
|
|
247
|
+
const grad = new Float64Array(dim);
|
|
248
|
+
const cur = new Float32Array(dim);
|
|
249
|
+
const oth = new Float32Array(dim);
|
|
250
|
+
for (let n = 0; n < nEpochs; n++) {
|
|
251
|
+
throwIfAborted(signal);
|
|
252
|
+
for (let i = 0; i < nEdges; i++) {
|
|
253
|
+
if (epochOfNextSample[i] > n) continue;
|
|
254
|
+
const j = head[i];
|
|
255
|
+
let k = tail[i];
|
|
256
|
+
const jBase = j * dim;
|
|
257
|
+
let kBase = k * dim;
|
|
258
|
+
for (let d = 0; d < dim; d++) {
|
|
259
|
+
cur[d] = headEmbedding[jBase + d];
|
|
260
|
+
oth[d] = tailEmbedding[kBase + d];
|
|
261
|
+
}
|
|
262
|
+
outputMetric(cur, oth, grad);
|
|
263
|
+
const gradCoeff = -(1 / (weight[i] * sigmas[k] + 1e-6));
|
|
264
|
+
for (let d = 0; d < dim; d++) {
|
|
265
|
+
const gradD = clip(gradCoeff * grad[d]);
|
|
266
|
+
headEmbedding[jBase + d] = headEmbedding[jBase + d] + gradD * alpha;
|
|
267
|
+
}
|
|
268
|
+
epochOfNextSample[i] = epochOfNextSample[i] + epochsPerSample[i];
|
|
269
|
+
const nNegSamples = Math.trunc((n - epochOfNextNegativeSample[i]) / epochsPerNegativeSample[i]);
|
|
270
|
+
for (let p = 0; p < nNegSamples; p++) {
|
|
271
|
+
const r = tauRandInt(rngState);
|
|
272
|
+
k = r % nVertices + (r % nVertices < 0 ? nVertices : 0);
|
|
273
|
+
kBase = k * dim;
|
|
274
|
+
for (let d = 0; d < dim; d++) {
|
|
275
|
+
cur[d] = headEmbedding[jBase + d];
|
|
276
|
+
oth[d] = tailEmbedding[kBase + d];
|
|
277
|
+
}
|
|
278
|
+
const distOutput = outputMetric(cur, oth, grad);
|
|
279
|
+
const wH = Math.exp(-Math.max(distOutput - rhos[k], 1e-6) / (sigmas[k] + 1e-6));
|
|
280
|
+
let denom = (1 - wH) * sigmas[k] + 1e-6;
|
|
281
|
+
if (!Number.isFinite(denom)) denom = 1e-6;
|
|
282
|
+
const gradCoeffNeg = -gamma * ((0 - wH) / denom);
|
|
283
|
+
for (let d = 0; d < dim; d++) {
|
|
284
|
+
const gradD = clip(gradCoeffNeg * grad[d]);
|
|
285
|
+
headEmbedding[jBase + d] = headEmbedding[jBase + d] + gradD * alpha;
|
|
286
|
+
}
|
|
287
|
+
}
|
|
288
|
+
epochOfNextNegativeSample[i] = epochOfNextNegativeSample[i] + nNegSamples * epochsPerNegativeSample[i];
|
|
289
|
+
}
|
|
290
|
+
alpha = initialAlpha * (1 - n / nEpochs);
|
|
291
|
+
}
|
|
292
|
+
return headEmbedding;
|
|
293
|
+
}
|
|
294
|
+
/** optimize_layout_generic — non-euclidean output metrics (notes §10). */
|
|
295
|
+
function optimizeLayoutGeneric(headEmbedding, tailEmbedding, dim, head, tail, nEpochs, nVertices, epochsPerSample, a, b, rngState, outputMetric, { gamma = 1, initialAlpha = 1, negativeSampleRate = 5, moveOther = true, onEpoch, signal, snapshotEvery = 0 } = {}) {
|
|
296
|
+
const nEdges = epochsPerSample.length;
|
|
297
|
+
let alpha = initialAlpha;
|
|
298
|
+
const epochsPerNegativeSample = new Float64Array(nEdges);
|
|
299
|
+
for (let i = 0; i < nEdges; i++) epochsPerNegativeSample[i] = epochsPerSample[i] / negativeSampleRate;
|
|
300
|
+
const epochOfNextNegativeSample = epochsPerNegativeSample.slice();
|
|
301
|
+
const epochOfNextSample = epochsPerSample.slice();
|
|
302
|
+
const vertexRng = perVertexRngStates(headEmbedding.length / dim, rngState, headEmbedding, dim);
|
|
303
|
+
const grad = new Float64Array(dim);
|
|
304
|
+
const revGrad = new Float64Array(dim);
|
|
305
|
+
const rngRow = new Int32Array(3);
|
|
306
|
+
const cur = new Float32Array(dim);
|
|
307
|
+
const oth = new Float32Array(dim);
|
|
308
|
+
for (let n = 0; n < nEpochs; n++) {
|
|
309
|
+
throwIfAborted(signal);
|
|
310
|
+
for (let i = 0; i < nEdges; i++) {
|
|
311
|
+
if (epochOfNextSample[i] > n) continue;
|
|
312
|
+
const j = head[i];
|
|
313
|
+
let k = tail[i];
|
|
314
|
+
const jBase = j * dim;
|
|
315
|
+
let kBase = k * dim;
|
|
316
|
+
for (let d = 0; d < dim; d++) {
|
|
317
|
+
cur[d] = headEmbedding[jBase + d];
|
|
318
|
+
oth[d] = tailEmbedding[kBase + d];
|
|
319
|
+
}
|
|
320
|
+
const distOutput = outputMetric(cur, oth, grad);
|
|
321
|
+
outputMetric(oth, cur, revGrad);
|
|
322
|
+
const wL = distOutput > 0 ? (1 + a * distOutput ** (2 * b)) ** -1 : 1;
|
|
323
|
+
const gradCoeff = 2 * b * (wL - 1) / (distOutput + 1e-6);
|
|
324
|
+
for (let d = 0; d < dim; d++) {
|
|
325
|
+
const gradD = clip(gradCoeff * grad[d]);
|
|
326
|
+
headEmbedding[jBase + d] = headEmbedding[jBase + d] + gradD * alpha;
|
|
327
|
+
if (moveOther) {
|
|
328
|
+
const revD = clip(gradCoeff * revGrad[d]);
|
|
329
|
+
tailEmbedding[kBase + d] = tailEmbedding[kBase + d] + revD * alpha;
|
|
330
|
+
}
|
|
331
|
+
}
|
|
332
|
+
epochOfNextSample[i] = epochOfNextSample[i] + epochsPerSample[i];
|
|
333
|
+
const nNegSamples = Math.trunc((n - epochOfNextNegativeSample[i]) / epochsPerNegativeSample[i]);
|
|
334
|
+
rngRow[0] = vertexRng[j * 3];
|
|
335
|
+
rngRow[1] = vertexRng[j * 3 + 1];
|
|
336
|
+
rngRow[2] = vertexRng[j * 3 + 2];
|
|
337
|
+
for (let p = 0; p < nNegSamples; p++) {
|
|
338
|
+
const r = tauRandInt(rngRow);
|
|
339
|
+
k = r % nVertices + (r % nVertices < 0 ? nVertices : 0);
|
|
340
|
+
kBase = k * dim;
|
|
341
|
+
for (let d = 0; d < dim; d++) {
|
|
342
|
+
cur[d] = headEmbedding[jBase + d];
|
|
343
|
+
oth[d] = tailEmbedding[kBase + d];
|
|
344
|
+
}
|
|
345
|
+
const dNeg = outputMetric(cur, oth, grad);
|
|
346
|
+
let wLNeg;
|
|
347
|
+
if (dNeg > 0) wLNeg = (1 + a * dNeg ** (2 * b)) ** -1;
|
|
348
|
+
else if (j === k) continue;
|
|
349
|
+
else wLNeg = 1;
|
|
350
|
+
const repCoeff = gamma * 2 * b * wLNeg / (dNeg + 1e-6);
|
|
351
|
+
for (let d = 0; d < dim; d++) {
|
|
352
|
+
const gradD = clip(repCoeff * grad[d]);
|
|
353
|
+
headEmbedding[jBase + d] = headEmbedding[jBase + d] + gradD * alpha;
|
|
354
|
+
}
|
|
355
|
+
}
|
|
356
|
+
vertexRng[j * 3] = rngRow[0];
|
|
357
|
+
vertexRng[j * 3 + 1] = rngRow[1];
|
|
358
|
+
vertexRng[j * 3 + 2] = rngRow[2];
|
|
359
|
+
epochOfNextNegativeSample[i] = epochOfNextNegativeSample[i] + nNegSamples * epochsPerNegativeSample[i];
|
|
360
|
+
}
|
|
361
|
+
alpha = initialAlpha * (1 - n / nEpochs);
|
|
362
|
+
if (onEpoch) {
|
|
363
|
+
const info = {
|
|
364
|
+
epoch: n,
|
|
365
|
+
nEpochs,
|
|
366
|
+
phase: "layout"
|
|
367
|
+
};
|
|
368
|
+
if (snapshotEvery > 0 && (n % snapshotEvery === 0 || n === nEpochs - 1)) info.embedding = headEmbedding;
|
|
369
|
+
onEpoch(info);
|
|
370
|
+
}
|
|
371
|
+
}
|
|
372
|
+
return headEmbedding;
|
|
373
|
+
}
|
|
374
|
+
|
|
375
|
+
//#endregion
|
|
376
|
+
//#region src/backends/cpu/pool.ts
|
|
377
|
+
const IS_NODE = typeof process !== "undefined" && !!process.versions?.node;
|
|
378
|
+
/**
|
|
379
|
+
* Runtime-only import of node builtins. The variable specifier plus the ignore
|
|
380
|
+
* comments keep browser bundlers (esbuild/Vite/webpack) from trying to resolve
|
|
381
|
+
* `node:*` modules they cannot ship; every call site is behind an IS_NODE
|
|
382
|
+
* guard so the import never evaluates in a browser.
|
|
383
|
+
*/
|
|
384
|
+
function nodeImport(specifier) {
|
|
385
|
+
return import(
|
|
386
|
+
/* @vite-ignore */
|
|
387
|
+
/* webpackIgnore: true */
|
|
388
|
+
specifier
|
|
389
|
+
);
|
|
390
|
+
}
|
|
391
|
+
function sabAvailable() {
|
|
392
|
+
if (typeof SharedArrayBuffer === "undefined") return false;
|
|
393
|
+
if (IS_NODE) return true;
|
|
394
|
+
return globalThis.crossOriginIsolated === true;
|
|
395
|
+
}
|
|
396
|
+
function hardwareConcurrency() {
|
|
397
|
+
return Math.max(1, globalThis.navigator?.hardwareConcurrency ?? 4);
|
|
398
|
+
}
|
|
399
|
+
/**
|
|
400
|
+
* Bundler-visible browser worker factory. webpack 5 and Vite (build) statically
|
|
401
|
+
* detect the exact inline `new Worker(new URL('./worker.js', import.meta.url),
|
|
402
|
+
* {type:'module'})` expression — including inside node_modules — and emit the
|
|
403
|
+
* worker as an asset of the consumer's build. Do NOT extract the URL into a
|
|
404
|
+
* variable or helper: detection depends on the literal shape of this
|
|
405
|
+
* expression (guarded by scripts/check-bundlers.mjs).
|
|
406
|
+
*/
|
|
407
|
+
function defaultBrowserWorker() {
|
|
408
|
+
return new Worker(new URL("./worker.js", import.meta.url), { type: "module" });
|
|
409
|
+
}
|
|
410
|
+
/**
|
|
411
|
+
* Cross-origin retry (CDN / import-map consumers): fetch the sibling worker
|
|
412
|
+
* source and spawn it from a blob: URL. Valid only because dist/worker.js is
|
|
413
|
+
* built fully self-contained (no sibling-chunk imports; enforced by
|
|
414
|
+
* scripts/check-worker-selfcontained.mjs). Requires `worker-src blob:` under
|
|
415
|
+
* a strict CSP (documented in README).
|
|
416
|
+
*/
|
|
417
|
+
async function blobWorkerUrl() {
|
|
418
|
+
try {
|
|
419
|
+
const res = await fetch(new URL("./worker.js", import.meta.url));
|
|
420
|
+
if (!res.ok) return null;
|
|
421
|
+
const text = await res.text();
|
|
422
|
+
return URL.createObjectURL(new Blob([text], { type: "text/javascript" }));
|
|
423
|
+
} catch {
|
|
424
|
+
return null;
|
|
425
|
+
}
|
|
426
|
+
}
|
|
427
|
+
async function spawnWorker(source) {
|
|
428
|
+
if (IS_NODE) {
|
|
429
|
+
const { Worker: Worker$1 } = await nodeImport("node:worker_threads");
|
|
430
|
+
const resolved = source ?? await workerEntryUrl();
|
|
431
|
+
const w$1 = new Worker$1(typeof resolved === "string" && resolved.startsWith("file:") ? new URL(resolved) : resolved);
|
|
432
|
+
return {
|
|
433
|
+
post: (msg, transfers) => w$1.postMessage(msg, transfers ?? []),
|
|
434
|
+
onMessage: (cb) => {
|
|
435
|
+
w$1.on("message", cb);
|
|
436
|
+
},
|
|
437
|
+
onFailure: (cb) => {
|
|
438
|
+
w$1.on("error", (e) => cb(e instanceof Error ? e : new Error(String(e))));
|
|
439
|
+
w$1.on("messageerror", (e) => cb(/* @__PURE__ */ new Error(`worker messageerror: ${e}`)));
|
|
440
|
+
w$1.on("exit", (code) => {
|
|
441
|
+
if (code !== 0) cb(/* @__PURE__ */ new Error(`worker exited with code ${code}`));
|
|
442
|
+
});
|
|
443
|
+
},
|
|
444
|
+
terminate: async () => {
|
|
445
|
+
await w$1.terminate();
|
|
446
|
+
}
|
|
447
|
+
};
|
|
448
|
+
}
|
|
449
|
+
const w = source === void 0 ? defaultBrowserWorker() : new Worker(source, { type: "module" });
|
|
450
|
+
return {
|
|
451
|
+
post: (msg, transfers) => w.postMessage(msg, { transfer: transfers ?? [] }),
|
|
452
|
+
onMessage: (cb) => {
|
|
453
|
+
w.onmessage = (e) => cb(e.data);
|
|
454
|
+
},
|
|
455
|
+
onFailure: (cb) => {
|
|
456
|
+
w.onerror = (e) => cb(/* @__PURE__ */ new Error(`worker error: ${e.message || "failed to load worker.js"}`));
|
|
457
|
+
w.onmessageerror = () => cb(/* @__PURE__ */ new Error("worker messageerror (unserializable message)"));
|
|
458
|
+
},
|
|
459
|
+
terminate: () => w.terminate()
|
|
460
|
+
};
|
|
461
|
+
}
|
|
462
|
+
var WorkerPool = class WorkerPool {
|
|
463
|
+
workers = [];
|
|
464
|
+
pending = /* @__PURE__ */ new Map();
|
|
465
|
+
nextId = 1;
|
|
466
|
+
failure = null;
|
|
467
|
+
terminating = false;
|
|
468
|
+
cleanup = null;
|
|
469
|
+
size;
|
|
470
|
+
constructor(size) {
|
|
471
|
+
this.size = size;
|
|
472
|
+
}
|
|
473
|
+
/**
|
|
474
|
+
* Spawn `size` workers from `source` (undefined → platform default resolution).
|
|
475
|
+
* `cleanup` runs once at terminate() (e.g. revoking a blob: URL). Partial
|
|
476
|
+
* spawn failures terminate the already-spawned workers before rethrowing.
|
|
477
|
+
*/
|
|
478
|
+
static async create(size, source, cleanup) {
|
|
479
|
+
const pool = new WorkerPool(size);
|
|
480
|
+
pool.cleanup = cleanup ?? null;
|
|
481
|
+
try {
|
|
482
|
+
for (let i = 0; i < size; i++) {
|
|
483
|
+
const w = await spawnWorker(source);
|
|
484
|
+
w.onMessage((msg) => pool.dispatch(msg));
|
|
485
|
+
w.onFailure((err) => pool.markFailed(err));
|
|
486
|
+
pool.workers.push(w);
|
|
487
|
+
}
|
|
488
|
+
} catch (e) {
|
|
489
|
+
await pool.terminate();
|
|
490
|
+
throw e;
|
|
491
|
+
}
|
|
492
|
+
return pool;
|
|
493
|
+
}
|
|
494
|
+
/** A worker died: poison the pool and reject everything in flight. */
|
|
495
|
+
markFailed(err) {
|
|
496
|
+
if (this.terminating || this.failure) return;
|
|
497
|
+
this.failure = err;
|
|
498
|
+
const inflight = [...this.pending.values()];
|
|
499
|
+
this.pending.clear();
|
|
500
|
+
const wrapped = new UmapBackendError(`worker failed: ${err.message}`);
|
|
501
|
+
for (const p of inflight) p.reject(wrapped);
|
|
502
|
+
}
|
|
503
|
+
dispatch(msg) {
|
|
504
|
+
const m = msg;
|
|
505
|
+
if (typeof m?.id !== "number") return;
|
|
506
|
+
const p = this.pending.get(m.id);
|
|
507
|
+
if (!p) return;
|
|
508
|
+
this.pending.delete(m.id);
|
|
509
|
+
if (m.error) p.reject(new UmapBackendError(`worker error: ${m.error}`));
|
|
510
|
+
else p.resolve(m.result);
|
|
511
|
+
}
|
|
512
|
+
call(worker, method, args, transfers) {
|
|
513
|
+
if (this.failure) return Promise.reject(new UmapBackendError(`worker pool is dead: ${this.failure.message}`));
|
|
514
|
+
if (this.terminating) return Promise.reject(new UmapBackendError("worker pool has been terminated"));
|
|
515
|
+
const id = this.nextId++;
|
|
516
|
+
const promise = new Promise((resolve, reject) => {
|
|
517
|
+
this.pending.set(id, {
|
|
518
|
+
resolve,
|
|
519
|
+
reject
|
|
520
|
+
});
|
|
521
|
+
});
|
|
522
|
+
this.workers[worker].post({
|
|
523
|
+
id,
|
|
524
|
+
method,
|
|
525
|
+
args
|
|
526
|
+
}, transfers);
|
|
527
|
+
return promise;
|
|
528
|
+
}
|
|
529
|
+
/** Run `method` on every worker with per-worker args; await all. */
|
|
530
|
+
mapAll(method, argsFor, transfersFor) {
|
|
531
|
+
return Promise.all(this.workers.map((_, w) => this.call(w, method, argsFor(w), transfersFor?.(w))));
|
|
532
|
+
}
|
|
533
|
+
async terminate() {
|
|
534
|
+
if (this.terminating) return;
|
|
535
|
+
this.terminating = true;
|
|
536
|
+
const inflight = [...this.pending.values()];
|
|
537
|
+
this.pending.clear();
|
|
538
|
+
for (const p of inflight) p.reject(new UmapBackendError("worker pool terminated with calls in flight"));
|
|
539
|
+
await Promise.all(this.workers.map((w) => w.terminate()));
|
|
540
|
+
this.workers = [];
|
|
541
|
+
if (this.cleanup) {
|
|
542
|
+
this.cleanup();
|
|
543
|
+
this.cleanup = null;
|
|
544
|
+
}
|
|
545
|
+
}
|
|
546
|
+
};
|
|
547
|
+
/**
|
|
548
|
+
* Node-only resolution of the worker entry: the bundled sibling (dist layout),
|
|
549
|
+
* then the package export (bundled-consumer deploys that keep node_modules),
|
|
550
|
+
* then the package's dist for the src/test context. Browser resolution happens
|
|
551
|
+
* in defaultBrowserWorker/blobWorkerUrl instead.
|
|
552
|
+
*/
|
|
553
|
+
async function workerEntryUrl() {
|
|
554
|
+
const local = new URL("./worker.js", import.meta.url);
|
|
555
|
+
if (IS_NODE && local.protocol === "file:") {
|
|
556
|
+
const { existsSync } = await nodeImport("node:fs");
|
|
557
|
+
const { fileURLToPath } = await nodeImport("node:url");
|
|
558
|
+
if (!existsSync(fileURLToPath(local))) {
|
|
559
|
+
try {
|
|
560
|
+
const viaExport = import.meta.resolve?.("umap-web/worker");
|
|
561
|
+
if (viaExport?.startsWith("file:") && existsSync(fileURLToPath(new URL(viaExport)))) return new URL(viaExport);
|
|
562
|
+
} catch {}
|
|
563
|
+
return new URL("../../../dist/worker.js", import.meta.url);
|
|
564
|
+
}
|
|
565
|
+
}
|
|
566
|
+
return local;
|
|
567
|
+
}
|
|
568
|
+
|
|
569
|
+
//#endregion
|
|
570
|
+
//#region src/backends/cpu/workers.ts
|
|
571
|
+
/** Copy data into a SAB-backed view when available. */
|
|
572
|
+
function shareData(X, sab) {
|
|
573
|
+
if (!sab) return X.data;
|
|
574
|
+
const shared = new Float32Array(new SharedArrayBuffer(X.data.length * 4));
|
|
575
|
+
shared.set(X.data);
|
|
576
|
+
return shared;
|
|
577
|
+
}
|
|
578
|
+
/** Cap on worker spawn + init; a stuck worker must degrade, never hang fit(). */
|
|
579
|
+
const WORKER_INIT_TIMEOUT_MS = 1e4;
|
|
580
|
+
async function tryCreatePool(size, source, cleanup, initArgs, failures) {
|
|
581
|
+
let pool = null;
|
|
582
|
+
try {
|
|
583
|
+
pool = await WorkerPool.create(size, source, cleanup);
|
|
584
|
+
const init = pool.mapAll("init", () => initArgs);
|
|
585
|
+
init.catch(() => {});
|
|
586
|
+
await raceTimeout(init, WORKER_INIT_TIMEOUT_MS, "worker init");
|
|
587
|
+
return pool;
|
|
588
|
+
} catch (e) {
|
|
589
|
+
failures.push(e instanceof Error ? e.message : String(e));
|
|
590
|
+
if (pool) await pool.terminate();
|
|
591
|
+
return null;
|
|
592
|
+
}
|
|
593
|
+
}
|
|
594
|
+
async function createParallelContext(X, metricName, metricOptions, requested, workerUrl, logger) {
|
|
595
|
+
const hw = hardwareConcurrency();
|
|
596
|
+
const size = requested === "auto" ? Math.max(1, Math.min(hw - 1, 16)) : requested;
|
|
597
|
+
if (size <= 1) return null;
|
|
598
|
+
const sab = sabAvailable();
|
|
599
|
+
const data = shareData(X, sab);
|
|
600
|
+
const initArgs = {
|
|
601
|
+
data,
|
|
602
|
+
rows: X.rows,
|
|
603
|
+
cols: X.cols,
|
|
604
|
+
metricName,
|
|
605
|
+
metricOptions
|
|
606
|
+
};
|
|
607
|
+
const failures = [];
|
|
608
|
+
let pool = null;
|
|
609
|
+
if (workerUrl !== void 0) pool = await tryCreatePool(size, workerUrl, void 0, initArgs, failures);
|
|
610
|
+
else {
|
|
611
|
+
pool = await tryCreatePool(size, void 0, void 0, initArgs, failures);
|
|
612
|
+
if (!pool && !IS_NODE) {
|
|
613
|
+
const blob = await blobWorkerUrl();
|
|
614
|
+
if (blob) pool = await tryCreatePool(size, blob, () => URL.revokeObjectURL(blob), initArgs, failures);
|
|
615
|
+
else failures.push("worker source fetch for blob retry failed");
|
|
616
|
+
}
|
|
617
|
+
}
|
|
618
|
+
if (!pool) {
|
|
619
|
+
logger?.warn(`worker pool unavailable (${failures.join("; ")}) — continuing single-threaded. Bundled apps: see the README section "Workers & bundlers" or pass options.workerUrl.`);
|
|
620
|
+
return null;
|
|
621
|
+
}
|
|
622
|
+
const internal = resolveInternalMetric(metricName, metricOptions);
|
|
623
|
+
return {
|
|
624
|
+
pool,
|
|
625
|
+
sab,
|
|
626
|
+
data,
|
|
627
|
+
rows: X.rows,
|
|
628
|
+
cols: X.cols,
|
|
629
|
+
dist: (a, b) => internal.dist(a, b)
|
|
630
|
+
};
|
|
631
|
+
}
|
|
632
|
+
function splitRange(total, parts) {
|
|
633
|
+
const out = [];
|
|
634
|
+
const step = Math.ceil(total / parts);
|
|
635
|
+
for (let i = 0; i < parts; i++) out.push([Math.min(i * step, total), Math.min((i + 1) * step, total)]);
|
|
636
|
+
return out;
|
|
637
|
+
}
|
|
638
|
+
function applyTriples(heap, buffers) {
|
|
639
|
+
let c = 0;
|
|
640
|
+
for (const buf of buffers) for (let t = 0; t < buf.length; t += 3) {
|
|
641
|
+
const p = buf[t];
|
|
642
|
+
const q = buf[t + 1];
|
|
643
|
+
const d = buf[t + 2];
|
|
644
|
+
c += checkedFlaggedHeapPush(heap, p, d, q, 1);
|
|
645
|
+
c += checkedFlaggedHeapPush(heap, q, d, p, 1);
|
|
646
|
+
}
|
|
647
|
+
return c;
|
|
648
|
+
}
|
|
649
|
+
function thresholdSnapshot(heap, out) {
|
|
650
|
+
const k = heap.size;
|
|
651
|
+
for (let i = 0; i < heap.n; i++) out[i] = heap.distances[i * k];
|
|
652
|
+
return out;
|
|
653
|
+
}
|
|
654
|
+
/**
|
|
655
|
+
* Parallel NN-descent (same candidate scheme/acceptance as the single-thread
|
|
656
|
+
* port; distance passes fan out to workers per block).
|
|
657
|
+
*/
|
|
658
|
+
async function parallelNnDescent(ctx, nNeighbors, rng, angular, { maxCandidates = 60, delta = .001 } = {}) {
|
|
659
|
+
const n = ctx.rows;
|
|
660
|
+
const { nTrees, nIters } = umapNnDescentParams(n);
|
|
661
|
+
const rngState = rng.tauState();
|
|
662
|
+
const searchRngState = rng.tauState();
|
|
663
|
+
for (let i = 0; i < 10; i++) tauRandInt(searchRngState);
|
|
664
|
+
const treeStates = [];
|
|
665
|
+
for (let t = 0; t < nTrees; t++) treeStates.push(rng.tauState());
|
|
666
|
+
const perWorker = splitRange(nTrees, ctx.pool.size);
|
|
667
|
+
const parts = (await ctx.pool.mapAll("buildTrees", (w) => ({
|
|
668
|
+
treeStates: treeStates.slice(perWorker[w][0], perWorker[w][1]),
|
|
669
|
+
nNeighbors,
|
|
670
|
+
angular
|
|
671
|
+
}))).filter((p) => p.rows > 0 && p.cols > 0);
|
|
672
|
+
const maxCols = Math.max(1, ...parts.map((p) => p.cols));
|
|
673
|
+
const totalRows = parts.reduce((s, p) => s + p.rows, 0);
|
|
674
|
+
const leafArray = new Int32Array(totalRows * maxCols).fill(-1);
|
|
675
|
+
{
|
|
676
|
+
let r = 0;
|
|
677
|
+
for (const p of parts) for (let i = 0; i < p.rows; i++) {
|
|
678
|
+
leafArray.set(p.data.subarray(i * p.cols, (i + 1) * p.cols), r * maxCols);
|
|
679
|
+
r++;
|
|
680
|
+
}
|
|
681
|
+
}
|
|
682
|
+
const heap = makeHeap(n, nNeighbors);
|
|
683
|
+
const thresholds = new Float32Array(n);
|
|
684
|
+
const LEAF_BLOCK = 4096;
|
|
685
|
+
for (let start = 0; start < totalRows; start += LEAF_BLOCK) {
|
|
686
|
+
const end = Math.min(totalRows, start + LEAF_BLOCK);
|
|
687
|
+
thresholdSnapshot(heap, thresholds);
|
|
688
|
+
const shards = splitRange(end - start, ctx.pool.size);
|
|
689
|
+
applyTriples(heap, await ctx.pool.mapAll("leafUpdates", (w) => ({
|
|
690
|
+
leafArray,
|
|
691
|
+
cols: maxCols,
|
|
692
|
+
rowStart: start + shards[w][0],
|
|
693
|
+
rowEnd: start + shards[w][1],
|
|
694
|
+
thresholds
|
|
695
|
+
})));
|
|
696
|
+
}
|
|
697
|
+
{
|
|
698
|
+
const d = ctx.cols;
|
|
699
|
+
const dist = (i, j) => ctx.dist(ctx.data.subarray(i * d, (i + 1) * d), ctx.data.subarray(j * d, (j + 1) * d));
|
|
700
|
+
const k = nNeighbors;
|
|
701
|
+
for (let i = 0; i < n; i++) if (heap.indices[i * k] < 0) {
|
|
702
|
+
let filled = 0;
|
|
703
|
+
for (let j = 0; j < k; j++) if (heap.indices[i * k + j] >= 0) filled++;
|
|
704
|
+
for (let j = 0; j < k - filled; j++) {
|
|
705
|
+
const idx = Math.abs(tauRandInt(rngState)) % n;
|
|
706
|
+
checkedFlaggedHeapPush(heap, i, dist(idx, i), idx, 1);
|
|
707
|
+
}
|
|
708
|
+
}
|
|
709
|
+
}
|
|
710
|
+
const VBLOCK = 16384;
|
|
711
|
+
for (let iter = 0; iter < nIters; iter++) {
|
|
712
|
+
const { newIdx, oldIdx } = buildCandidatesMain(heap, maxCandidates, rngState);
|
|
713
|
+
let c = 0;
|
|
714
|
+
for (let vs = 0; vs < n; vs += VBLOCK) {
|
|
715
|
+
const ve = Math.min(n, vs + VBLOCK);
|
|
716
|
+
thresholdSnapshot(heap, thresholds);
|
|
717
|
+
const shards = splitRange(ve - vs, ctx.pool.size);
|
|
718
|
+
const results = await ctx.pool.mapAll("joinUpdates", (w) => ({
|
|
719
|
+
newIdx: newIdx.subarray((vs + shards[w][0]) * maxCandidates, (vs + shards[w][1]) * maxCandidates),
|
|
720
|
+
oldIdx: oldIdx.subarray((vs + shards[w][0]) * maxCandidates, (vs + shards[w][1]) * maxCandidates),
|
|
721
|
+
maxCandidates,
|
|
722
|
+
vStart: 0,
|
|
723
|
+
vEnd: shards[w][1] - shards[w][0],
|
|
724
|
+
thresholds
|
|
725
|
+
}));
|
|
726
|
+
c += applyTriples(heap, results);
|
|
727
|
+
}
|
|
728
|
+
if (c <= delta * nNeighbors * n) break;
|
|
729
|
+
}
|
|
730
|
+
deheapSort(heap);
|
|
731
|
+
return {
|
|
732
|
+
knnInternal: {
|
|
733
|
+
indices: heap.indices,
|
|
734
|
+
distances: heap.distances,
|
|
735
|
+
k: nNeighbors
|
|
736
|
+
},
|
|
737
|
+
searchRngState
|
|
738
|
+
};
|
|
739
|
+
}
|
|
740
|
+
/** new_build_candidates on the main thread (heap pushes are cheap). */
|
|
741
|
+
function buildCandidatesMain(heap, maxCandidates, masterRng) {
|
|
742
|
+
const n = heap.n;
|
|
743
|
+
const k = heap.size;
|
|
744
|
+
const newIdx = new Int32Array(n * maxCandidates).fill(-1);
|
|
745
|
+
const newPri = new Float32Array(n * maxCandidates).fill(Infinity);
|
|
746
|
+
const oldIdx = new Int32Array(n * maxCandidates).fill(-1);
|
|
747
|
+
const oldPri = new Float32Array(n * maxCandidates).fill(Infinity);
|
|
748
|
+
const localRng = Int32Array.from(masterRng);
|
|
749
|
+
for (let i = 0; i < n; i++) for (let j = 0; j < k; j++) {
|
|
750
|
+
const idx = heap.indices[i * k + j];
|
|
751
|
+
if (idx < 0) continue;
|
|
752
|
+
const d = tauRand(localRng);
|
|
753
|
+
if (heap.flags[i * k + j] === 1) {
|
|
754
|
+
checkedHeapPush(newPri, newIdx, i * maxCandidates, maxCandidates, d, idx);
|
|
755
|
+
checkedHeapPush(newPri, newIdx, idx * maxCandidates, maxCandidates, d, i);
|
|
756
|
+
} else {
|
|
757
|
+
checkedHeapPush(oldPri, oldIdx, i * maxCandidates, maxCandidates, d, idx);
|
|
758
|
+
checkedHeapPush(oldPri, oldIdx, idx * maxCandidates, maxCandidates, d, i);
|
|
759
|
+
}
|
|
760
|
+
}
|
|
761
|
+
for (let i = 0; i < n; i++) for (let j = 0; j < k; j++) {
|
|
762
|
+
const idx = heap.indices[i * k + j];
|
|
763
|
+
if (idx < 0 || heap.flags[i * k + j] !== 1) continue;
|
|
764
|
+
const cbase = i * maxCandidates;
|
|
765
|
+
for (let c = 0; c < maxCandidates; c++) if (newIdx[cbase + c] === idx) {
|
|
766
|
+
heap.flags[i * k + j] = 0;
|
|
767
|
+
break;
|
|
768
|
+
}
|
|
769
|
+
}
|
|
770
|
+
return {
|
|
771
|
+
newIdx,
|
|
772
|
+
oldIdx
|
|
773
|
+
};
|
|
774
|
+
}
|
|
775
|
+
/**
|
|
776
|
+
* Hogwild epoch-sharded SGD across the pool. SAB mode writes concurrently into
|
|
777
|
+
* the shared embedding (races tolerated, like umap's parallel=True); transfer
|
|
778
|
+
* mode applies per-worker deltas at each epoch barrier (stale-by-one-epoch reads).
|
|
779
|
+
*/
|
|
780
|
+
async function parallelOptimizeLayout(ctx, embedding, dim, head, tail, nEpochs, nVertices, epochsPerSample, a, b, rngState, opts) {
|
|
781
|
+
const nEdges = epochsPerSample.length;
|
|
782
|
+
const shards = splitRange(nEdges, ctx.pool.size);
|
|
783
|
+
let shared = null;
|
|
784
|
+
if (ctx.sab) {
|
|
785
|
+
shared = new Float32Array(new SharedArrayBuffer(embedding.length * 4));
|
|
786
|
+
shared.set(embedding);
|
|
787
|
+
}
|
|
788
|
+
const embForWorkers = shared ?? embedding;
|
|
789
|
+
await ctx.pool.mapAll("layoutInit", (w) => ({
|
|
790
|
+
embedding: embForWorkers,
|
|
791
|
+
tailEmbedding: null,
|
|
792
|
+
head,
|
|
793
|
+
tail,
|
|
794
|
+
epochsPerSample,
|
|
795
|
+
rngState,
|
|
796
|
+
params: {
|
|
797
|
+
a,
|
|
798
|
+
b,
|
|
799
|
+
gamma: opts.gamma,
|
|
800
|
+
initialAlpha: opts.initialAlpha,
|
|
801
|
+
negativeSampleRate: opts.negativeSampleRate,
|
|
802
|
+
dim,
|
|
803
|
+
nVertices,
|
|
804
|
+
moveOther: opts.moveOther,
|
|
805
|
+
edgeStart: shards[w][0],
|
|
806
|
+
edgeEnd: shards[w][1],
|
|
807
|
+
sharedEmbedding: ctx.sab
|
|
808
|
+
}
|
|
809
|
+
}));
|
|
810
|
+
let alpha = opts.initialAlpha;
|
|
811
|
+
for (let epoch = 0; epoch < nEpochs; epoch++) {
|
|
812
|
+
if (opts.signal?.aborted) throw opts.signal.reason ?? new DOMException("The operation was aborted.", "AbortError");
|
|
813
|
+
if (ctx.sab) await ctx.pool.mapAll("layoutEpoch", () => ({
|
|
814
|
+
epoch,
|
|
815
|
+
alpha
|
|
816
|
+
}));
|
|
817
|
+
else {
|
|
818
|
+
const deltas = await ctx.pool.mapAll("layoutEpoch", () => ({
|
|
819
|
+
epoch,
|
|
820
|
+
alpha,
|
|
821
|
+
embedding
|
|
822
|
+
}));
|
|
823
|
+
for (const d of deltas) for (let i = 0; i < embedding.length; i++) embedding[i] = embedding[i] + d[i];
|
|
824
|
+
}
|
|
825
|
+
alpha = opts.initialAlpha * (1 - epoch / nEpochs);
|
|
826
|
+
if (opts.onEpoch) {
|
|
827
|
+
const info = {
|
|
828
|
+
epoch,
|
|
829
|
+
nEpochs,
|
|
830
|
+
phase: "layout"
|
|
831
|
+
};
|
|
832
|
+
if (opts.snapshotEvery && (epoch % opts.snapshotEvery === 0 || epoch === nEpochs - 1)) {
|
|
833
|
+
if (shared) embedding.set(shared);
|
|
834
|
+
info.embedding = embedding;
|
|
835
|
+
}
|
|
836
|
+
opts.onEpoch(info);
|
|
837
|
+
}
|
|
838
|
+
}
|
|
839
|
+
if (shared) embedding.set(shared);
|
|
840
|
+
return embedding;
|
|
841
|
+
}
|
|
842
|
+
/** Build the prepared search index from a parallel kNN result (main thread). */
|
|
843
|
+
function prepareIndexFromParallel(X, knnInternal, metricName, metricOptions, nNeighbors, searchRngState, angular) {
|
|
844
|
+
const internal = resolveInternalMetric(metricName, metricOptions);
|
|
845
|
+
const pairDist = (x, y) => internal.dist(x, y);
|
|
846
|
+
const searchIndex = prepareSearchIndex(X, knnInternal, pairDist, internal.correction, angular ?? internal.angularTrees, nNeighbors, searchRngState);
|
|
847
|
+
const outDst = new Float32Array(knnInternal.distances.length);
|
|
848
|
+
for (let t = 0; t < outDst.length; t++) {
|
|
849
|
+
const d = knnInternal.distances[t];
|
|
850
|
+
outDst[t] = d === Infinity ? Infinity : internal.correction(d);
|
|
851
|
+
}
|
|
852
|
+
return {
|
|
853
|
+
searchIndex,
|
|
854
|
+
knn: {
|
|
855
|
+
indices: knnInternal.indices.slice(),
|
|
856
|
+
distances: outDst,
|
|
857
|
+
k: knnInternal.k
|
|
858
|
+
}
|
|
859
|
+
};
|
|
860
|
+
}
|
|
861
|
+
|
|
862
|
+
//#endregion
|
|
863
|
+
//#region src/graph/supervised.ts
|
|
864
|
+
/**
|
|
865
|
+
* fast_intersection semantics in-place over a COPY of the graph:
|
|
866
|
+
* unknown (−1) edges ×= e^{−unknown_dist}; mismatched labels ×= e^{−far_dist}.
|
|
867
|
+
*/
|
|
868
|
+
function fastIntersection(graph, target, unknownDist = 1, farDist = 5) {
|
|
869
|
+
const out = {
|
|
870
|
+
data: graph.data.slice(),
|
|
871
|
+
indices: graph.indices,
|
|
872
|
+
indptr: graph.indptr,
|
|
873
|
+
rows: graph.rows,
|
|
874
|
+
cols: graph.cols
|
|
875
|
+
};
|
|
876
|
+
const eUnknown = Math.exp(-unknownDist);
|
|
877
|
+
const eFar = Math.exp(-farDist);
|
|
878
|
+
for (let i = 0; i < out.rows; i++) {
|
|
879
|
+
const ti = target[i];
|
|
880
|
+
for (let p = out.indptr[i]; p < out.indptr[i + 1]; p++) {
|
|
881
|
+
const tj = target[out.indices[p]];
|
|
882
|
+
if (ti === -1 || tj === -1) out.data[p] = out.data[p] * eUnknown;
|
|
883
|
+
else if (ti !== tj) out.data[p] = out.data[p] * eFar;
|
|
884
|
+
}
|
|
885
|
+
}
|
|
886
|
+
return out;
|
|
887
|
+
}
|
|
888
|
+
/** reset_local_connectivity (umap_.py:749-777): row-max normalize + fuzzy union. */
|
|
889
|
+
function resetLocalConnectivity(graph) {
|
|
890
|
+
const normed = csrRowNormalizeMax({
|
|
891
|
+
data: graph.data.slice(),
|
|
892
|
+
indices: graph.indices,
|
|
893
|
+
indptr: graph.indptr,
|
|
894
|
+
rows: graph.rows,
|
|
895
|
+
cols: graph.cols
|
|
896
|
+
});
|
|
897
|
+
return csrEliminateZeros(csrBinaryOp(normed, csrTranspose(normed), (a, b) => a + b - a * b));
|
|
898
|
+
}
|
|
899
|
+
/**
|
|
900
|
+
* discrete_metric_simplicial_set_intersection with metric=None (categorical path):
|
|
901
|
+
* fast_intersection then reset_local_connectivity (umap_.py:780-855).
|
|
902
|
+
*/
|
|
903
|
+
function categoricalSimplicialSetIntersection(graph, target, { unknownDist = 1, farDist = 5 } = {}) {
|
|
904
|
+
return resetLocalConnectivity(csrEliminateZeros(fastIntersection(graph, target, unknownDist, farDist)));
|
|
905
|
+
}
|
|
906
|
+
/**
|
|
907
|
+
* general_simplicial_set_intersection (umap_.py:858-883 + sparse.py:146-197):
|
|
908
|
+
* result sparsity = union (A + B); values re-mixed by weight where either side
|
|
909
|
+
* exceeds its floor.
|
|
910
|
+
*/
|
|
911
|
+
function generalSimplicialSetIntersection(A, B, weight) {
|
|
912
|
+
let leftMin = Infinity;
|
|
913
|
+
for (let t = 0; t < A.data.length; t++) if (A.data[t] < leftMin) leftMin = A.data[t];
|
|
914
|
+
leftMin = Math.max(leftMin / 2, 1e-8);
|
|
915
|
+
let rightMin = Infinity;
|
|
916
|
+
for (let t = 0; t < B.data.length; t++) if (B.data[t] < rightMin) rightMin = B.data[t];
|
|
917
|
+
rightMin = Math.min(Math.max(rightMin / 2, 1e-8), 1e-4);
|
|
918
|
+
return csrBinaryOp(A, B, (a, b) => {
|
|
919
|
+
const leftVal = a !== 0 ? a : leftMin;
|
|
920
|
+
const rightVal = b !== 0 ? b : rightMin;
|
|
921
|
+
if (leftVal > leftMin || rightVal > rightMin) {
|
|
922
|
+
if (weight < .5) return leftVal * rightVal ** (weight / (1 - weight));
|
|
923
|
+
return leftVal ** ((1 - weight) / weight) * rightVal;
|
|
924
|
+
}
|
|
925
|
+
return a + b;
|
|
926
|
+
});
|
|
927
|
+
}
|
|
928
|
+
/** far_dist rule from target_weight (umap_.py:2706-2713). */
|
|
929
|
+
function farDistFromWeight(targetWeight) {
|
|
930
|
+
return targetWeight < 1 ? 2.5 * (1 / (1 - targetWeight)) : 0xe8d4a51000;
|
|
931
|
+
}
|
|
932
|
+
|
|
933
|
+
//#endregion
|
|
934
|
+
//#region src/layout/densmap.ts
|
|
935
|
+
/** Binary-search lookup in a canonical CSR row; absent → 0 (DOK semantics). */
|
|
936
|
+
function csrLookup(m, row, col) {
|
|
937
|
+
let lo = m.indptr[row];
|
|
938
|
+
let hi = m.indptr[row + 1] - 1;
|
|
939
|
+
while (lo <= hi) {
|
|
940
|
+
const mid = lo + hi >> 1;
|
|
941
|
+
const c = m.indices[mid];
|
|
942
|
+
if (c === col) return m.data[mid];
|
|
943
|
+
if (c < col) lo = mid + 1;
|
|
944
|
+
else hi = mid - 1;
|
|
945
|
+
}
|
|
946
|
+
return 0;
|
|
947
|
+
}
|
|
948
|
+
/**
|
|
949
|
+
* Original radii over the PRUNED graph edges, using the symmetrized kNN distance
|
|
950
|
+
* matrix `dists`. `head`/`tail`/`mu` are the pruned COO arrays.
|
|
951
|
+
*/
|
|
952
|
+
function computeOrigRadii(head, tail, mu, dists, nVertices) {
|
|
953
|
+
const ro = new Float32Array(nVertices);
|
|
954
|
+
const muSum = new Float32Array(nVertices);
|
|
955
|
+
for (let i = 0; i < head.length; i++) {
|
|
956
|
+
const j = head[i];
|
|
957
|
+
const k = tail[i];
|
|
958
|
+
const dRaw = csrLookup(dists, j, k);
|
|
959
|
+
const D = dRaw * dRaw;
|
|
960
|
+
const m = mu[i];
|
|
961
|
+
ro[j] = ro[j] + m * D;
|
|
962
|
+
ro[k] = ro[k] + m * D;
|
|
963
|
+
muSum[j] = muSum[j] + m;
|
|
964
|
+
muSum[k] = muSum[k] + m;
|
|
965
|
+
}
|
|
966
|
+
const R = new Float32Array(nVertices);
|
|
967
|
+
let mean = 0;
|
|
968
|
+
for (let i = 0; i < nVertices; i++) {
|
|
969
|
+
ro[i] = Math.log(1e-8 + ro[i] / muSum[i]);
|
|
970
|
+
mean += ro[i];
|
|
971
|
+
}
|
|
972
|
+
mean /= nVertices;
|
|
973
|
+
let variance = 0;
|
|
974
|
+
for (let i = 0; i < nVertices; i++) variance += (ro[i] - mean) ** 2;
|
|
975
|
+
variance /= nVertices;
|
|
976
|
+
const std = Math.sqrt(variance);
|
|
977
|
+
for (let i = 0; i < nVertices; i++) R[i] = (ro[i] - mean) / std;
|
|
978
|
+
return {
|
|
979
|
+
ro,
|
|
980
|
+
muSum,
|
|
981
|
+
R
|
|
982
|
+
};
|
|
983
|
+
}
|
|
984
|
+
/**
|
|
985
|
+
* Embedding radii for outputDens: mu-weighted mean of UNsquared embedding kNN
|
|
986
|
+
* distances over a fresh fuzzy set built on the embedding (asymmetric with
|
|
987
|
+
* rad_orig, which squares — pinned quirk).
|
|
988
|
+
*/
|
|
989
|
+
function computeEmbRadii(graph, dists, nVertices) {
|
|
990
|
+
const re = new Float32Array(nVertices);
|
|
991
|
+
const muSum = new Float32Array(nVertices);
|
|
992
|
+
for (let j = 0; j < nVertices; j++) for (let p = graph.indptr[j]; p < graph.indptr[j + 1]; p++) {
|
|
993
|
+
const k = graph.indices[p];
|
|
994
|
+
const m = graph.data[p];
|
|
995
|
+
const D = csrLookup(dists, j, k);
|
|
996
|
+
re[j] = re[j] + m * D;
|
|
997
|
+
re[k] = re[k] + m * D;
|
|
998
|
+
muSum[j] = muSum[j] + m;
|
|
999
|
+
muSum[k] = muSum[k] + m;
|
|
1000
|
+
}
|
|
1001
|
+
for (let i = 0; i < nVertices; i++) re[i] = Math.log(1e-8 + re[i] / muSum[i]);
|
|
1002
|
+
return re;
|
|
1003
|
+
}
|
|
1004
|
+
|
|
1005
|
+
//#endregion
|
|
1006
|
+
//#region src/layout/outputMetrics.ts
|
|
1007
|
+
function euclideanGrad(x, y, grad) {
|
|
1008
|
+
let s = 0;
|
|
1009
|
+
for (let i = 0; i < x.length; i++) {
|
|
1010
|
+
const d = x[i] - y[i];
|
|
1011
|
+
s += d * d;
|
|
1012
|
+
}
|
|
1013
|
+
const dist = Math.sqrt(s);
|
|
1014
|
+
for (let i = 0; i < x.length; i++) grad[i] = (x[i] - y[i]) / (1e-6 + dist);
|
|
1015
|
+
return dist;
|
|
1016
|
+
}
|
|
1017
|
+
/** haversine_grad (distances.py:504-531) — 2-D only, latitudes offset by π/2. */
|
|
1018
|
+
function haversineGrad(x, y, grad) {
|
|
1019
|
+
if (x.length !== 2) throw new UmapValidationError("haversine is only defined for 2 dimensional data");
|
|
1020
|
+
const sinLat = Math.sin(.5 * (x[0] - y[0]));
|
|
1021
|
+
const cosLat = Math.cos(.5 * (x[0] - y[0]));
|
|
1022
|
+
const sinLong = Math.sin(.5 * (x[1] - y[1]));
|
|
1023
|
+
const cosLong = Math.cos(.5 * (x[1] - y[1]));
|
|
1024
|
+
const a1 = Math.cos(x[0] + Math.PI / 2) * Math.cos(y[0] + Math.PI / 2) * sinLong * sinLong + sinLat * sinLat;
|
|
1025
|
+
const d = 2 * Math.asin(Math.sqrt(Math.min(Math.max(Math.abs(a1), 0), 1)));
|
|
1026
|
+
const denom = Math.sqrt(Math.abs(a1 - 1)) * Math.sqrt(Math.abs(a1));
|
|
1027
|
+
grad[0] = (sinLat * cosLat - Math.sin(x[0] + Math.PI / 2) * Math.cos(y[0] + Math.PI / 2) * sinLong * sinLong) / (denom + 1e-6);
|
|
1028
|
+
grad[1] = Math.cos(x[0] + Math.PI / 2) * Math.cos(y[0] + Math.PI / 2) * sinLong * cosLong / (denom + 1e-6);
|
|
1029
|
+
return d;
|
|
1030
|
+
}
|
|
1031
|
+
/** hyperboloid_grad (distances.py:205-225) — reduced hyperboloid coordinates. */
|
|
1032
|
+
function hyperboloidGrad(x, y, grad) {
|
|
1033
|
+
let sx = 1;
|
|
1034
|
+
let sy = 1;
|
|
1035
|
+
let dot = 0;
|
|
1036
|
+
for (let i = 0; i < x.length; i++) {
|
|
1037
|
+
sx += x[i] * x[i];
|
|
1038
|
+
sy += y[i] * y[i];
|
|
1039
|
+
dot += x[i] * y[i];
|
|
1040
|
+
}
|
|
1041
|
+
const s = Math.sqrt(sx);
|
|
1042
|
+
const t = Math.sqrt(sy);
|
|
1043
|
+
let B = s * t - dot;
|
|
1044
|
+
if (B <= 1) B = 1.00000001;
|
|
1045
|
+
const gradCoeff = 1 / (Math.sqrt(B - 1) * Math.sqrt(B + 1));
|
|
1046
|
+
for (let i = 0; i < x.length; i++) grad[i] = gradCoeff * (x[i] * t / s - y[i]);
|
|
1047
|
+
return Math.acosh(B);
|
|
1048
|
+
}
|
|
1049
|
+
const OUTPUT_METRIC_GRADS = {
|
|
1050
|
+
euclidean: euclideanGrad,
|
|
1051
|
+
l2: euclideanGrad,
|
|
1052
|
+
haversine: haversineGrad,
|
|
1053
|
+
hyperboloid: hyperboloidGrad
|
|
1054
|
+
};
|
|
1055
|
+
function resolveOutputMetricGrad(name) {
|
|
1056
|
+
const fn = OUTPUT_METRIC_GRADS[name];
|
|
1057
|
+
if (!fn) throw new UmapValidationError(`gradient function is not implemented for output metric '${name}' (supported: ${Object.keys(OUTPUT_METRIC_GRADS).join(", ")})`);
|
|
1058
|
+
return fn;
|
|
1059
|
+
}
|
|
1060
|
+
function isEuclideanOutput(name) {
|
|
1061
|
+
return name === "euclidean" || name === "l2";
|
|
1062
|
+
}
|
|
1063
|
+
|
|
1064
|
+
//#endregion
|
|
1065
|
+
//#region src/estimator/options.ts
|
|
1066
|
+
/** DISCONNECTION_DISTANCES (umap_.py:61-68). */
|
|
1067
|
+
const DISCONNECTION_DISTANCES = {
|
|
1068
|
+
correlation: 2,
|
|
1069
|
+
cosine: 2,
|
|
1070
|
+
hellinger: 1,
|
|
1071
|
+
jaccard: 1,
|
|
1072
|
+
bit_jaccard: 1,
|
|
1073
|
+
dice: 1
|
|
1074
|
+
};
|
|
1075
|
+
|
|
1076
|
+
//#endregion
|
|
1077
|
+
//#region src/estimator/serialize.ts
|
|
1078
|
+
const MAGIC = "UMAPWEB1";
|
|
1079
|
+
const BYTES_PER = {
|
|
1080
|
+
f32: 4,
|
|
1081
|
+
f64: 8,
|
|
1082
|
+
i32: 4,
|
|
1083
|
+
u8: 1
|
|
1084
|
+
};
|
|
1085
|
+
function dtypeOf(arr) {
|
|
1086
|
+
if (arr instanceof Float32Array) return "f32";
|
|
1087
|
+
if (arr instanceof Float64Array) return "f64";
|
|
1088
|
+
if (arr instanceof Int32Array) return "i32";
|
|
1089
|
+
return "u8";
|
|
1090
|
+
}
|
|
1091
|
+
function serializeModel(model) {
|
|
1092
|
+
const m = model;
|
|
1093
|
+
if (!m._embedding) throw new UmapError("cannot serialize an unfitted model");
|
|
1094
|
+
if (typeof model.options.metric === "function") throw new UmapValidationError("models with custom metric functions cannot be serialized");
|
|
1095
|
+
if (m._rawCsr) throw new UmapValidationError("sparse-fitted model serialization is not supported");
|
|
1096
|
+
if (!m._graph || !m._sigmas || !m._rhos) throw new UmapError("cannot serialize a partially fitted model");
|
|
1097
|
+
const sections = [];
|
|
1098
|
+
const push = (name, arr) => {
|
|
1099
|
+
if (arr) sections.push({
|
|
1100
|
+
name,
|
|
1101
|
+
arr
|
|
1102
|
+
});
|
|
1103
|
+
};
|
|
1104
|
+
push("embedding", m._embedding.data);
|
|
1105
|
+
push("embeddingUnique", m._embeddingUnique);
|
|
1106
|
+
push("graph.data", m._graph.data);
|
|
1107
|
+
push("graph.indices", m._graph.indices);
|
|
1108
|
+
push("graph.indptr", m._graph.indptr);
|
|
1109
|
+
push("sigmas", m._sigmas);
|
|
1110
|
+
push("rhos", m._rhos);
|
|
1111
|
+
push("knn.indices", m._knn?.indices);
|
|
1112
|
+
push("knn.distances", m._knn?.distances);
|
|
1113
|
+
push("rawData", m._rawData?.data);
|
|
1114
|
+
push("uniqueIndex", m._uniqueIndex);
|
|
1115
|
+
push("uniqueInverse", m._uniqueInverse);
|
|
1116
|
+
push("radOrig", m._radOrig);
|
|
1117
|
+
push("radEmb", m._radEmb);
|
|
1118
|
+
const si = m._searchIndex;
|
|
1119
|
+
if (si) {
|
|
1120
|
+
push("si.indptr", si.indptr);
|
|
1121
|
+
push("si.indices", si.indices);
|
|
1122
|
+
push("si.vertexOrder", si.vertexOrder);
|
|
1123
|
+
push("si.data", si.data.data);
|
|
1124
|
+
push("si.tree.hyperplanes", si.tree.hyperplanes);
|
|
1125
|
+
push("si.tree.offsets", si.tree.offsets);
|
|
1126
|
+
push("si.tree.children", si.tree.children);
|
|
1127
|
+
push("si.tree.indices", si.tree.indices);
|
|
1128
|
+
push("si.searchRngState", si.searchRngState);
|
|
1129
|
+
}
|
|
1130
|
+
let offset = 0;
|
|
1131
|
+
const sectionMeta = [];
|
|
1132
|
+
for (const s of sections) {
|
|
1133
|
+
offset = Math.ceil(offset / 8) * 8;
|
|
1134
|
+
sectionMeta.push({
|
|
1135
|
+
name: s.name,
|
|
1136
|
+
dtype: dtypeOf(s.arr),
|
|
1137
|
+
offset,
|
|
1138
|
+
length: s.arr.length
|
|
1139
|
+
});
|
|
1140
|
+
offset += s.arr.length * BYTES_PER[dtypeOf(s.arr)];
|
|
1141
|
+
}
|
|
1142
|
+
const { onEpoch: _onEpoch, verbose: _verbose, device: _device, precomputedKnn: _precomputedKnn,...serializableOptions } = model.options;
|
|
1143
|
+
const header = {
|
|
1144
|
+
magic: MAGIC,
|
|
1145
|
+
version: 1,
|
|
1146
|
+
options: serializableOptions,
|
|
1147
|
+
scalars: {
|
|
1148
|
+
a: m._a,
|
|
1149
|
+
b: m._b,
|
|
1150
|
+
nNeighborsUsed: m._nNeighborsUsed,
|
|
1151
|
+
smallData: m._smallData,
|
|
1152
|
+
disconnection: m._disconnection === Infinity ? "inf" : m._disconnection,
|
|
1153
|
+
supervised: m._supervised,
|
|
1154
|
+
inputHash: m._inputHash,
|
|
1155
|
+
embedding: {
|
|
1156
|
+
rows: m._embedding.rows,
|
|
1157
|
+
cols: m._embedding.cols
|
|
1158
|
+
},
|
|
1159
|
+
graph: {
|
|
1160
|
+
rows: m._graph.rows,
|
|
1161
|
+
cols: m._graph.cols
|
|
1162
|
+
},
|
|
1163
|
+
knnK: m._knn?.k ?? null,
|
|
1164
|
+
rawData: m._rawData ? {
|
|
1165
|
+
rows: m._rawData.rows,
|
|
1166
|
+
cols: m._rawData.cols
|
|
1167
|
+
} : null,
|
|
1168
|
+
searchIndex: si ? {
|
|
1169
|
+
angular: si.angular,
|
|
1170
|
+
minDistance: si.minDistance,
|
|
1171
|
+
nNeighbors: si.nNeighbors,
|
|
1172
|
+
metricName: model.options.metric ?? "euclidean",
|
|
1173
|
+
tree: {
|
|
1174
|
+
leafSize: si.tree.leafSize,
|
|
1175
|
+
nNodes: si.tree.nNodes,
|
|
1176
|
+
dim: si.tree.dim
|
|
1177
|
+
},
|
|
1178
|
+
dataRows: si.data.rows,
|
|
1179
|
+
dataCols: si.data.cols
|
|
1180
|
+
} : null
|
|
1181
|
+
},
|
|
1182
|
+
sections: sectionMeta
|
|
1183
|
+
};
|
|
1184
|
+
const headerBytes = new TextEncoder().encode(JSON.stringify(header));
|
|
1185
|
+
const base = Math.ceil((12 + headerBytes.length) / 8) * 8;
|
|
1186
|
+
const total = base + offset;
|
|
1187
|
+
const out = new Uint8Array(total);
|
|
1188
|
+
const dv = new DataView(out.buffer);
|
|
1189
|
+
out.set(new TextEncoder().encode(MAGIC), 0);
|
|
1190
|
+
dv.setUint32(8, headerBytes.length, true);
|
|
1191
|
+
out.set(headerBytes, 12);
|
|
1192
|
+
for (let i = 0; i < sections.length; i++) {
|
|
1193
|
+
const s = sections[i];
|
|
1194
|
+
const metaS = sectionMeta[i];
|
|
1195
|
+
const bytes = new Uint8Array(s.arr.buffer, s.arr.byteOffset, s.arr.byteLength);
|
|
1196
|
+
out.set(bytes, base + metaS.offset);
|
|
1197
|
+
}
|
|
1198
|
+
return out;
|
|
1199
|
+
}
|
|
1200
|
+
function fail(why) {
|
|
1201
|
+
throw new UmapValidationError(`invalid serialized model: ${why}`);
|
|
1202
|
+
}
|
|
1203
|
+
function reqInt(v, name, min) {
|
|
1204
|
+
if (typeof v !== "number" || !Number.isInteger(v) || v < min) fail(`${name} must be an integer >= ${min}`);
|
|
1205
|
+
return v;
|
|
1206
|
+
}
|
|
1207
|
+
function reqFinite(v, name) {
|
|
1208
|
+
if (typeof v !== "number" || !Number.isFinite(v)) fail(`${name} must be a finite number`);
|
|
1209
|
+
return v;
|
|
1210
|
+
}
|
|
1211
|
+
function reqBool(v, name) {
|
|
1212
|
+
if (typeof v !== "boolean") fail(`${name} must be a boolean`);
|
|
1213
|
+
return v;
|
|
1214
|
+
}
|
|
1215
|
+
/** indptr shape: length n+1, starts at 0, monotone, ends at nnz. */
|
|
1216
|
+
function checkIndptr(indptr, n, nnz, name) {
|
|
1217
|
+
if (indptr.length !== n + 1) fail(`${name} length must be ${n + 1}`);
|
|
1218
|
+
if (indptr[0] !== 0) fail(`${name} must start at 0`);
|
|
1219
|
+
for (let i = 0; i < n; i++) if (indptr[i + 1] < indptr[i]) fail(`${name} must be non-decreasing`);
|
|
1220
|
+
if (indptr[n] !== nnz) fail(`${name} end (${indptr[n]}) must equal nnz (${nnz})`);
|
|
1221
|
+
}
|
|
1222
|
+
function checkIndexRange(arr, lo, hi, name) {
|
|
1223
|
+
for (let i = 0; i < arr.length; i++) {
|
|
1224
|
+
const v = arr[i];
|
|
1225
|
+
if (v < lo || v >= hi) fail(`${name} contains out-of-range index ${v}`);
|
|
1226
|
+
}
|
|
1227
|
+
}
|
|
1228
|
+
function deserializeModel(blob, UMAPClass) {
|
|
1229
|
+
if (!(blob instanceof Uint8Array)) fail("expected a Uint8Array");
|
|
1230
|
+
if (blob.byteLength < 12) fail("blob too short for a header");
|
|
1231
|
+
if (new TextDecoder().decode(blob.subarray(0, 8)) !== MAGIC) fail("bad magic (not a umap-web serialized model)");
|
|
1232
|
+
const headerLen = new DataView(blob.buffer, blob.byteOffset, blob.byteLength).getUint32(8, true);
|
|
1233
|
+
if (12 + headerLen > blob.byteLength) fail("declared header length exceeds the blob");
|
|
1234
|
+
let header;
|
|
1235
|
+
try {
|
|
1236
|
+
header = JSON.parse(new TextDecoder().decode(blob.subarray(12, 12 + headerLen)));
|
|
1237
|
+
} catch {
|
|
1238
|
+
fail("corrupt JSON header");
|
|
1239
|
+
}
|
|
1240
|
+
if (header?.version !== 1) throw new UmapValidationError(`unsupported serialization version ${header?.version}`);
|
|
1241
|
+
const sc = header.scalars;
|
|
1242
|
+
if (typeof sc !== "object" || sc === null || !Array.isArray(header.sections)) fail("malformed header");
|
|
1243
|
+
const base = Math.ceil((12 + headerLen) / 8) * 8;
|
|
1244
|
+
const metas = /* @__PURE__ */ new Map();
|
|
1245
|
+
for (const s of header.sections) {
|
|
1246
|
+
if (typeof s?.name !== "string" || !(s?.dtype in BYTES_PER)) fail("malformed section entry");
|
|
1247
|
+
reqInt(s.offset, `section ${s.name} offset`, 0);
|
|
1248
|
+
reqInt(s.length, `section ${s.name} length`, 0);
|
|
1249
|
+
if (base + s.offset + s.length * BYTES_PER[s.dtype] > blob.byteLength) fail(`section ${s.name} exceeds the blob bounds`);
|
|
1250
|
+
if (metas.has(s.name)) fail(`duplicate section ${s.name}`);
|
|
1251
|
+
metas.set(s.name, s);
|
|
1252
|
+
}
|
|
1253
|
+
const view = (s) => {
|
|
1254
|
+
const start = blob.byteOffset + base + s.offset;
|
|
1255
|
+
switch (s.dtype) {
|
|
1256
|
+
case "f32": return new Float32Array(blob.buffer.slice(start, start + s.length * 4));
|
|
1257
|
+
case "f64": return new Float64Array(blob.buffer.slice(start, start + s.length * 8));
|
|
1258
|
+
case "i32": return new Int32Array(blob.buffer.slice(start, start + s.length * 4));
|
|
1259
|
+
default: return new Uint8Array(blob.buffer.slice(start, start + s.length));
|
|
1260
|
+
}
|
|
1261
|
+
};
|
|
1262
|
+
const arrays = /* @__PURE__ */ new Map();
|
|
1263
|
+
for (const s of metas.values()) arrays.set(s.name, view(s));
|
|
1264
|
+
/** Fetch a section, enforcing its declared dtype (a lying header must not type-pun). */
|
|
1265
|
+
const need = (name, dtype) => {
|
|
1266
|
+
const meta = metas.get(name);
|
|
1267
|
+
const a = arrays.get(name);
|
|
1268
|
+
if (!meta || !a) fail(`missing section ${name}`);
|
|
1269
|
+
if (meta.dtype !== dtype) fail(`section ${name} has dtype ${meta.dtype}, expected ${dtype}`);
|
|
1270
|
+
return a;
|
|
1271
|
+
};
|
|
1272
|
+
const optional = (name, dtype) => metas.has(name) ? need(name, dtype) : null;
|
|
1273
|
+
const { onEpoch: _onEpoch, verbose: _verbose, device: _device, precomputedKnn: _precomputedKnn,...options } = typeof header.options === "object" && header.options !== null ? header.options : {};
|
|
1274
|
+
const model = new UMAPClass(options);
|
|
1275
|
+
const m = model;
|
|
1276
|
+
m._a = reqFinite(sc.a, "scalars.a");
|
|
1277
|
+
m._b = reqFinite(sc.b, "scalars.b");
|
|
1278
|
+
if (sc.disconnection !== "inf") reqFinite(sc.disconnection, "scalars.disconnection");
|
|
1279
|
+
m._disconnection = sc.disconnection === "inf" ? Infinity : sc.disconnection;
|
|
1280
|
+
m._smallData = reqBool(sc.smallData, "scalars.smallData");
|
|
1281
|
+
m._supervised = reqBool(sc.supervised, "scalars.supervised");
|
|
1282
|
+
if (sc.inputHash !== null && typeof sc.inputHash !== "string") fail("scalars.inputHash must be a string or null");
|
|
1283
|
+
m._inputHash = sc.inputHash;
|
|
1284
|
+
const embRows = reqInt(sc.embedding?.rows, "embedding.rows", 1);
|
|
1285
|
+
const embCols = reqInt(sc.embedding?.cols, "embedding.cols", 1);
|
|
1286
|
+
const embData = need("embedding", "f32");
|
|
1287
|
+
if (embData.length !== embRows * embCols) fail("embedding data does not match rows*cols");
|
|
1288
|
+
m._embedding = {
|
|
1289
|
+
data: embData,
|
|
1290
|
+
rows: embRows,
|
|
1291
|
+
cols: embCols
|
|
1292
|
+
};
|
|
1293
|
+
const gRows = reqInt(sc.graph?.rows, "graph.rows", 1);
|
|
1294
|
+
const gCols = reqInt(sc.graph?.cols, "graph.cols", 1);
|
|
1295
|
+
if (gRows !== gCols) fail("graph must be square");
|
|
1296
|
+
const gData = need("graph.data", "f64");
|
|
1297
|
+
const gIndices = need("graph.indices", "i32");
|
|
1298
|
+
const gIndptr = need("graph.indptr", "i32");
|
|
1299
|
+
if (gIndices.length !== gData.length) fail("graph indices/data length mismatch");
|
|
1300
|
+
checkIndptr(gIndptr, gRows, gData.length, "graph.indptr");
|
|
1301
|
+
checkIndexRange(gIndices, 0, gCols, "graph.indices");
|
|
1302
|
+
m._graph = {
|
|
1303
|
+
data: gData,
|
|
1304
|
+
indices: gIndices,
|
|
1305
|
+
indptr: gIndptr,
|
|
1306
|
+
rows: gRows,
|
|
1307
|
+
cols: gCols
|
|
1308
|
+
};
|
|
1309
|
+
m._nNeighborsUsed = reqInt(sc.nNeighborsUsed, "scalars.nNeighborsUsed", 1);
|
|
1310
|
+
if (m._nNeighborsUsed > gRows) fail("nNeighborsUsed exceeds the number of training rows");
|
|
1311
|
+
m._sigmas = need("sigmas", "f32");
|
|
1312
|
+
m._rhos = need("rhos", "f32");
|
|
1313
|
+
if (m._sigmas.length !== gRows || m._rhos.length !== gRows) fail("sigmas/rhos length must equal graph.rows");
|
|
1314
|
+
const embUnique = optional("embeddingUnique", "f32");
|
|
1315
|
+
if (embUnique && embUnique.length !== gRows * embCols) fail("embeddingUnique length must equal graph.rows * embedding.cols");
|
|
1316
|
+
m._embeddingUnique = embUnique ?? m._embedding.data;
|
|
1317
|
+
if (sc.knnK != null) {
|
|
1318
|
+
const k = reqInt(sc.knnK, "scalars.knnK", 1);
|
|
1319
|
+
const ki = need("knn.indices", "i32");
|
|
1320
|
+
const kd = need("knn.distances", "f32");
|
|
1321
|
+
if (ki.length !== gRows * k || kd.length !== gRows * k) fail("knn arrays must have graph.rows * k entries");
|
|
1322
|
+
checkIndexRange(ki, -1, gRows, "knn.indices");
|
|
1323
|
+
m._knn = {
|
|
1324
|
+
indices: ki,
|
|
1325
|
+
distances: kd,
|
|
1326
|
+
k
|
|
1327
|
+
};
|
|
1328
|
+
}
|
|
1329
|
+
if (sc.rawData != null) {
|
|
1330
|
+
const rRows = reqInt(sc.rawData.rows, "rawData.rows", 1);
|
|
1331
|
+
const rCols = reqInt(sc.rawData.cols, "rawData.cols", 1);
|
|
1332
|
+
const rData = need("rawData", "f32");
|
|
1333
|
+
if (rData.length !== rRows * rCols) fail("rawData does not match rows*cols");
|
|
1334
|
+
if (rRows !== embRows) fail("rawData.rows must equal embedding.rows");
|
|
1335
|
+
m._rawData = {
|
|
1336
|
+
data: rData,
|
|
1337
|
+
rows: rRows,
|
|
1338
|
+
cols: rCols
|
|
1339
|
+
};
|
|
1340
|
+
}
|
|
1341
|
+
const uniqueIndex = optional("uniqueIndex", "i32");
|
|
1342
|
+
const uniqueInverse = optional("uniqueInverse", "i32");
|
|
1343
|
+
if (uniqueIndex === null !== (uniqueInverse === null)) fail("uniqueIndex/uniqueInverse must be present together");
|
|
1344
|
+
if (uniqueIndex && uniqueInverse) {
|
|
1345
|
+
if (uniqueIndex.length !== gRows) fail("uniqueIndex length must equal graph.rows");
|
|
1346
|
+
if (uniqueInverse.length !== embRows) fail("uniqueInverse length must equal embedding.rows");
|
|
1347
|
+
if (m._rawData) checkIndexRange(uniqueIndex, 0, m._rawData.rows, "uniqueIndex");
|
|
1348
|
+
checkIndexRange(uniqueInverse, 0, gRows, "uniqueInverse");
|
|
1349
|
+
}
|
|
1350
|
+
m._uniqueIndex = uniqueIndex;
|
|
1351
|
+
m._uniqueInverse = uniqueInverse;
|
|
1352
|
+
const radOrig = optional("radOrig", "f32");
|
|
1353
|
+
const radEmb = optional("radEmb", "f32");
|
|
1354
|
+
if (radOrig && radOrig.length !== embRows) fail("radOrig length must equal embedding.rows");
|
|
1355
|
+
if (radEmb && radEmb.length !== embRows) fail("radEmb length must equal embedding.rows");
|
|
1356
|
+
m._radOrig = radOrig;
|
|
1357
|
+
m._radEmb = radEmb;
|
|
1358
|
+
if (sc.searchIndex != null) {
|
|
1359
|
+
const s = sc.searchIndex;
|
|
1360
|
+
if (typeof s.metricName !== "string") fail("searchIndex.metricName must be a string");
|
|
1361
|
+
const angular = reqBool(s.angular, "searchIndex.angular");
|
|
1362
|
+
if (typeof s.minDistance !== "number" || Number.isNaN(s.minDistance)) fail("searchIndex.minDistance must be a number");
|
|
1363
|
+
const siNeighbors = reqInt(s.nNeighbors, "searchIndex.nNeighbors", 1);
|
|
1364
|
+
const dataRows = reqInt(s.dataRows, "searchIndex.dataRows", 1);
|
|
1365
|
+
const dataCols = reqInt(s.dataCols, "searchIndex.dataCols", 1);
|
|
1366
|
+
if (dataRows !== gRows) fail("searchIndex.dataRows must equal graph.rows");
|
|
1367
|
+
if (m._rawData && dataCols !== m._rawData.cols) fail("searchIndex.dataCols must equal rawData.cols");
|
|
1368
|
+
const siData = need("si.data", "f32");
|
|
1369
|
+
if (siData.length !== dataRows * dataCols) fail("si.data does not match rows*cols");
|
|
1370
|
+
const leafSize = reqInt(s.tree?.leafSize, "searchIndex.tree.leafSize", 1);
|
|
1371
|
+
const nNodes = reqInt(s.tree?.nNodes, "searchIndex.tree.nNodes", 1);
|
|
1372
|
+
const dim = reqInt(s.tree?.dim, "searchIndex.tree.dim", 1);
|
|
1373
|
+
if (dim !== dataCols) fail("searchIndex.tree.dim must equal dataCols");
|
|
1374
|
+
const hyperplanes = need("si.tree.hyperplanes", "f32");
|
|
1375
|
+
const offsets = need("si.tree.offsets", "f32");
|
|
1376
|
+
const children = need("si.tree.children", "i32");
|
|
1377
|
+
const treeIndices = need("si.tree.indices", "i32");
|
|
1378
|
+
if (hyperplanes.length !== nNodes * dim) fail("si.tree.hyperplanes size mismatch");
|
|
1379
|
+
if (offsets.length !== nNodes) fail("si.tree.offsets size mismatch");
|
|
1380
|
+
if (children.length !== nNodes * 2) fail("si.tree.children size mismatch");
|
|
1381
|
+
if (treeIndices.length !== dataRows) fail("si.tree.indices size mismatch");
|
|
1382
|
+
checkIndexRange(treeIndices, 0, dataRows, "si.tree.indices");
|
|
1383
|
+
for (let i = 0; i < nNodes; i++) {
|
|
1384
|
+
const c0 = children[i * 2];
|
|
1385
|
+
const c1 = children[i * 2 + 1];
|
|
1386
|
+
if (c0 > 0) {
|
|
1387
|
+
if (c0 <= i || c0 >= nNodes || c1 <= i || c1 >= nNodes) fail(`si.tree.children node ${i} violates preorder traversal`);
|
|
1388
|
+
} else {
|
|
1389
|
+
const ls = -c0;
|
|
1390
|
+
const le = -c1;
|
|
1391
|
+
if (ls < 0 || le < ls || le > treeIndices.length) fail(`si.tree.children node ${i} has an invalid leaf range`);
|
|
1392
|
+
}
|
|
1393
|
+
}
|
|
1394
|
+
const tree = {
|
|
1395
|
+
hyperplanes,
|
|
1396
|
+
offsets,
|
|
1397
|
+
children,
|
|
1398
|
+
indices: treeIndices,
|
|
1399
|
+
leafSize,
|
|
1400
|
+
nNodes,
|
|
1401
|
+
dim
|
|
1402
|
+
};
|
|
1403
|
+
const siIndptr = need("si.indptr", "i32");
|
|
1404
|
+
const siIndices = need("si.indices", "i32");
|
|
1405
|
+
checkIndptr(siIndptr, dataRows, siIndices.length, "si.indptr");
|
|
1406
|
+
checkIndexRange(siIndices, 0, dataRows, "si.indices");
|
|
1407
|
+
const vertexOrder = need("si.vertexOrder", "i32");
|
|
1408
|
+
if (vertexOrder.length !== dataRows) fail("si.vertexOrder size mismatch");
|
|
1409
|
+
checkIndexRange(vertexOrder, 0, dataRows, "si.vertexOrder");
|
|
1410
|
+
const searchRngState = need("si.searchRngState", "i32");
|
|
1411
|
+
if (searchRngState.length !== 3) fail("si.searchRngState must have 3 words");
|
|
1412
|
+
const internal = resolveInternalMetric(s.metricName, model.options.metricOptions ?? {});
|
|
1413
|
+
m._searchIndex = {
|
|
1414
|
+
indptr: siIndptr,
|
|
1415
|
+
indices: siIndices,
|
|
1416
|
+
tree,
|
|
1417
|
+
vertexOrder,
|
|
1418
|
+
data: {
|
|
1419
|
+
data: siData,
|
|
1420
|
+
rows: dataRows,
|
|
1421
|
+
cols: dataCols
|
|
1422
|
+
},
|
|
1423
|
+
dist: (a, b) => internal.dist(a, b),
|
|
1424
|
+
correction: internal.correction,
|
|
1425
|
+
angular,
|
|
1426
|
+
minDistance: s.minDistance,
|
|
1427
|
+
searchRngState,
|
|
1428
|
+
nNeighbors: siNeighbors
|
|
1429
|
+
};
|
|
1430
|
+
}
|
|
1431
|
+
return model;
|
|
1432
|
+
}
|
|
1433
|
+
|
|
1434
|
+
//#endregion
|
|
1435
|
+
//#region src/estimator/transform-init.ts
|
|
1436
|
+
/**
|
|
1437
|
+
* init_graph_transform (umap_.py:1337-1375): weighted mean of train embedding rows.
|
|
1438
|
+
* Rows with no nonzeros → NaN; a weight == 1.0 copies that train row exactly.
|
|
1439
|
+
*/
|
|
1440
|
+
function initGraphTransform(graph, embedding, dim) {
|
|
1441
|
+
const nNew = graph.rows;
|
|
1442
|
+
const out = new Float32Array(nNew * dim);
|
|
1443
|
+
for (let row = 0; row < nNew; row++) {
|
|
1444
|
+
const start = graph.indptr[row];
|
|
1445
|
+
const end = graph.indptr[row + 1];
|
|
1446
|
+
if (start === end) {
|
|
1447
|
+
for (let d = 0; d < dim; d++) out[row * dim + d] = NaN;
|
|
1448
|
+
continue;
|
|
1449
|
+
}
|
|
1450
|
+
let rowSum = 0;
|
|
1451
|
+
let exactCol = -1;
|
|
1452
|
+
for (let p = start; p < end; p++) {
|
|
1453
|
+
rowSum += graph.data[p];
|
|
1454
|
+
if (graph.data[p] === 1 && exactCol === -1) exactCol = graph.indices[p];
|
|
1455
|
+
}
|
|
1456
|
+
if (exactCol !== -1) {
|
|
1457
|
+
for (let d = 0; d < dim; d++) out[row * dim + d] = embedding[exactCol * dim + d];
|
|
1458
|
+
continue;
|
|
1459
|
+
}
|
|
1460
|
+
for (let p = start; p < end; p++) {
|
|
1461
|
+
const w = graph.data[p] / rowSum;
|
|
1462
|
+
const col = graph.indices[p];
|
|
1463
|
+
for (let d = 0; d < dim; d++) out[row * dim + d] = out[row * dim + d] + w * embedding[col * dim + d];
|
|
1464
|
+
}
|
|
1465
|
+
}
|
|
1466
|
+
return out;
|
|
1467
|
+
}
|
|
1468
|
+
/**
|
|
1469
|
+
* init_transform (umap_.py:1304-1334): constant-degree weighted mean of raw data
|
|
1470
|
+
* rows (used by inverseTransform).
|
|
1471
|
+
*/
|
|
1472
|
+
function initTransform(indices, weights, degree, data, dim) {
|
|
1473
|
+
const nNew = indices.length / degree;
|
|
1474
|
+
const out = new Float32Array(nNew * dim);
|
|
1475
|
+
for (let i = 0; i < nNew; i++) for (let j = 0; j < degree; j++) {
|
|
1476
|
+
const w = weights[i * degree + j];
|
|
1477
|
+
const col = indices[i * degree + j];
|
|
1478
|
+
for (let d = 0; d < dim; d++) out[i * dim + d] = out[i * dim + d] + w * data[col * dim + d];
|
|
1479
|
+
}
|
|
1480
|
+
return out;
|
|
1481
|
+
}
|
|
1482
|
+
/**
|
|
1483
|
+
* init_update (umap_.py:1378-1390) for update(): each new point gets the summed
|
|
1484
|
+
* coordinates of its OLD-point neighbors divided by (count × dim) — the pinned
|
|
1485
|
+
* quirk divides by the extra dim factor.
|
|
1486
|
+
*/
|
|
1487
|
+
function initUpdate(embedding, dim, originalSize, knnIndices, k) {
|
|
1488
|
+
const nTotal = embedding.length / dim;
|
|
1489
|
+
for (let i = originalSize; i < nTotal; i++) {
|
|
1490
|
+
let count = 0;
|
|
1491
|
+
for (let j = 0; j < k; j++) {
|
|
1492
|
+
const nb = knnIndices[i * k + j];
|
|
1493
|
+
if (nb >= 0 && nb < originalSize) {
|
|
1494
|
+
count += dim;
|
|
1495
|
+
for (let d = 0; d < dim; d++) embedding[i * dim + d] = embedding[i * dim + d] + embedding[nb * dim + d];
|
|
1496
|
+
}
|
|
1497
|
+
}
|
|
1498
|
+
if (count > 0) for (let d = 0; d < dim; d++) embedding[i * dim + d] = embedding[i * dim + d] / count;
|
|
1499
|
+
}
|
|
1500
|
+
}
|
|
1501
|
+
|
|
1502
|
+
//#endregion
|
|
1503
|
+
//#region src/estimator/umap.ts
|
|
1504
|
+
const yieldToLoop = typeof setImmediate === "function" ? () => new Promise((r) => setImmediate(r)) : () => new Promise((r) => setTimeout(r, 0));
|
|
1505
|
+
/** Bound on requestAdapter/requestDevice — a broken GPU stack must fall back, not hang. */
|
|
1506
|
+
const GPU_ACQUIRE_TIMEOUT_MS = 5e3;
|
|
1507
|
+
function entropySeed() {
|
|
1508
|
+
const cryptoObj = globalThis.crypto;
|
|
1509
|
+
if (typeof cryptoObj?.getRandomValues !== "function") throw new UmapError("no global crypto available for entropy seeding — pass options.seed explicitly");
|
|
1510
|
+
const buf = new Uint32Array(2);
|
|
1511
|
+
cryptoObj.getRandomValues(buf);
|
|
1512
|
+
return buf[0] + (buf[1] & 2097151) * 4294967296;
|
|
1513
|
+
}
|
|
1514
|
+
/** FNV-1a over bytes — cheap input-identity hash (same-data transform shortcut). */
|
|
1515
|
+
function contentHash(view) {
|
|
1516
|
+
let h1 = 2166136261;
|
|
1517
|
+
let h2 = 3421674724;
|
|
1518
|
+
for (let i = 0; i < view.length; i++) {
|
|
1519
|
+
h1 = Math.imul(h1 ^ view[i], 16777619) >>> 0;
|
|
1520
|
+
h2 = Math.imul(h2 ^ view[i], 16777623) >>> 0;
|
|
1521
|
+
}
|
|
1522
|
+
return `${view.length}:${h1.toString(16)}:${h2.toString(16)}`;
|
|
1523
|
+
}
|
|
1524
|
+
/** Bitwise equality (confirms a contentHash match; Float32Array offsets are 4-aligned). */
|
|
1525
|
+
function sameFloat32Bits(a, b) {
|
|
1526
|
+
if (a.length !== b.length) return false;
|
|
1527
|
+
const ai = new Int32Array(a.buffer, a.byteOffset, a.length);
|
|
1528
|
+
const bi = new Int32Array(b.buffer, b.byteOffset, b.length);
|
|
1529
|
+
for (let i = 0; i < ai.length; i++) if (ai[i] !== bi[i]) return false;
|
|
1530
|
+
return true;
|
|
1531
|
+
}
|
|
1532
|
+
var UMAP = class UMAP {
|
|
1533
|
+
options;
|
|
1534
|
+
logger;
|
|
1535
|
+
_embedding = null;
|
|
1536
|
+
_embeddingUnique = null;
|
|
1537
|
+
_graph = null;
|
|
1538
|
+
_sigmas = null;
|
|
1539
|
+
_rhos = null;
|
|
1540
|
+
_knn = null;
|
|
1541
|
+
_searchIndex = null;
|
|
1542
|
+
_rawData = null;
|
|
1543
|
+
_rawCsr = null;
|
|
1544
|
+
_uniqueIndex = null;
|
|
1545
|
+
_uniqueInverse = null;
|
|
1546
|
+
_a = 0;
|
|
1547
|
+
_b = 0;
|
|
1548
|
+
_nNeighborsUsed = 15;
|
|
1549
|
+
_smallData = false;
|
|
1550
|
+
_disconnection = Infinity;
|
|
1551
|
+
_supervised = false;
|
|
1552
|
+
_inputHash = null;
|
|
1553
|
+
_radOrig = null;
|
|
1554
|
+
_radEmb = null;
|
|
1555
|
+
_backendInfo = {
|
|
1556
|
+
knn: "none",
|
|
1557
|
+
layout: "none"
|
|
1558
|
+
};
|
|
1559
|
+
_timings = {};
|
|
1560
|
+
_parallel = null;
|
|
1561
|
+
_gpu = null;
|
|
1562
|
+
_gpuCtx = null;
|
|
1563
|
+
_useGpuLayout = false;
|
|
1564
|
+
/** name of the in-flight operation; guards against concurrent use */
|
|
1565
|
+
_busy = null;
|
|
1566
|
+
constructor(options = {}) {
|
|
1567
|
+
this.options = { ...options };
|
|
1568
|
+
this.logger = typeof options.verbose === "object" ? options.verbose : options.verbose ? consoleLogger : silentLogger;
|
|
1569
|
+
this.validateOptions();
|
|
1570
|
+
}
|
|
1571
|
+
validateOptions() {
|
|
1572
|
+
const o = this.options;
|
|
1573
|
+
const positive = (name, v) => {
|
|
1574
|
+
if (v !== void 0 && !(v > 0)) throw new UmapValidationError(`${name} must be > 0, got ${v}`);
|
|
1575
|
+
};
|
|
1576
|
+
const positiveInt = (name, v) => {
|
|
1577
|
+
if (v !== void 0 && (!Number.isInteger(v) || v <= 0)) throw new UmapValidationError(`${name} must be a positive integer, got ${v}`);
|
|
1578
|
+
};
|
|
1579
|
+
const inUnit = (name, v) => {
|
|
1580
|
+
if (v !== void 0 && !(v >= 0 && v <= 1)) throw new UmapValidationError(`${name} must be in [0, 1], got ${v}`);
|
|
1581
|
+
};
|
|
1582
|
+
positiveInt("nNeighbors", o.nNeighbors);
|
|
1583
|
+
if (o.nNeighbors !== void 0 && o.nNeighbors < 2) throw new UmapValidationError("nNeighbors must be >= 2 (matching umap-learn)");
|
|
1584
|
+
positiveInt("nComponents", o.nComponents);
|
|
1585
|
+
positiveInt("negativeSampleRate", o.negativeSampleRate);
|
|
1586
|
+
if (o.targetNNeighbors !== void 0 && o.targetNNeighbors !== -1) positiveInt("targetNNeighbors", o.targetNNeighbors);
|
|
1587
|
+
positive("learningRate", o.learningRate);
|
|
1588
|
+
positive("spread", o.spread);
|
|
1589
|
+
positive("localConnectivity", o.localConnectivity);
|
|
1590
|
+
if (o.repulsionStrength !== void 0 && !(o.repulsionStrength >= 0)) throw new UmapValidationError(`repulsionStrength must be >= 0, got ${o.repulsionStrength}`);
|
|
1591
|
+
if (o.minDist !== void 0 && o.minDist < 0) throw new UmapValidationError(`minDist must be >= 0, got ${o.minDist}`);
|
|
1592
|
+
if (o.minDist !== void 0 && o.spread !== void 0 && o.minDist > o.spread) throw new UmapValidationError("minDist must be <= spread");
|
|
1593
|
+
inUnit("setOpMixRatio", o.setOpMixRatio);
|
|
1594
|
+
inUnit("targetWeight", o.targetWeight);
|
|
1595
|
+
inUnit("densFrac", o.densFrac);
|
|
1596
|
+
inUnit("densVarShift", o.densVarShift);
|
|
1597
|
+
if (o.nEpochs !== void 0 && (o.nEpochs < 0 || !Number.isInteger(o.nEpochs))) throw new UmapValidationError("nEpochs must be a nonnegative integer");
|
|
1598
|
+
if (o.a === void 0 !== (o.b === void 0)) throw new UmapValidationError("a and b must be provided together (or both derived)");
|
|
1599
|
+
if (o.concurrency !== void 0 && o.concurrency !== "auto") positiveInt("concurrency", o.concurrency);
|
|
1600
|
+
if (o.backend !== void 0 && o.backend !== "auto" && o.backend !== "cpu" && o.backend !== "webgpu") throw new UmapValidationError(`backend must be 'auto' | 'cpu' | 'webgpu' (custom Backend instances are not supported yet), got ${JSON.stringify(o.backend)}`);
|
|
1601
|
+
if (typeof o.metric === "string" && o.metric !== "precomputed" && !isNamedMetric(o.metric)) throw new UmapValidationError(`unknown metric '${o.metric}'`);
|
|
1602
|
+
if (o.targetMetric !== void 0 && o.targetMetric !== "categorical" && !isNamedMetric(o.targetMetric)) throw new UmapValidationError(`unknown targetMetric '${o.targetMetric}'`);
|
|
1603
|
+
if (o.outputMetric === "precomputed") throw new UmapValidationError("outputMetric cannot be precomputed");
|
|
1604
|
+
if (o.outputMetric !== void 0 && !isEuclideanOutput(o.outputMetric)) resolveOutputMetricGrad(o.outputMetric);
|
|
1605
|
+
if (o.outputMetricOptions && Object.keys(o.outputMetricOptions).length > 0) throw new UmapValidationError("outputMetricOptions is not consumed by any named output metric in this port; remove it rather than have it silently ignored");
|
|
1606
|
+
if (o.densMap && o.outputMetric && !isEuclideanOutput(o.outputMetric)) throw new UmapValidationError("densMap requires euclidean output metric");
|
|
1607
|
+
if (o.unique && o.metric === "precomputed") throw new UmapValidationError("unique is poorly defined on a precomputed metric");
|
|
1608
|
+
if (o.unique && o.precomputedKnn) throw new UmapValidationError("unique cannot be used with precomputedKnn");
|
|
1609
|
+
}
|
|
1610
|
+
get embedding() {
|
|
1611
|
+
if (!this._embedding) throw new UmapError("model is not fitted");
|
|
1612
|
+
return this._embedding;
|
|
1613
|
+
}
|
|
1614
|
+
/** CSR fuzzy graph over the training data (unique rows when unique=true). */
|
|
1615
|
+
get graph() {
|
|
1616
|
+
if (!this._graph) throw new UmapError("model is not fitted");
|
|
1617
|
+
return this._graph;
|
|
1618
|
+
}
|
|
1619
|
+
get densities() {
|
|
1620
|
+
if (!this._radOrig || !this._radEmb) throw new UmapError("densities require fitting with outputDens: true");
|
|
1621
|
+
return {
|
|
1622
|
+
original: this._radOrig,
|
|
1623
|
+
embedded: this._radEmb
|
|
1624
|
+
};
|
|
1625
|
+
}
|
|
1626
|
+
/** Which backend actually executed each stage (§5.2). */
|
|
1627
|
+
get backendInfo() {
|
|
1628
|
+
return { ...this._backendInfo };
|
|
1629
|
+
}
|
|
1630
|
+
get phaseTimings() {
|
|
1631
|
+
return { ...this._timings };
|
|
1632
|
+
}
|
|
1633
|
+
resolveMetricName() {
|
|
1634
|
+
const m = this.options.metric ?? "euclidean";
|
|
1635
|
+
return typeof m === "function" ? "<custom>" : m;
|
|
1636
|
+
}
|
|
1637
|
+
/** Serialize operations on one instance: concurrent calls are a caller bug. */
|
|
1638
|
+
async withBusy(name, fn) {
|
|
1639
|
+
if (this._busy) throw new UmapError(`cannot ${name}: ${this._busy}() is already in progress on this estimator`);
|
|
1640
|
+
this._busy = name;
|
|
1641
|
+
try {
|
|
1642
|
+
return await fn();
|
|
1643
|
+
} finally {
|
|
1644
|
+
this._busy = null;
|
|
1645
|
+
}
|
|
1646
|
+
}
|
|
1647
|
+
snapshotFittedState() {
|
|
1648
|
+
return {
|
|
1649
|
+
_embedding: this._embedding,
|
|
1650
|
+
_embeddingUnique: this._embeddingUnique,
|
|
1651
|
+
_graph: this._graph,
|
|
1652
|
+
_sigmas: this._sigmas,
|
|
1653
|
+
_rhos: this._rhos,
|
|
1654
|
+
_knn: this._knn,
|
|
1655
|
+
_searchIndex: this._searchIndex,
|
|
1656
|
+
_rawData: this._rawData,
|
|
1657
|
+
_rawCsr: this._rawCsr,
|
|
1658
|
+
_uniqueIndex: this._uniqueIndex,
|
|
1659
|
+
_uniqueInverse: this._uniqueInverse,
|
|
1660
|
+
_radOrig: this._radOrig,
|
|
1661
|
+
_radEmb: this._radEmb,
|
|
1662
|
+
_inputHash: this._inputHash,
|
|
1663
|
+
_a: this._a,
|
|
1664
|
+
_b: this._b,
|
|
1665
|
+
_nNeighborsUsed: this._nNeighborsUsed,
|
|
1666
|
+
_smallData: this._smallData,
|
|
1667
|
+
_disconnection: this._disconnection,
|
|
1668
|
+
_supervised: this._supervised,
|
|
1669
|
+
_useGpuLayout: this._useGpuLayout,
|
|
1670
|
+
_backendInfo: { ...this._backendInfo },
|
|
1671
|
+
_timings: { ...this._timings }
|
|
1672
|
+
};
|
|
1673
|
+
}
|
|
1674
|
+
restoreFittedState(snapshot) {
|
|
1675
|
+
Object.assign(this, snapshot);
|
|
1676
|
+
}
|
|
1677
|
+
/** Discard fitted state so a running fit never exposes a mixed old/new model. */
|
|
1678
|
+
resetFittedState() {
|
|
1679
|
+
this._embedding = null;
|
|
1680
|
+
this._embeddingUnique = null;
|
|
1681
|
+
this._graph = null;
|
|
1682
|
+
this._sigmas = null;
|
|
1683
|
+
this._rhos = null;
|
|
1684
|
+
this._knn = null;
|
|
1685
|
+
this._searchIndex = null;
|
|
1686
|
+
this._rawData = null;
|
|
1687
|
+
this._rawCsr = null;
|
|
1688
|
+
this._uniqueIndex = null;
|
|
1689
|
+
this._uniqueInverse = null;
|
|
1690
|
+
this._radOrig = null;
|
|
1691
|
+
this._radEmb = null;
|
|
1692
|
+
this._inputHash = null;
|
|
1693
|
+
this._supervised = false;
|
|
1694
|
+
this._useGpuLayout = false;
|
|
1695
|
+
this._backendInfo = {
|
|
1696
|
+
knn: "none",
|
|
1697
|
+
layout: "none"
|
|
1698
|
+
};
|
|
1699
|
+
this._timings = {};
|
|
1700
|
+
}
|
|
1701
|
+
/** Stop using the GPU for the rest of this fit (typed failure → CPU, D-017). */
|
|
1702
|
+
dropGpu() {
|
|
1703
|
+
if (this._gpuCtx && !this._gpuCtx.injected) try {
|
|
1704
|
+
this._gpuCtx.device.destroy();
|
|
1705
|
+
} catch {}
|
|
1706
|
+
this._gpuCtx = null;
|
|
1707
|
+
this._gpu = null;
|
|
1708
|
+
this._useGpuLayout = false;
|
|
1709
|
+
}
|
|
1710
|
+
/** Always runs after fit (success, error, or abort): no leaked workers/devices. */
|
|
1711
|
+
async releaseBackends() {
|
|
1712
|
+
if (this._parallel) {
|
|
1713
|
+
const pool = this._parallel.pool;
|
|
1714
|
+
this._parallel = null;
|
|
1715
|
+
await pool.terminate();
|
|
1716
|
+
}
|
|
1717
|
+
if (this._gpuCtx && !this._gpuCtx.injected) {
|
|
1718
|
+
this._gpuCtx.device.destroy();
|
|
1719
|
+
this._gpuCtx = null;
|
|
1720
|
+
this._gpu = null;
|
|
1721
|
+
this._useGpuLayout = false;
|
|
1722
|
+
}
|
|
1723
|
+
}
|
|
1724
|
+
async fit(X, y, { signal } = {}) {
|
|
1725
|
+
return this.withBusy("fit", async () => {
|
|
1726
|
+
throwIfAborted(signal);
|
|
1727
|
+
const snapshot = this.snapshotFittedState();
|
|
1728
|
+
this.resetFittedState();
|
|
1729
|
+
try {
|
|
1730
|
+
return await this.fitInner(X, y, signal);
|
|
1731
|
+
} catch (e) {
|
|
1732
|
+
this.restoreFittedState(snapshot);
|
|
1733
|
+
throw e;
|
|
1734
|
+
} finally {
|
|
1735
|
+
await this.releaseBackends();
|
|
1736
|
+
}
|
|
1737
|
+
});
|
|
1738
|
+
}
|
|
1739
|
+
async fitInner(X, y, signal) {
|
|
1740
|
+
const o = this.options;
|
|
1741
|
+
const nComponents = o.nComponents ?? 2;
|
|
1742
|
+
const metric = o.metric ?? "euclidean";
|
|
1743
|
+
const minDist = o.minDist ?? .1;
|
|
1744
|
+
const spread = o.spread ?? 1;
|
|
1745
|
+
const densMap = o.densMap ?? false;
|
|
1746
|
+
const outputDens = o.outputDens ?? false;
|
|
1747
|
+
const outputMetric = o.outputMetric ?? "euclidean";
|
|
1748
|
+
const seed = o.seed ?? entropySeed();
|
|
1749
|
+
const rng = Pcg32.fromSeed(seed);
|
|
1750
|
+
let denseFull = null;
|
|
1751
|
+
let sparse = null;
|
|
1752
|
+
if (isCsrInput(X)) {
|
|
1753
|
+
sparse = validateCsr(X);
|
|
1754
|
+
this._rawCsr = sparse;
|
|
1755
|
+
} else {
|
|
1756
|
+
denseFull = toMatrix(X);
|
|
1757
|
+
this._rawData = denseFull;
|
|
1758
|
+
}
|
|
1759
|
+
const nFull = denseFull ? denseFull.rows : sparse.rows;
|
|
1760
|
+
if (nFull <= 1) throw new UmapValidationError("need at least 2 samples");
|
|
1761
|
+
if (typeof metric === "string" && o.metricOptions) validateMetricOptionSizes(metric, o.metricOptions, denseFull ? denseFull.cols : sparse.cols);
|
|
1762
|
+
this._inputHash = denseFull ? contentHash(new Uint8Array(denseFull.data.buffer, denseFull.data.byteOffset, denseFull.data.byteLength)) : contentHash(new Uint8Array(sparse.data.buffer, sparse.data.byteOffset, sparse.data.byteLength));
|
|
1763
|
+
let dense = denseFull;
|
|
1764
|
+
if (o.unique) {
|
|
1765
|
+
if (!denseFull) throw new UmapValidationError("unique=true requires dense input");
|
|
1766
|
+
const { index, inverse } = uniqueRows(denseFull);
|
|
1767
|
+
this._uniqueIndex = index;
|
|
1768
|
+
this._uniqueInverse = inverse;
|
|
1769
|
+
if (index.length < nFull) this.logger.info(`unique: ${nFull} rows → ${index.length} unique`);
|
|
1770
|
+
const d = denseFull.cols;
|
|
1771
|
+
const data = new Float32Array(index.length * d);
|
|
1772
|
+
for (let i = 0; i < index.length; i++) data.set(denseFull.data.subarray(index[i] * d, (index[i] + 1) * d), i * d);
|
|
1773
|
+
dense = {
|
|
1774
|
+
data,
|
|
1775
|
+
rows: index.length,
|
|
1776
|
+
cols: d
|
|
1777
|
+
};
|
|
1778
|
+
}
|
|
1779
|
+
const n = dense ? dense.rows : sparse.rows;
|
|
1780
|
+
let nNeighbors = o.nNeighbors ?? 15;
|
|
1781
|
+
if (n - 1 < nNeighbors) {
|
|
1782
|
+
this.logger.warn(`nNeighbors ${nNeighbors} is larger than dataset size ${n}; reduced to ${n - 1}`);
|
|
1783
|
+
nNeighbors = n - 1;
|
|
1784
|
+
}
|
|
1785
|
+
this._nNeighborsUsed = nNeighbors;
|
|
1786
|
+
if (o.a !== void 0 && o.b !== void 0) {
|
|
1787
|
+
this._a = o.a;
|
|
1788
|
+
this._b = o.b;
|
|
1789
|
+
} else {
|
|
1790
|
+
const ab = findABParams(spread, minDist);
|
|
1791
|
+
this._a = ab.a;
|
|
1792
|
+
this._b = ab.b;
|
|
1793
|
+
}
|
|
1794
|
+
const metricName = this.resolveMetricName();
|
|
1795
|
+
this._disconnection = o.disconnectionDistance ?? DISCONNECTION_DISTANCES[metricName] ?? Infinity;
|
|
1796
|
+
this._smallData = n < SMALL_DATA_CUTOFF && !(o.forceApproximation ?? false);
|
|
1797
|
+
const backendOpt = o.backend ?? "auto";
|
|
1798
|
+
let useGpuKnn = false;
|
|
1799
|
+
let useGpuLayout = false;
|
|
1800
|
+
if (backendOpt === "webgpu" && dense === null) this.logger.warn("backend \"webgpu\" requires dense input; this fit runs on the CPU");
|
|
1801
|
+
if (backendOpt === "webgpu" && (o.deterministic ?? false)) this.logger.warn("deterministic: true forces the reproducible CPU path; backend \"webgpu\" is ignored (§5.5)");
|
|
1802
|
+
if ((backendOpt === "webgpu" || backendOpt === "auto") && !(o.deterministic ?? false) && dense !== null) {
|
|
1803
|
+
const hasGpuApi = !!globalThis.navigator?.gpu || o.device !== void 0;
|
|
1804
|
+
if (backendOpt === "webgpu" && !hasGpuApi) throw new UmapBackendError("backend \"webgpu\" requested but WebGPU is unavailable");
|
|
1805
|
+
if (hasGpuApi) {
|
|
1806
|
+
const mod = await import("./webgpu.js");
|
|
1807
|
+
const support = mod.gpuSupportsConfig({
|
|
1808
|
+
...typeof metric === "string" ? { metric } : { metric },
|
|
1809
|
+
outputMetric,
|
|
1810
|
+
densMap,
|
|
1811
|
+
nNeighbors,
|
|
1812
|
+
nComponents
|
|
1813
|
+
});
|
|
1814
|
+
if (backendOpt === "webgpu" || support.knn || support.layout) {
|
|
1815
|
+
const acquire = mod.acquireGpu(o.device);
|
|
1816
|
+
acquire.catch(() => {});
|
|
1817
|
+
try {
|
|
1818
|
+
this._gpuCtx = await raceTimeout(acquire, GPU_ACQUIRE_TIMEOUT_MS, "WebGPU device acquisition");
|
|
1819
|
+
this._gpu = mod;
|
|
1820
|
+
useGpuKnn = support.knn;
|
|
1821
|
+
useGpuLayout = support.layout;
|
|
1822
|
+
this._useGpuLayout = support.layout;
|
|
1823
|
+
if (support.reason) this.logger.warn(support.reason);
|
|
1824
|
+
this.logger.info(`WebGPU: ${this._gpuCtx.adapterInfo.vendor}/${this._gpuCtx.adapterInfo.architecture}`);
|
|
1825
|
+
} catch (e) {
|
|
1826
|
+
acquire.then((ctx) => {
|
|
1827
|
+
if (ctx && !ctx.injected) ctx.device.destroy();
|
|
1828
|
+
}).catch(() => {});
|
|
1829
|
+
if (backendOpt === "webgpu") throw e;
|
|
1830
|
+
this.logger.warn(`WebGPU unavailable (${e.message}); using CPU`);
|
|
1831
|
+
}
|
|
1832
|
+
}
|
|
1833
|
+
}
|
|
1834
|
+
}
|
|
1835
|
+
if (!(useGpuKnn && useGpuLayout) && (o.backend === void 0 || o.backend === "auto" || o.backend === "cpu") && (o.concurrency ?? 1) !== 1 && !(o.deterministic ?? false) && dense !== null && typeof metric === "string" && isNamedMetric(metric) && !o.precomputedKnn) {
|
|
1836
|
+
this._parallel = await createParallelContext(dense, metric, o.metricOptions ?? {}, o.concurrency ?? "auto", o.workerUrl, this.logger);
|
|
1837
|
+
if (this._parallel) this.logger.info(`worker pool: ${this._parallel.pool.size} workers (${this._parallel.sab ? "SharedArrayBuffer" : "transfer fallback"})`);
|
|
1838
|
+
}
|
|
1839
|
+
const t0 = performance.now();
|
|
1840
|
+
let knn;
|
|
1841
|
+
if (o.precomputedKnn) {
|
|
1842
|
+
const pk = o.precomputedKnn;
|
|
1843
|
+
if (pk.k < nNeighbors) throw new UmapValidationError(`precomputedKnn.k=${pk.k} is smaller than nNeighbors=${nNeighbors}`);
|
|
1844
|
+
knn = sliceKnn(pk, n, nNeighbors);
|
|
1845
|
+
if (this._disconnection !== Infinity && knn.indices === pk.indices) knn = {
|
|
1846
|
+
indices: knn.indices.slice(),
|
|
1847
|
+
distances: knn.distances.slice(),
|
|
1848
|
+
k: knn.k
|
|
1849
|
+
};
|
|
1850
|
+
applyDisconnectionToKnn(knn, this._disconnection);
|
|
1851
|
+
this._backendInfo.knn = "precomputed";
|
|
1852
|
+
this._smallData = false;
|
|
1853
|
+
} else if (metric === "precomputed") {
|
|
1854
|
+
if (!dense) throw new UmapValidationError("metric 'precomputed' requires a dense matrix");
|
|
1855
|
+
knn = knnFromPrecomputed(applyDisconnectionToMatrix(dense, this._disconnection), nNeighbors);
|
|
1856
|
+
this._backendInfo.knn = "precomputed-matrix";
|
|
1857
|
+
this._smallData = true;
|
|
1858
|
+
} else if (sparse) {
|
|
1859
|
+
const res = nearestNeighborsSparse(sparse, nNeighbors, metric, rng, {
|
|
1860
|
+
forceApproximation: o.forceApproximation ?? false,
|
|
1861
|
+
angularRpForest: o.angularRpForest ?? false,
|
|
1862
|
+
disconnectionDistance: this._disconnection
|
|
1863
|
+
});
|
|
1864
|
+
knn = res.knn;
|
|
1865
|
+
this._backendInfo.knn = `cpu-${res.method}`;
|
|
1866
|
+
} else {
|
|
1867
|
+
const runCpuKnn = () => {
|
|
1868
|
+
const res = nearestNeighborsDense(dense, nNeighbors, metric, rng, {
|
|
1869
|
+
metricOptions: o.metricOptions ?? {},
|
|
1870
|
+
forceApproximation: o.forceApproximation ?? false,
|
|
1871
|
+
angularRpForest: o.angularRpForest ?? false,
|
|
1872
|
+
disconnectionDistance: this._disconnection
|
|
1873
|
+
});
|
|
1874
|
+
this._searchIndex = res.searchIndex;
|
|
1875
|
+
this._backendInfo.knn = `cpu-${res.method}`;
|
|
1876
|
+
return res.knn;
|
|
1877
|
+
};
|
|
1878
|
+
if (typeof metric === "function") {
|
|
1879
|
+
knn = exactKnnDense(dense, nNeighbors, metric, { disconnectionDistance: this._disconnection });
|
|
1880
|
+
this._backendInfo.knn = "cpu-exact-custom";
|
|
1881
|
+
this._smallData = true;
|
|
1882
|
+
} else if (useGpuKnn && this._gpuCtx && (!this._smallData || backendOpt === "webgpu")) {
|
|
1883
|
+
let gpuKnnResult = null;
|
|
1884
|
+
try {
|
|
1885
|
+
const plan = this._gpu.gpuMetricPlan(metric);
|
|
1886
|
+
gpuKnnResult = await this._gpu.gpuKnn(this._gpuCtx, dense, null, nNeighbors, plan, { ...signal ? { signal } : {} });
|
|
1887
|
+
} catch (e) {
|
|
1888
|
+
if (backendOpt === "webgpu" || !(e instanceof UmapBackendError)) throw e;
|
|
1889
|
+
this.logger.warn(`WebGPU kNN failed (${e.message}); falling back to CPU`);
|
|
1890
|
+
this.dropGpu();
|
|
1891
|
+
}
|
|
1892
|
+
if (gpuKnnResult) {
|
|
1893
|
+
knn = gpuKnnResult;
|
|
1894
|
+
applyDisconnectionToKnn(knn, this._disconnection);
|
|
1895
|
+
this._backendInfo.knn = "webgpu";
|
|
1896
|
+
} else knn = runCpuKnn();
|
|
1897
|
+
} else if (this._parallel && !this._smallData) {
|
|
1898
|
+
const angular = resolveInternalMetric(metric, o.metricOptions ?? {}).angularTrees || (o.angularRpForest ?? false);
|
|
1899
|
+
const par = await parallelNnDescent(this._parallel, nNeighbors, rng, angular);
|
|
1900
|
+
const prepared = prepareIndexFromParallel(dense, par.knnInternal, metric, o.metricOptions ?? {}, nNeighbors, par.searchRngState, angular);
|
|
1901
|
+
knn = prepared.knn;
|
|
1902
|
+
applyDisconnectionToKnn(knn, this._disconnection);
|
|
1903
|
+
this._searchIndex = prepared.searchIndex;
|
|
1904
|
+
this._backendInfo.knn = `cpu-workers(${this._parallel.pool.size})`;
|
|
1905
|
+
} else knn = runCpuKnn();
|
|
1906
|
+
}
|
|
1907
|
+
this._knn = knn;
|
|
1908
|
+
this._timings.knn = performance.now() - t0;
|
|
1909
|
+
this.logger.info(`kNN done in ${this._timings.knn.toFixed(0)}ms (${this._backendInfo.knn})`);
|
|
1910
|
+
throwIfAborted(signal);
|
|
1911
|
+
o.onEpoch?.({
|
|
1912
|
+
epoch: 0,
|
|
1913
|
+
nEpochs: 0,
|
|
1914
|
+
phase: "knn"
|
|
1915
|
+
});
|
|
1916
|
+
const t1 = performance.now();
|
|
1917
|
+
const fss = fuzzySimplicialSet(n, nNeighbors, knn.indices, knn.distances, {
|
|
1918
|
+
setOpMixRatio: o.setOpMixRatio ?? 1,
|
|
1919
|
+
localConnectivity: o.localConnectivity ?? 1,
|
|
1920
|
+
returnDists: densMap || outputDens
|
|
1921
|
+
});
|
|
1922
|
+
let graph = fss.graph;
|
|
1923
|
+
this._sigmas = fss.sigmas;
|
|
1924
|
+
this._rhos = fss.rhos;
|
|
1925
|
+
this._supervised = y != null;
|
|
1926
|
+
if (y != null) {
|
|
1927
|
+
const yArr = normalizeY(y, nFull);
|
|
1928
|
+
const yUnique = this._uniqueIndex ? Float64Array.from(this._uniqueIndex, (i) => yArr[i]) : yArr;
|
|
1929
|
+
if (yUnique.length !== n) throw new UmapValidationError(`y length ${yArr.length} != number of samples ${nFull}`);
|
|
1930
|
+
if ((o.targetMetric ?? "categorical") === "categorical") {
|
|
1931
|
+
const farDist = farDistFromWeight(o.targetWeight ?? .5);
|
|
1932
|
+
graph = categoricalSimplicialSetIntersection(graph, yUnique, { farDist });
|
|
1933
|
+
} else {
|
|
1934
|
+
const targetK = (o.targetNNeighbors ?? -1) === -1 ? nNeighbors : o.targetNNeighbors;
|
|
1935
|
+
const tKnn = exactKnnDense({
|
|
1936
|
+
data: Float32Array.from(yUnique),
|
|
1937
|
+
rows: n,
|
|
1938
|
+
cols: 1
|
|
1939
|
+
}, Math.min(targetK, n - 1), "euclidean");
|
|
1940
|
+
const tFss = fuzzySimplicialSet(n, Math.min(targetK, n - 1), tKnn.indices, tKnn.distances, {
|
|
1941
|
+
setOpMixRatio: 1,
|
|
1942
|
+
localConnectivity: 1,
|
|
1943
|
+
applySetOperations: false
|
|
1944
|
+
});
|
|
1945
|
+
graph = generalSimplicialSetIntersection(graph, tFss.graph, o.targetWeight ?? .5);
|
|
1946
|
+
graph = resetLocalConnectivity(graph);
|
|
1947
|
+
}
|
|
1948
|
+
}
|
|
1949
|
+
this._graph = graph;
|
|
1950
|
+
this._timings.graph = performance.now() - t1;
|
|
1951
|
+
throwIfAborted(signal);
|
|
1952
|
+
await this.embedGraph(graph, fss.dists, dense, nComponents, outputMetric, rng, {
|
|
1953
|
+
densMap,
|
|
1954
|
+
outputDens,
|
|
1955
|
+
...signal ? { signal } : {}
|
|
1956
|
+
});
|
|
1957
|
+
if (this._uniqueInverse && this._embedding) {
|
|
1958
|
+
const inv = this._uniqueInverse;
|
|
1959
|
+
const emb = this._embedding.data;
|
|
1960
|
+
this._embeddingUnique = emb;
|
|
1961
|
+
const full = new Float32Array(nFull * nComponents);
|
|
1962
|
+
for (let i = 0; i < nFull; i++) full.set(emb.subarray(inv[i] * nComponents, (inv[i] + 1) * nComponents), i * nComponents);
|
|
1963
|
+
this._embedding = {
|
|
1964
|
+
data: full,
|
|
1965
|
+
rows: nFull,
|
|
1966
|
+
cols: nComponents
|
|
1967
|
+
};
|
|
1968
|
+
const expand = (arr) => {
|
|
1969
|
+
if (!arr) return null;
|
|
1970
|
+
const out = new Float32Array(nFull);
|
|
1971
|
+
for (let i = 0; i < nFull; i++) out[i] = arr[inv[i]];
|
|
1972
|
+
return out;
|
|
1973
|
+
};
|
|
1974
|
+
this._radOrig = expand(this._radOrig);
|
|
1975
|
+
this._radEmb = expand(this._radEmb);
|
|
1976
|
+
} else this._embeddingUnique = this._embedding?.data ?? null;
|
|
1977
|
+
return this;
|
|
1978
|
+
}
|
|
1979
|
+
async embedGraph(graph, graphDists, dense, nComponents, outputMetric, rng, { densMap, outputDens, signal, initOverride, nEpochsOverride }) {
|
|
1980
|
+
const o = this.options;
|
|
1981
|
+
const n = graph.rows;
|
|
1982
|
+
const defEpochs = defaultNEpochs(n, densMap);
|
|
1983
|
+
const nEpochs = nEpochsOverride ?? o.nEpochs ?? defEpochs;
|
|
1984
|
+
const pruned = pruneGraphForEpochs(graph, nEpochs, defEpochs);
|
|
1985
|
+
const nnz = pruned.indptr[n];
|
|
1986
|
+
const head = new Int32Array(nnz);
|
|
1987
|
+
const tail = new Int32Array(nnz);
|
|
1988
|
+
const weights = new Float32Array(nnz);
|
|
1989
|
+
for (let r = 0; r < n; r++) for (let p = pruned.indptr[r]; p < pruned.indptr[r + 1]; p++) {
|
|
1990
|
+
head[p] = r;
|
|
1991
|
+
tail[p] = pruned.indices[p];
|
|
1992
|
+
weights[p] = Math.fround(pruned.data[p]);
|
|
1993
|
+
}
|
|
1994
|
+
const t2 = performance.now();
|
|
1995
|
+
let emb;
|
|
1996
|
+
if (initOverride) emb = initOverride;
|
|
1997
|
+
else {
|
|
1998
|
+
const initMethod = o.init ?? "spectral";
|
|
1999
|
+
let gpuInit = null;
|
|
2000
|
+
if (initMethod === "spectral" && this._useGpuLayout && this._gpuCtx) {
|
|
2001
|
+
const mod = await import("./webgpu.js");
|
|
2002
|
+
const rngBackup = rng.clone();
|
|
2003
|
+
try {
|
|
2004
|
+
gpuInit = await mod.gpuSpectralInit(this._gpuCtx, pruned, nComponents, rng, this.logger, signal);
|
|
2005
|
+
} catch (e) {
|
|
2006
|
+
if (!(e instanceof UmapBackendError) || (o.backend ?? "auto") === "webgpu") throw e;
|
|
2007
|
+
this.logger.warn(`WebGPU spectral init failed (${e.message}); using the CPU init`);
|
|
2008
|
+
rng.restoreFrom(rngBackup);
|
|
2009
|
+
this.dropGpu();
|
|
2010
|
+
}
|
|
2011
|
+
}
|
|
2012
|
+
emb = gpuInit?.embedding ?? initializeEmbedding(initMethod, dense, pruned, nComponents, rng, { logger: this.logger }).embedding;
|
|
2013
|
+
}
|
|
2014
|
+
this._timings.init = performance.now() - t2;
|
|
2015
|
+
throwIfAborted(signal);
|
|
2016
|
+
o.onEpoch?.({
|
|
2017
|
+
epoch: 0,
|
|
2018
|
+
nEpochs,
|
|
2019
|
+
phase: "init"
|
|
2020
|
+
});
|
|
2021
|
+
const epochsPerSample = makeEpochsPerSample(weights, nEpochs);
|
|
2022
|
+
const rngState = rng.tauState();
|
|
2023
|
+
let densParams = null;
|
|
2024
|
+
if ((densMap || outputDens) && graphDists) {
|
|
2025
|
+
const radii = computeOrigRadii(head, tail, weights, graphDists, n);
|
|
2026
|
+
this._radOrig = radii.ro;
|
|
2027
|
+
if (densMap) densParams = {
|
|
2028
|
+
lambda: o.densLambda ?? 2,
|
|
2029
|
+
frac: o.densFrac ?? .3,
|
|
2030
|
+
varShift: o.densVarShift ?? .1,
|
|
2031
|
+
mu: weights,
|
|
2032
|
+
muSum: radii.muSum,
|
|
2033
|
+
R: radii.R
|
|
2034
|
+
};
|
|
2035
|
+
}
|
|
2036
|
+
const t3 = performance.now();
|
|
2037
|
+
const onEpoch = o.onEpoch;
|
|
2038
|
+
const layoutOpts = {
|
|
2039
|
+
gamma: o.repulsionStrength ?? 1,
|
|
2040
|
+
initialAlpha: o.learningRate ?? 1,
|
|
2041
|
+
negativeSampleRate: o.negativeSampleRate ?? 5,
|
|
2042
|
+
moveOther: true,
|
|
2043
|
+
densmap: densParams,
|
|
2044
|
+
snapshotEvery: onEpoch ? 10 : 0,
|
|
2045
|
+
...onEpoch ? { onEpoch: (info) => onEpoch(info) } : {},
|
|
2046
|
+
...signal ? { signal } : {}
|
|
2047
|
+
};
|
|
2048
|
+
let gpuLayoutDone = false;
|
|
2049
|
+
if (this._useGpuLayout && this._gpuCtx && isEuclideanOutput(outputMetric) && !densParams) {
|
|
2050
|
+
const gpuRng = rng;
|
|
2051
|
+
const rngBackup = rng.clone();
|
|
2052
|
+
try {
|
|
2053
|
+
await this._gpu.gpuOptimizeLayout(this._gpuCtx, emb, nComponents, head, tail, nEpochs, n, epochsPerSample, this._a, this._b, gpuRng, {
|
|
2054
|
+
gamma: o.repulsionStrength ?? 1,
|
|
2055
|
+
initialAlpha: o.learningRate ?? 1,
|
|
2056
|
+
negativeSampleRate: o.negativeSampleRate ?? 5,
|
|
2057
|
+
moveOther: true,
|
|
2058
|
+
...o.gpuLayoutStrategy ? { strategy: o.gpuLayoutStrategy } : {},
|
|
2059
|
+
...onEpoch ? {
|
|
2060
|
+
onEpoch,
|
|
2061
|
+
snapshotEvery: 10
|
|
2062
|
+
} : {},
|
|
2063
|
+
...signal ? { signal } : {}
|
|
2064
|
+
});
|
|
2065
|
+
this._backendInfo.layout = "webgpu";
|
|
2066
|
+
gpuLayoutDone = true;
|
|
2067
|
+
} catch (e) {
|
|
2068
|
+
if (!(e instanceof UmapBackendError) || (o.backend ?? "auto") === "webgpu") throw e;
|
|
2069
|
+
this.logger.warn(`WebGPU layout failed (${e.message}); falling back to CPU layout`);
|
|
2070
|
+
rng.restoreFrom(rngBackup);
|
|
2071
|
+
this.dropGpu();
|
|
2072
|
+
}
|
|
2073
|
+
}
|
|
2074
|
+
if (gpuLayoutDone) {} else if (this._parallel && isEuclideanOutput(outputMetric) && !densParams) {
|
|
2075
|
+
await parallelOptimizeLayout(this._parallel, emb, nComponents, head, tail, nEpochs, n, epochsPerSample, this._a, this._b, rngState, {
|
|
2076
|
+
gamma: layoutOpts.gamma,
|
|
2077
|
+
initialAlpha: layoutOpts.initialAlpha,
|
|
2078
|
+
negativeSampleRate: layoutOpts.negativeSampleRate,
|
|
2079
|
+
moveOther: true,
|
|
2080
|
+
...onEpoch ? { onEpoch: layoutOpts.onEpoch } : {},
|
|
2081
|
+
...signal ? { signal } : {},
|
|
2082
|
+
snapshotEvery: layoutOpts.snapshotEvery
|
|
2083
|
+
});
|
|
2084
|
+
this._backendInfo.layout = `cpu-workers(${this._parallel.pool.size})`;
|
|
2085
|
+
} else if (isEuclideanOutput(outputMetric)) {
|
|
2086
|
+
await yieldToLoop();
|
|
2087
|
+
optimizeLayoutEuclidean(emb, emb, nComponents, head, tail, nEpochs, n, epochsPerSample, this._a, this._b, rngState, layoutOpts);
|
|
2088
|
+
} else {
|
|
2089
|
+
const grad = resolveOutputMetricGrad(outputMetric);
|
|
2090
|
+
await yieldToLoop();
|
|
2091
|
+
optimizeLayoutGeneric(emb, emb, nComponents, head, tail, nEpochs, n, epochsPerSample, this._a, this._b, rngState, grad, layoutOpts);
|
|
2092
|
+
}
|
|
2093
|
+
this._timings.layout = performance.now() - t3;
|
|
2094
|
+
if (!this._backendInfo.layout.startsWith("cpu-workers") && this._backendInfo.layout !== "webgpu") this._backendInfo.layout = "cpu-single";
|
|
2095
|
+
this._embedding = {
|
|
2096
|
+
data: emb,
|
|
2097
|
+
rows: n,
|
|
2098
|
+
cols: nComponents
|
|
2099
|
+
};
|
|
2100
|
+
if (outputDens) {
|
|
2101
|
+
const embMatrix = {
|
|
2102
|
+
data: emb,
|
|
2103
|
+
rows: n,
|
|
2104
|
+
cols: nComponents
|
|
2105
|
+
};
|
|
2106
|
+
const kDens = Math.min(o.nNeighbors ?? 15, n - 1);
|
|
2107
|
+
const res = nearestNeighborsDense(embMatrix, kDens, "euclidean", Pcg32.fromSeed((o.seed ?? 0) + 2654435769), {});
|
|
2108
|
+
const fss2 = fuzzySimplicialSet(n, kDens, res.knn.indices, res.knn.distances, { returnDists: true });
|
|
2109
|
+
this._radEmb = computeEmbRadii(fss2.graph, fss2.dists, n);
|
|
2110
|
+
}
|
|
2111
|
+
}
|
|
2112
|
+
async fitTransform(X, y, opts = {}) {
|
|
2113
|
+
await this.fit(X, y, opts);
|
|
2114
|
+
return this.embedding;
|
|
2115
|
+
}
|
|
2116
|
+
async transform(Xnew, { signal } = {}) {
|
|
2117
|
+
return this.withBusy("transform", async () => {
|
|
2118
|
+
const graphOrEmb = await this.transformInner(Xnew, signal);
|
|
2119
|
+
if (graphOrEmb.kind === "embedding") return graphOrEmb.embedding;
|
|
2120
|
+
throw new UmapError("transformMode is \"graph\"; use transformGraph() for the CSR result");
|
|
2121
|
+
});
|
|
2122
|
+
}
|
|
2123
|
+
/** transformMode 'graph': the bipartite membership CSR (nNew × nTrain). */
|
|
2124
|
+
async transformGraph(Xnew, { signal } = {}) {
|
|
2125
|
+
return this.withBusy("transformGraph", async () => {
|
|
2126
|
+
const res = await this.transformInner(Xnew, signal, true);
|
|
2127
|
+
if (res.kind === "graph") return res.graph;
|
|
2128
|
+
throw new UmapError("internal: expected graph result");
|
|
2129
|
+
});
|
|
2130
|
+
}
|
|
2131
|
+
async transformInner(Xnew, signal, forceGraph = false) {
|
|
2132
|
+
if (!this._embedding) throw new UmapError("model is not fitted");
|
|
2133
|
+
if (this.options.densMap) throw new UmapError("transform is unsupported with densMap");
|
|
2134
|
+
if (this._rawCsr) throw new UmapError("transform on sparse-fitted models is not supported");
|
|
2135
|
+
if (!this._rawData) throw new UmapError("transform requires the training data");
|
|
2136
|
+
const o = this.options;
|
|
2137
|
+
const metric = o.metric ?? "euclidean";
|
|
2138
|
+
const Q = toMatrix(Xnew, "Xnew");
|
|
2139
|
+
if (metric !== "precomputed" && Q.cols !== this._rawData.cols) throw new UmapValidationError(`Xnew has ${Q.cols} features but the model was fitted with ${this._rawData.cols}`);
|
|
2140
|
+
throwIfAborted(signal);
|
|
2141
|
+
const graphMode = forceGraph || o.transformMode === "graph";
|
|
2142
|
+
const qHash = contentHash(new Uint8Array(Q.data.buffer, Q.data.byteOffset, Q.data.byteLength));
|
|
2143
|
+
if (!graphMode && qHash === this._inputHash && sameFloat32Bits(Q.data, this._rawData.data)) return {
|
|
2144
|
+
kind: "embedding",
|
|
2145
|
+
embedding: this.embedding
|
|
2146
|
+
};
|
|
2147
|
+
const k = this._nNeighborsUsed;
|
|
2148
|
+
const nTrainUnique = this._graph.rows;
|
|
2149
|
+
const trainEmb = this._embeddingUnique;
|
|
2150
|
+
const nComponents = this._embedding.cols;
|
|
2151
|
+
const transformSeed = o.transformSeed ?? 42;
|
|
2152
|
+
const rng = Pcg32.fromSeed(transformSeed);
|
|
2153
|
+
if (!this._searchIndex && !this._smallData && this._knn && typeof metric === "string") this._searchIndex = this.buildSearchIndexFromKnn(metric);
|
|
2154
|
+
let qKnn;
|
|
2155
|
+
if (metric === "precomputed") {
|
|
2156
|
+
if (Q.cols !== nTrainUnique) throw new UmapValidationError(`precomputed transform input must be (nNew × ${nTrainUnique})`);
|
|
2157
|
+
qKnn = precomputedQueryKnn(Q, k);
|
|
2158
|
+
} else if (this._smallData || !this._searchIndex) qKnn = smallDataQueryKnn(Q, this.uniqueTrainMatrix(), k, typeof metric === "function" ? metric : resolveMetric(metric, o.metricOptions ?? {}));
|
|
2159
|
+
else {
|
|
2160
|
+
const epsilon = this._searchIndex.angular ? .24 : .12;
|
|
2161
|
+
qKnn = searchKnn(this._searchIndex, Q, k, epsilon);
|
|
2162
|
+
}
|
|
2163
|
+
throwIfAborted(signal);
|
|
2164
|
+
for (let t = 0; t < qKnn.distances.length; t++) if (qKnn.distances[t] >= this._disconnection) {
|
|
2165
|
+
qKnn.indices[t] = -1;
|
|
2166
|
+
qKnn.distances[t] = Infinity;
|
|
2167
|
+
}
|
|
2168
|
+
const adjustedLc = Math.max(0, (o.localConnectivity ?? 1) - 1);
|
|
2169
|
+
const { sigmas, rhos } = smoothKnnDist(qKnn.distances, k, k, { localConnectivity: adjustedLc });
|
|
2170
|
+
const { rows, cols, vals } = computeMembershipStrengths(qKnn.indices, qKnn.distances, sigmas, rhos, k, false, true);
|
|
2171
|
+
let graph = csrFromCoo(Q.rows, nTrainUnique, rows, cols, vals);
|
|
2172
|
+
graph = csrEliminateZeros(graph);
|
|
2173
|
+
if (graphMode) return {
|
|
2174
|
+
kind: "graph",
|
|
2175
|
+
graph
|
|
2176
|
+
};
|
|
2177
|
+
const emb = initGraphTransform(graph, trainEmb, nComponents);
|
|
2178
|
+
let nEpochs;
|
|
2179
|
+
if (o.nEpochs === void 0) nEpochs = Q.rows <= 1e4 ? 100 : 30;
|
|
2180
|
+
else nEpochs = Math.floor(o.nEpochs / 3);
|
|
2181
|
+
let wMax = -Infinity;
|
|
2182
|
+
for (let t = 0; t < graph.data.length; t++) if (graph.data[t] > wMax) wMax = graph.data[t];
|
|
2183
|
+
const thresh = wMax / nEpochs;
|
|
2184
|
+
const gData = graph.data.slice();
|
|
2185
|
+
for (let t = 0; t < gData.length; t++) if (gData[t] < thresh) gData[t] = 0;
|
|
2186
|
+
graph = csrEliminateZeros({
|
|
2187
|
+
...graph,
|
|
2188
|
+
data: gData
|
|
2189
|
+
});
|
|
2190
|
+
const nnz = graph.indptr[Q.rows];
|
|
2191
|
+
const head = new Int32Array(nnz);
|
|
2192
|
+
const tail = new Int32Array(nnz);
|
|
2193
|
+
const weights = new Float32Array(nnz);
|
|
2194
|
+
for (let r = 0; r < Q.rows; r++) for (let p = graph.indptr[r]; p < graph.indptr[r + 1]; p++) {
|
|
2195
|
+
head[p] = r;
|
|
2196
|
+
tail[p] = graph.indices[p];
|
|
2197
|
+
weights[p] = Math.fround(graph.data[p]);
|
|
2198
|
+
}
|
|
2199
|
+
const epochsPerSample = makeEpochsPerSample(weights, nEpochs);
|
|
2200
|
+
const rngState = rng.tauState();
|
|
2201
|
+
const tailEmbedding = trainEmb.slice();
|
|
2202
|
+
const outputMetric = o.outputMetric ?? "euclidean";
|
|
2203
|
+
const layoutOpts = {
|
|
2204
|
+
gamma: o.repulsionStrength ?? 1,
|
|
2205
|
+
initialAlpha: (o.learningRate ?? 1) / 4,
|
|
2206
|
+
negativeSampleRate: o.negativeSampleRate ?? 5,
|
|
2207
|
+
moveOther: false,
|
|
2208
|
+
...signal ? { signal } : {}
|
|
2209
|
+
};
|
|
2210
|
+
await yieldToLoop();
|
|
2211
|
+
if (isEuclideanOutput(outputMetric)) optimizeLayoutEuclidean(emb, tailEmbedding, nComponents, head, tail, nEpochs, nTrainUnique, epochsPerSample, this._a, this._b, rngState, layoutOpts);
|
|
2212
|
+
else optimizeLayoutGeneric(emb, tailEmbedding, nComponents, head, tail, nEpochs, nTrainUnique, epochsPerSample, this._a, this._b, rngState, resolveOutputMetricGrad(outputMetric), layoutOpts);
|
|
2213
|
+
return {
|
|
2214
|
+
kind: "embedding",
|
|
2215
|
+
embedding: {
|
|
2216
|
+
data: emb,
|
|
2217
|
+
rows: Q.rows,
|
|
2218
|
+
cols: nComponents
|
|
2219
|
+
}
|
|
2220
|
+
};
|
|
2221
|
+
}
|
|
2222
|
+
/**
|
|
2223
|
+
* Build the prepared search structure from the stored kNN graph (D-018):
|
|
2224
|
+
* used when the fit produced no index (GPU kNN path, deserialized blobs
|
|
2225
|
+
* without one). Internal-space distances are recomputed exactly via the
|
|
2226
|
+
* internal metric — n·k evaluations, no correction inversion. The rng for
|
|
2227
|
+
* hub-tree tie-breaks/query descent derives from transformSeed on its own
|
|
2228
|
+
* PCG stream, so fit-time determinism (D-013) is untouched and transforms
|
|
2229
|
+
* stay reproducible for a given transformSeed.
|
|
2230
|
+
*/
|
|
2231
|
+
buildSearchIndexFromKnn(metricName) {
|
|
2232
|
+
if (!this._knn || !this._rawData || !isNamedMetric(metricName)) return null;
|
|
2233
|
+
const t0 = performance.now();
|
|
2234
|
+
const X = this.uniqueTrainMatrix();
|
|
2235
|
+
const internal = resolveInternalMetric(metricName, this.options.metricOptions ?? {});
|
|
2236
|
+
const k = this._knn.k;
|
|
2237
|
+
const n = X.rows;
|
|
2238
|
+
const d = X.cols;
|
|
2239
|
+
const indices = this._knn.indices;
|
|
2240
|
+
const distInternal = new Float32Array(n * k);
|
|
2241
|
+
for (let i = 0; i < n; i++) {
|
|
2242
|
+
const xi = X.data.subarray(i * d, (i + 1) * d);
|
|
2243
|
+
for (let j = 0; j < k; j++) {
|
|
2244
|
+
const t = i * k + j;
|
|
2245
|
+
const nb = indices[t];
|
|
2246
|
+
distInternal[t] = nb < 0 ? Infinity : internal.dist(xi, X.data.subarray(nb * d, (nb + 1) * d));
|
|
2247
|
+
}
|
|
2248
|
+
}
|
|
2249
|
+
const knnInternal = {
|
|
2250
|
+
indices,
|
|
2251
|
+
distances: distInternal,
|
|
2252
|
+
k
|
|
2253
|
+
};
|
|
2254
|
+
const searchRngState = Pcg32.fromSeed(this.options.transformSeed ?? 42, 63).tauState();
|
|
2255
|
+
const searchIndex = prepareSearchIndex(X, knnInternal, (a, b) => internal.dist(a, b), internal.correction, internal.angularTrees || (this.options.angularRpForest ?? false), this._nNeighborsUsed, searchRngState);
|
|
2256
|
+
this.logger.info(`built the transform search index lazily in ${(performance.now() - t0).toFixed(0)}ms`);
|
|
2257
|
+
return searchIndex;
|
|
2258
|
+
}
|
|
2259
|
+
uniqueTrainMatrix() {
|
|
2260
|
+
if (!this._uniqueIndex || !this._rawData) return this._rawData;
|
|
2261
|
+
const d = this._rawData.cols;
|
|
2262
|
+
const data = new Float32Array(this._uniqueIndex.length * d);
|
|
2263
|
+
for (let i = 0; i < this._uniqueIndex.length; i++) data.set(this._rawData.data.subarray(this._uniqueIndex[i] * d, (this._uniqueIndex[i] + 1) * d), i * d);
|
|
2264
|
+
return {
|
|
2265
|
+
data,
|
|
2266
|
+
rows: this._uniqueIndex.length,
|
|
2267
|
+
cols: d
|
|
2268
|
+
};
|
|
2269
|
+
}
|
|
2270
|
+
async inverseTransform(points, { signal } = {}) {
|
|
2271
|
+
return this.withBusy("inverseTransform", () => this.inverseTransformInner(points, signal));
|
|
2272
|
+
}
|
|
2273
|
+
async inverseTransformInner(points, signal) {
|
|
2274
|
+
throwIfAborted(signal);
|
|
2275
|
+
if (!this._embedding) throw new UmapError("model is not fitted");
|
|
2276
|
+
if (this._rawCsr) throw new UmapValidationError("inverseTransform is unsupported for sparse-fitted models");
|
|
2277
|
+
if (this.options.densMap) throw new UmapValidationError("inverseTransform is unsupported with densMap");
|
|
2278
|
+
if (!isEuclideanOutput(this.options.outputMetric ?? "euclidean")) throw new UmapValidationError("inverseTransform requires euclidean output");
|
|
2279
|
+
const metric = this.options.metric ?? "euclidean";
|
|
2280
|
+
if (typeof metric !== "function" && !["euclidean", "l2"].includes(metric)) throw new UmapValidationError(`inverseTransform supports euclidean data metric in this port, got '${metric}'`);
|
|
2281
|
+
const P = toMatrix(points, "points");
|
|
2282
|
+
const trainEmb = this._embeddingUnique;
|
|
2283
|
+
const raw = this.uniqueTrainMatrix();
|
|
2284
|
+
const nTrain = raw.rows;
|
|
2285
|
+
const dim = this._embedding.cols;
|
|
2286
|
+
const dataDim = raw.cols;
|
|
2287
|
+
if (P.cols !== dim) throw new UmapValidationError(`points must be (n × ${dim})`);
|
|
2288
|
+
const rng = Pcg32.fromSeed(this.options.transformSeed ?? 42);
|
|
2289
|
+
const minVertices = Math.min(dataDim, nTrain);
|
|
2290
|
+
const nn = smallDataQueryKnn(P, {
|
|
2291
|
+
data: trainEmb,
|
|
2292
|
+
rows: nTrain,
|
|
2293
|
+
cols: dim
|
|
2294
|
+
}, minVertices, resolveMetric("euclidean"));
|
|
2295
|
+
throwIfAborted(signal);
|
|
2296
|
+
const weights = new Float32Array(P.rows * minVertices);
|
|
2297
|
+
for (let i = 0; i < P.rows; i++) {
|
|
2298
|
+
let rowSum = 0;
|
|
2299
|
+
for (let j = 0; j < minVertices; j++) {
|
|
2300
|
+
const d = nn.distances[i * minVertices + j];
|
|
2301
|
+
const w = 1 / (1 + this._a * d ** (2 * this._b));
|
|
2302
|
+
weights[i * minVertices + j] = w;
|
|
2303
|
+
rowSum += w;
|
|
2304
|
+
}
|
|
2305
|
+
for (let j = 0; j < minVertices; j++) weights[i * minVertices + j] = weights[i * minVertices + j] / rowSum;
|
|
2306
|
+
}
|
|
2307
|
+
const inv = initTransform(nn.indices, weights, minVertices, raw.data, dataDim);
|
|
2308
|
+
let nEpochs;
|
|
2309
|
+
if (this.options.nEpochs === void 0) nEpochs = P.rows <= 1e4 ? 100 : 30;
|
|
2310
|
+
else nEpochs = Math.floor(this.options.nEpochs / 3);
|
|
2311
|
+
const kernelW = new Float32Array(P.rows * minVertices);
|
|
2312
|
+
for (let i = 0; i < P.rows; i++) for (let j = 0; j < minVertices; j++) {
|
|
2313
|
+
const d = nn.distances[i * minVertices + j];
|
|
2314
|
+
kernelW[i * minVertices + j] = Math.fround(1 / (1 + this._a * d ** (2 * this._b)));
|
|
2315
|
+
}
|
|
2316
|
+
const head = new Int32Array(P.rows * minVertices);
|
|
2317
|
+
const tail = new Int32Array(P.rows * minVertices);
|
|
2318
|
+
for (let i = 0; i < P.rows; i++) for (let j = 0; j < minVertices; j++) {
|
|
2319
|
+
head[i * minVertices + j] = i;
|
|
2320
|
+
tail[i * minVertices + j] = nn.indices[i * minVertices + j];
|
|
2321
|
+
}
|
|
2322
|
+
const epochsPerSample = makeEpochsPerSample(kernelW, nEpochs);
|
|
2323
|
+
const rngState = rng.tauState();
|
|
2324
|
+
await yieldToLoop();
|
|
2325
|
+
optimizeLayoutInverse(inv, raw.data, dataDim, head, tail, kernelW, this._sigmas, this._rhos, nEpochs, nTrain, epochsPerSample, rngState, resolveOutputMetricGrad("euclidean"), {
|
|
2326
|
+
gamma: this.options.repulsionStrength ?? 1,
|
|
2327
|
+
initialAlpha: (this.options.learningRate ?? 1) / 4,
|
|
2328
|
+
negativeSampleRate: this.options.negativeSampleRate ?? 5,
|
|
2329
|
+
...signal ? { signal } : {}
|
|
2330
|
+
});
|
|
2331
|
+
return {
|
|
2332
|
+
data: inv,
|
|
2333
|
+
rows: P.rows,
|
|
2334
|
+
cols: dataDim
|
|
2335
|
+
};
|
|
2336
|
+
}
|
|
2337
|
+
async update(Xmore, { signal } = {}) {
|
|
2338
|
+
return this.withBusy("update", async () => {
|
|
2339
|
+
const snapshot = this.snapshotFittedState();
|
|
2340
|
+
try {
|
|
2341
|
+
return await this.updateInner(Xmore, signal);
|
|
2342
|
+
} catch (e) {
|
|
2343
|
+
this.restoreFittedState(snapshot);
|
|
2344
|
+
throw e;
|
|
2345
|
+
}
|
|
2346
|
+
});
|
|
2347
|
+
}
|
|
2348
|
+
async updateInner(Xmore, signal) {
|
|
2349
|
+
throwIfAborted(signal);
|
|
2350
|
+
if (!this._embedding) throw new UmapError("model is not fitted");
|
|
2351
|
+
if (this.options.metric === "precomputed") throw new UmapValidationError("update is unsupported with a precomputed metric");
|
|
2352
|
+
if (this._supervised) throw new UmapValidationError("update is unsupported for supervised models");
|
|
2353
|
+
if (this._uniqueIndex) throw new UmapError("update after unique=true is unsupported");
|
|
2354
|
+
if (this._rawCsr) throw new UmapError("update on sparse-fitted models is unsupported");
|
|
2355
|
+
const o = this.options;
|
|
2356
|
+
const Xn = toMatrix(Xmore, "Xmore");
|
|
2357
|
+
if (Xn.cols !== this._rawData.cols) throw new UmapValidationError("Xmore feature count differs from training data");
|
|
2358
|
+
const originalSize = this._rawData.rows;
|
|
2359
|
+
const merged = new Float32Array((originalSize + Xn.rows) * Xn.cols);
|
|
2360
|
+
merged.set(this._rawData.data, 0);
|
|
2361
|
+
merged.set(Xn.data, originalSize * Xn.cols);
|
|
2362
|
+
const X = {
|
|
2363
|
+
data: merged,
|
|
2364
|
+
rows: originalSize + Xn.rows,
|
|
2365
|
+
cols: Xn.cols
|
|
2366
|
+
};
|
|
2367
|
+
this._rawData = X;
|
|
2368
|
+
const n = X.rows;
|
|
2369
|
+
const rng = Pcg32.fromSeed(this.options.transformSeed ?? 42);
|
|
2370
|
+
const metric = o.metric ?? "euclidean";
|
|
2371
|
+
const nNeighbors = this._nNeighborsUsed;
|
|
2372
|
+
const res = nearestNeighborsDense(X, nNeighbors, metric, rng, {
|
|
2373
|
+
metricOptions: o.metricOptions ?? {},
|
|
2374
|
+
forceApproximation: (o.forceApproximation ?? false) || !this._smallData,
|
|
2375
|
+
angularRpForest: o.angularRpForest ?? false,
|
|
2376
|
+
disconnectionDistance: this._disconnection
|
|
2377
|
+
});
|
|
2378
|
+
this._smallData = res.method === "exact";
|
|
2379
|
+
this._searchIndex = res.searchIndex;
|
|
2380
|
+
this._knn = res.knn;
|
|
2381
|
+
const fss = fuzzySimplicialSet(n, nNeighbors, res.knn.indices, res.knn.distances, {
|
|
2382
|
+
setOpMixRatio: o.setOpMixRatio ?? 1,
|
|
2383
|
+
localConnectivity: o.localConnectivity ?? 1,
|
|
2384
|
+
returnDists: (o.densMap ?? false) || (o.outputDens ?? false)
|
|
2385
|
+
});
|
|
2386
|
+
this._graph = fss.graph;
|
|
2387
|
+
this._sigmas = fss.sigmas;
|
|
2388
|
+
this._rhos = fss.rhos;
|
|
2389
|
+
const nComponents = this._embedding.cols;
|
|
2390
|
+
const init = new Float32Array(n * nComponents);
|
|
2391
|
+
init.set(this._embedding.data, 0);
|
|
2392
|
+
initUpdate(init, nComponents, originalSize, res.knn.indices, nNeighbors);
|
|
2393
|
+
const nEpochs = o.nEpochs === void 0 ? 0 : o.nEpochs;
|
|
2394
|
+
await this.embedGraph(fss.graph, fss.dists, X, nComponents, o.outputMetric ?? "euclidean", rng, {
|
|
2395
|
+
densMap: o.densMap ?? false,
|
|
2396
|
+
outputDens: o.outputDens ?? false,
|
|
2397
|
+
initOverride: init,
|
|
2398
|
+
nEpochsOverride: nEpochs,
|
|
2399
|
+
...signal ? { signal } : {}
|
|
2400
|
+
});
|
|
2401
|
+
this._embeddingUnique = this._embedding.data;
|
|
2402
|
+
this._inputHash = contentHash(new Uint8Array(X.data.buffer, X.data.byteOffset, X.data.byteLength));
|
|
2403
|
+
return this;
|
|
2404
|
+
}
|
|
2405
|
+
serialize() {
|
|
2406
|
+
return serializeModel(this);
|
|
2407
|
+
}
|
|
2408
|
+
static deserialize(blob) {
|
|
2409
|
+
return deserializeModel(blob, UMAP);
|
|
2410
|
+
}
|
|
2411
|
+
};
|
|
2412
|
+
function normalizeY(y, n) {
|
|
2413
|
+
const arr = new Float64Array(n);
|
|
2414
|
+
const src = y;
|
|
2415
|
+
if (src.length !== n) throw new UmapValidationError(`y length ${src.length} != number of samples ${n}`);
|
|
2416
|
+
for (let i = 0; i < n; i++) {
|
|
2417
|
+
const v = src[i];
|
|
2418
|
+
if (!Number.isFinite(v)) throw new UmapValidationError("y contains a non-finite value");
|
|
2419
|
+
arr[i] = v;
|
|
2420
|
+
}
|
|
2421
|
+
return arr;
|
|
2422
|
+
}
|
|
2423
|
+
function uniqueRows(X) {
|
|
2424
|
+
const seen = /* @__PURE__ */ new Map();
|
|
2425
|
+
const index = [];
|
|
2426
|
+
const inverse = new Int32Array(X.rows);
|
|
2427
|
+
for (let i = 0; i < X.rows; i++) {
|
|
2428
|
+
const key = Array.prototype.join.call(X.data.subarray(i * X.cols, (i + 1) * X.cols), ",");
|
|
2429
|
+
let u = seen.get(key);
|
|
2430
|
+
if (u === void 0) {
|
|
2431
|
+
u = index.length;
|
|
2432
|
+
seen.set(key, u);
|
|
2433
|
+
index.push(i);
|
|
2434
|
+
}
|
|
2435
|
+
inverse[i] = u;
|
|
2436
|
+
}
|
|
2437
|
+
return {
|
|
2438
|
+
index: Int32Array.from(index),
|
|
2439
|
+
inverse
|
|
2440
|
+
};
|
|
2441
|
+
}
|
|
2442
|
+
function sliceKnn(pk, n, k) {
|
|
2443
|
+
if (pk.indices.length !== n * pk.k || pk.distances.length !== n * pk.k) throw new UmapValidationError(`precomputedKnn arrays must have n*k = ${n * pk.k} entries (indices: ${pk.indices.length}, distances: ${pk.distances.length})`);
|
|
2444
|
+
if (pk.k === k) return {
|
|
2445
|
+
indices: pk.indices,
|
|
2446
|
+
distances: pk.distances,
|
|
2447
|
+
k
|
|
2448
|
+
};
|
|
2449
|
+
const indices = new Int32Array(n * k);
|
|
2450
|
+
const distances = new Float32Array(n * k);
|
|
2451
|
+
for (let i = 0; i < n; i++) {
|
|
2452
|
+
indices.set(pk.indices.subarray(i * pk.k, i * pk.k + k), i * k);
|
|
2453
|
+
distances.set(pk.distances.subarray(i * pk.k, i * pk.k + k), i * k);
|
|
2454
|
+
}
|
|
2455
|
+
return {
|
|
2456
|
+
indices,
|
|
2457
|
+
distances,
|
|
2458
|
+
k
|
|
2459
|
+
};
|
|
2460
|
+
}
|
|
2461
|
+
function applyDisconnectionToKnn(knn, dc) {
|
|
2462
|
+
if (dc === Infinity) return;
|
|
2463
|
+
for (let t = 0; t < knn.distances.length; t++) if (knn.distances[t] >= dc) {
|
|
2464
|
+
knn.distances[t] = Infinity;
|
|
2465
|
+
knn.indices[t] = -1;
|
|
2466
|
+
}
|
|
2467
|
+
}
|
|
2468
|
+
function applyDisconnectionToMatrix(X, dc) {
|
|
2469
|
+
if (dc === Infinity) return X;
|
|
2470
|
+
const data = X.data.slice();
|
|
2471
|
+
for (let i = 0; i < data.length; i++) if (data[i] >= dc) data[i] = Infinity;
|
|
2472
|
+
return {
|
|
2473
|
+
data,
|
|
2474
|
+
rows: X.rows,
|
|
2475
|
+
cols: X.cols
|
|
2476
|
+
};
|
|
2477
|
+
}
|
|
2478
|
+
/** Query kNN of Q against T by brute force (transform small-data path). */
|
|
2479
|
+
function smallDataQueryKnn(Q, T, k, metric) {
|
|
2480
|
+
const nq = Q.rows;
|
|
2481
|
+
const nt = T.rows;
|
|
2482
|
+
const d = Q.cols;
|
|
2483
|
+
const indices = new Int32Array(nq * k);
|
|
2484
|
+
const distances = new Float32Array(nq * k);
|
|
2485
|
+
const row = new Float64Array(nt);
|
|
2486
|
+
const heapIdx = new Int32Array(k);
|
|
2487
|
+
const heapDst = new Float64Array(k);
|
|
2488
|
+
for (let i = 0; i < nq; i++) {
|
|
2489
|
+
const qi = Q.data.subarray(i * d, (i + 1) * d);
|
|
2490
|
+
for (let j = 0; j < nt; j++) row[j] = metric(qi, T.data.subarray(j * d, (j + 1) * d));
|
|
2491
|
+
selectK(row, nt, k, heapIdx, heapDst);
|
|
2492
|
+
for (let j = 0; j < k; j++) {
|
|
2493
|
+
indices[i * k + j] = heapIdx[j];
|
|
2494
|
+
distances[i * k + j] = heapDst[j];
|
|
2495
|
+
}
|
|
2496
|
+
}
|
|
2497
|
+
return {
|
|
2498
|
+
indices,
|
|
2499
|
+
distances,
|
|
2500
|
+
k
|
|
2501
|
+
};
|
|
2502
|
+
}
|
|
2503
|
+
function selectK(rowDist, n, k, outIdx, outDst) {
|
|
2504
|
+
const gt = (d1, i1, d2, i2) => d1 > d2 || d1 === d2 && i1 > i2;
|
|
2505
|
+
let count = 0;
|
|
2506
|
+
for (let j = 0; j < n; j++) {
|
|
2507
|
+
const dj = rowDist[j];
|
|
2508
|
+
if (count < k) {
|
|
2509
|
+
let c = count++;
|
|
2510
|
+
outDst[c] = dj;
|
|
2511
|
+
outIdx[c] = j;
|
|
2512
|
+
while (c > 0) {
|
|
2513
|
+
const parent = c - 1 >> 1;
|
|
2514
|
+
if (gt(outDst[c], outIdx[c], outDst[parent], outIdx[parent])) {
|
|
2515
|
+
const td = outDst[c];
|
|
2516
|
+
outDst[c] = outDst[parent];
|
|
2517
|
+
outDst[parent] = td;
|
|
2518
|
+
const ti = outIdx[c];
|
|
2519
|
+
outIdx[c] = outIdx[parent];
|
|
2520
|
+
outIdx[parent] = ti;
|
|
2521
|
+
c = parent;
|
|
2522
|
+
} else break;
|
|
2523
|
+
}
|
|
2524
|
+
} else if (gt(outDst[0], outIdx[0], dj, j)) {
|
|
2525
|
+
outDst[0] = dj;
|
|
2526
|
+
outIdx[0] = j;
|
|
2527
|
+
let c = 0;
|
|
2528
|
+
for (;;) {
|
|
2529
|
+
const l = 2 * c + 1;
|
|
2530
|
+
const r = l + 1;
|
|
2531
|
+
let m = c;
|
|
2532
|
+
if (l < k && gt(outDst[l], outIdx[l], outDst[m], outIdx[m])) m = l;
|
|
2533
|
+
if (r < k && gt(outDst[r], outIdx[r], outDst[m], outIdx[m])) m = r;
|
|
2534
|
+
if (m === c) break;
|
|
2535
|
+
const td = outDst[c];
|
|
2536
|
+
outDst[c] = outDst[m];
|
|
2537
|
+
outDst[m] = td;
|
|
2538
|
+
const ti = outIdx[c];
|
|
2539
|
+
outIdx[c] = outIdx[m];
|
|
2540
|
+
outIdx[m] = ti;
|
|
2541
|
+
c = m;
|
|
2542
|
+
}
|
|
2543
|
+
}
|
|
2544
|
+
}
|
|
2545
|
+
for (let j = count; j < k; j++) {
|
|
2546
|
+
outDst[j] = Infinity;
|
|
2547
|
+
outIdx[j] = -1;
|
|
2548
|
+
}
|
|
2549
|
+
const idxArr = Array.from({ length: k }, (_, j) => j).sort((a, b) => outDst[a] - outDst[b] || outIdx[a] - outIdx[b]);
|
|
2550
|
+
const tmpD = Float64Array.from(idxArr, (j) => outDst[j]);
|
|
2551
|
+
const tmpI = Int32Array.from(idxArr, (j) => outIdx[j]);
|
|
2552
|
+
outDst.set(tmpD);
|
|
2553
|
+
outIdx.set(tmpI);
|
|
2554
|
+
}
|
|
2555
|
+
/** transform with metric='precomputed': rows are distances to train points. */
|
|
2556
|
+
function precomputedQueryKnn(Q, k) {
|
|
2557
|
+
const indices = new Int32Array(Q.rows * k);
|
|
2558
|
+
const distances = new Float32Array(Q.rows * k);
|
|
2559
|
+
const row = new Float64Array(Q.cols);
|
|
2560
|
+
const heapIdx = new Int32Array(k);
|
|
2561
|
+
const heapDst = new Float64Array(k);
|
|
2562
|
+
for (let i = 0; i < Q.rows; i++) {
|
|
2563
|
+
for (let j = 0; j < Q.cols; j++) row[j] = Q.data[i * Q.cols + j];
|
|
2564
|
+
selectK(row, Q.cols, k, heapIdx, heapDst);
|
|
2565
|
+
for (let j = 0; j < k; j++) {
|
|
2566
|
+
indices[i * k + j] = heapIdx[j];
|
|
2567
|
+
distances[i * k + j] = heapDst[j];
|
|
2568
|
+
}
|
|
2569
|
+
}
|
|
2570
|
+
return {
|
|
2571
|
+
indices,
|
|
2572
|
+
distances,
|
|
2573
|
+
k
|
|
2574
|
+
};
|
|
2575
|
+
}
|
|
2576
|
+
|
|
2577
|
+
//#endregion
|
|
2578
|
+
//#region src/estimator/aligned.ts
|
|
2579
|
+
function relGet(rel, k) {
|
|
2580
|
+
return rel instanceof Map ? rel.get(k) : rel[k];
|
|
2581
|
+
}
|
|
2582
|
+
/** expand_relations (aligned_umap.py:43-82). */
|
|
2583
|
+
function expandRelations(relations, windowSize, sliceSizes) {
|
|
2584
|
+
const nSlices = relations.length + 1;
|
|
2585
|
+
const maxN = Math.max(...sliceSizes);
|
|
2586
|
+
const W = 2 * windowSize + 1;
|
|
2587
|
+
const out = [];
|
|
2588
|
+
for (let i = 0; i < nSlices; i++) {
|
|
2589
|
+
const arr = new Int32Array(W * maxN).fill(-1);
|
|
2590
|
+
for (let off = -windowSize; off <= windowSize; off++) {
|
|
2591
|
+
const t = i + off;
|
|
2592
|
+
if (t < 0 || t >= nSlices || off === 0) {
|
|
2593
|
+
if (off === 0) for (let v = 0; v < sliceSizes[i]; v++) arr[(windowSize + 0) * maxN + v] = v;
|
|
2594
|
+
continue;
|
|
2595
|
+
}
|
|
2596
|
+
for (let v = 0; v < sliceSizes[i]; v++) {
|
|
2597
|
+
let cur = v;
|
|
2598
|
+
let ok = true;
|
|
2599
|
+
if (off > 0) for (let s = i; s < i + off; s++) {
|
|
2600
|
+
const nxt = relGet(relations[s], cur);
|
|
2601
|
+
if (nxt === void 0) {
|
|
2602
|
+
ok = false;
|
|
2603
|
+
break;
|
|
2604
|
+
}
|
|
2605
|
+
cur = nxt;
|
|
2606
|
+
}
|
|
2607
|
+
else for (let s = i - 1; s >= i + off; s--) {
|
|
2608
|
+
const rel = relations[s];
|
|
2609
|
+
let found = -1;
|
|
2610
|
+
if (rel instanceof Map) {
|
|
2611
|
+
for (const [k2, v2] of rel) if (v2 === cur) {
|
|
2612
|
+
found = k2;
|
|
2613
|
+
break;
|
|
2614
|
+
}
|
|
2615
|
+
} else for (const k2 of Object.keys(rel)) if (rel[k2] === cur) {
|
|
2616
|
+
found = Number(k2);
|
|
2617
|
+
break;
|
|
2618
|
+
}
|
|
2619
|
+
if (found < 0) {
|
|
2620
|
+
ok = false;
|
|
2621
|
+
break;
|
|
2622
|
+
}
|
|
2623
|
+
cur = found;
|
|
2624
|
+
}
|
|
2625
|
+
if (ok) arr[(windowSize + off) * maxN + v] = cur;
|
|
2626
|
+
}
|
|
2627
|
+
}
|
|
2628
|
+
out.push(arr);
|
|
2629
|
+
}
|
|
2630
|
+
return out;
|
|
2631
|
+
}
|
|
2632
|
+
/** build_neighborhood_similarities (aligned_umap.py:85-129): Jaccard of mapped kNN sets. */
|
|
2633
|
+
function buildNeighborhoodSimilarities(graphs, relations, windowSize, maxN) {
|
|
2634
|
+
const nSlices = graphs.length;
|
|
2635
|
+
const W = 2 * windowSize + 1;
|
|
2636
|
+
const out = [];
|
|
2637
|
+
for (let i = 0; i < nSlices; i++) {
|
|
2638
|
+
const weights = new Float32Array(W * maxN);
|
|
2639
|
+
for (let off = -windowSize; off <= windowSize; off++) {
|
|
2640
|
+
const t = i + off;
|
|
2641
|
+
if (t < 0 || t >= nSlices || off === 0) continue;
|
|
2642
|
+
const rel = relations[i];
|
|
2643
|
+
const image = /* @__PURE__ */ new Set();
|
|
2644
|
+
for (let v = 0; v < maxN; v++) {
|
|
2645
|
+
const m = rel[(windowSize + off) * maxN + v];
|
|
2646
|
+
if (m >= 0) image.add(m);
|
|
2647
|
+
}
|
|
2648
|
+
for (let v = 0; v < graphs[i].rows; v++) {
|
|
2649
|
+
const counterpart = rel[(windowSize + off) * maxN + v];
|
|
2650
|
+
if (counterpart < 0) continue;
|
|
2651
|
+
const mapped = /* @__PURE__ */ new Set();
|
|
2652
|
+
for (let p = graphs[i].indptr[v]; p < graphs[i].indptr[v + 1]; p++) {
|
|
2653
|
+
const nb = graphs[i].indices[p];
|
|
2654
|
+
if (nb >= maxN) continue;
|
|
2655
|
+
const m = rel[(windowSize + off) * maxN + nb];
|
|
2656
|
+
if (m >= 0) mapped.add(m);
|
|
2657
|
+
}
|
|
2658
|
+
const other = /* @__PURE__ */ new Set();
|
|
2659
|
+
for (let p = graphs[t].indptr[counterpart]; p < graphs[t].indptr[counterpart + 1]; p++) {
|
|
2660
|
+
const nb = graphs[t].indices[p];
|
|
2661
|
+
if (image.has(nb)) other.add(nb);
|
|
2662
|
+
}
|
|
2663
|
+
let inter = 0;
|
|
2664
|
+
for (const m of mapped) if (other.has(m)) inter++;
|
|
2665
|
+
const union = mapped.size + other.size - inter;
|
|
2666
|
+
weights[(windowSize + off) * maxN + v] = union === 0 ? 1 : inter / union;
|
|
2667
|
+
}
|
|
2668
|
+
}
|
|
2669
|
+
out.push(weights);
|
|
2670
|
+
}
|
|
2671
|
+
return out;
|
|
2672
|
+
}
|
|
2673
|
+
/** procrustes_align (aligned_umap.py:33-40): rotation-only via SVD of subset2ᵀ·subset1. */
|
|
2674
|
+
function procrustesRotate(embedding, dim, anchorsSelf, anchorsPrev, prev) {
|
|
2675
|
+
const d = dim;
|
|
2676
|
+
const M = new Float64Array(d * d);
|
|
2677
|
+
for (let t = 0; t < anchorsSelf.length; t++) {
|
|
2678
|
+
const a = anchorsSelf[t];
|
|
2679
|
+
const b = anchorsPrev[t];
|
|
2680
|
+
for (let r = 0; r < d; r++) for (let c = 0; c < d; c++) M[r * d + c] = M[r * d + c] + embedding[a * d + r] * prev[b * d + c];
|
|
2681
|
+
}
|
|
2682
|
+
const MtM = new Float64Array(d * d);
|
|
2683
|
+
for (let r = 0; r < d; r++) for (let c = 0; c < d; c++) {
|
|
2684
|
+
let s = 0;
|
|
2685
|
+
for (let k2 = 0; k2 < d; k2++) s += M[k2 * d + r] * M[k2 * d + c];
|
|
2686
|
+
MtM[r * d + c] = s;
|
|
2687
|
+
}
|
|
2688
|
+
const eig = symmetricEigen(MtM, d);
|
|
2689
|
+
const U = new Float64Array(d * d);
|
|
2690
|
+
const V = new Float64Array(d * d);
|
|
2691
|
+
for (let c = 0; c < d; c++) {
|
|
2692
|
+
const col = d - 1 - c;
|
|
2693
|
+
const sigma = Math.sqrt(Math.max(eig.values[col], 1e-30));
|
|
2694
|
+
for (let r = 0; r < d; r++) V[r * d + c] = eig.vectors[col * d + r];
|
|
2695
|
+
for (let r = 0; r < d; r++) {
|
|
2696
|
+
let s = 0;
|
|
2697
|
+
for (let k2 = 0; k2 < d; k2++) s += M[r * d + k2] * V[k2 * d + c];
|
|
2698
|
+
U[r * d + c] = s / sigma;
|
|
2699
|
+
}
|
|
2700
|
+
}
|
|
2701
|
+
const R = new Float64Array(d * d);
|
|
2702
|
+
for (let r = 0; r < d; r++) for (let c = 0; c < d; c++) {
|
|
2703
|
+
let s = 0;
|
|
2704
|
+
for (let k2 = 0; k2 < d; k2++) s += U[r * d + k2] * V[c * d + k2];
|
|
2705
|
+
R[r * d + c] = s;
|
|
2706
|
+
}
|
|
2707
|
+
const n = embedding.length / d;
|
|
2708
|
+
const tmp = new Float64Array(d);
|
|
2709
|
+
for (let i = 0; i < n; i++) {
|
|
2710
|
+
for (let c = 0; c < d; c++) {
|
|
2711
|
+
let s = 0;
|
|
2712
|
+
for (let k2 = 0; k2 < d; k2++) s += embedding[i * d + k2] * R[k2 * d + c];
|
|
2713
|
+
tmp[c] = s;
|
|
2714
|
+
}
|
|
2715
|
+
for (let c = 0; c < d; c++) embedding[i * d + c] = tmp[c];
|
|
2716
|
+
}
|
|
2717
|
+
}
|
|
2718
|
+
var AlignedUMAP = class {
|
|
2719
|
+
options;
|
|
2720
|
+
_embeddings = null;
|
|
2721
|
+
_mappers = [];
|
|
2722
|
+
_relations = [];
|
|
2723
|
+
constructor(options = {}) {
|
|
2724
|
+
this.options = { ...options };
|
|
2725
|
+
}
|
|
2726
|
+
get embeddings() {
|
|
2727
|
+
if (!this._embeddings) throw new UmapError("model is not fitted");
|
|
2728
|
+
return this._embeddings;
|
|
2729
|
+
}
|
|
2730
|
+
async fit(XList, yList, relations, { signal } = {}) {
|
|
2731
|
+
throwIfAborted(signal);
|
|
2732
|
+
if (!relations || relations.length !== XList.length - 1) throw new UmapValidationError("relations must have length len(X) - 1");
|
|
2733
|
+
this._relations = relations;
|
|
2734
|
+
const o = this.options;
|
|
2735
|
+
const windowSize = o.alignmentWindowSize ?? 3;
|
|
2736
|
+
const lambda = o.alignmentRegularisation ?? .01;
|
|
2737
|
+
const nComponents = o.nComponents ?? 2;
|
|
2738
|
+
const nEpochs = o.nEpochs ?? 200;
|
|
2739
|
+
const slices = XList.map((X, i) => toMatrix(X, `X[${i}]`));
|
|
2740
|
+
const sliceSizes = slices.map((s) => s.rows);
|
|
2741
|
+
const maxN = Math.max(...sliceSizes);
|
|
2742
|
+
this._mappers = [];
|
|
2743
|
+
for (let i = 0; i < slices.length; i++) {
|
|
2744
|
+
const { alignmentRegularisation: _ar, alignmentWindowSize: _aw, nEpochs: _ne,...umapOpts } = o;
|
|
2745
|
+
const mapper = new UMAP({
|
|
2746
|
+
...umapOpts,
|
|
2747
|
+
nComponents
|
|
2748
|
+
});
|
|
2749
|
+
await mapper.fit(slices[i], yList?.[i] ?? null, signal ? { signal } : {});
|
|
2750
|
+
this._mappers.push(mapper);
|
|
2751
|
+
}
|
|
2752
|
+
const graphs = this._mappers.map((m) => m.graph);
|
|
2753
|
+
throwIfAborted(signal);
|
|
2754
|
+
const expanded = expandRelations(relations, windowSize, sliceSizes);
|
|
2755
|
+
const regWeights = buildNeighborhoodSimilarities(graphs, expanded, windowSize, maxN);
|
|
2756
|
+
const heads = [];
|
|
2757
|
+
const tails = [];
|
|
2758
|
+
const epochsPerSamples = [];
|
|
2759
|
+
for (const g of graphs) {
|
|
2760
|
+
const nnz = g.indptr[g.rows];
|
|
2761
|
+
const h = new Int32Array(nnz);
|
|
2762
|
+
const t = new Int32Array(nnz);
|
|
2763
|
+
const w = new Float32Array(nnz);
|
|
2764
|
+
for (let r = 0; r < g.rows; r++) for (let p = g.indptr[r]; p < g.indptr[r + 1]; p++) {
|
|
2765
|
+
h[p] = r;
|
|
2766
|
+
t[p] = g.indices[p];
|
|
2767
|
+
w[p] = Math.fround(g.data[p]);
|
|
2768
|
+
}
|
|
2769
|
+
heads.push(h);
|
|
2770
|
+
tails.push(t);
|
|
2771
|
+
epochsPerSamples.push(makeEpochsPerSample(w, nEpochs));
|
|
2772
|
+
}
|
|
2773
|
+
const seed = o.seed ?? 42;
|
|
2774
|
+
const rng = Pcg32.fromSeed(seed);
|
|
2775
|
+
const embeddings = [];
|
|
2776
|
+
for (let i = 0; i < slices.length; i++) {
|
|
2777
|
+
const res = spectralLayout(slices[i], graphs[i], nComponents, rng);
|
|
2778
|
+
let mx = 0;
|
|
2779
|
+
for (const v of res.embedding) mx = Math.max(mx, Math.abs(v));
|
|
2780
|
+
const expansion = mx > 0 ? 10 / mx : 1;
|
|
2781
|
+
const emb = new Float32Array(res.embedding.length);
|
|
2782
|
+
for (let t = 0; t < emb.length; t++) emb[t] = Math.fround(res.embedding[t] * expansion);
|
|
2783
|
+
if (i > 0) {
|
|
2784
|
+
const anchorsPrev = [];
|
|
2785
|
+
const anchorsSelf = [];
|
|
2786
|
+
const rel = relations[i - 1];
|
|
2787
|
+
const entries = rel instanceof Map ? rel.entries() : Object.entries(rel);
|
|
2788
|
+
for (const [kRaw, v] of entries) {
|
|
2789
|
+
anchorsPrev.push(Number(kRaw));
|
|
2790
|
+
anchorsSelf.push(v);
|
|
2791
|
+
}
|
|
2792
|
+
procrustesRotate(emb, nComponents, Int32Array.from(anchorsSelf), Int32Array.from(anchorsPrev), embeddings[i - 1]);
|
|
2793
|
+
}
|
|
2794
|
+
embeddings.push(emb);
|
|
2795
|
+
}
|
|
2796
|
+
optimizeLayoutAlignedEuclidean(embeddings, nComponents, heads, tails, nEpochs, epochsPerSamples, regWeights, expanded, windowSize, maxN, Pcg32.fromSeed(o.transformSeed ?? 42).tauState(), {
|
|
2797
|
+
lambda,
|
|
2798
|
+
a: 1.576943460405378,
|
|
2799
|
+
b: .8950608781227859,
|
|
2800
|
+
gamma: o.repulsionStrength ?? 1,
|
|
2801
|
+
initialAlpha: o.learningRate ?? 1,
|
|
2802
|
+
negativeSampleRate: o.negativeSampleRate ?? 5,
|
|
2803
|
+
moveOther: true,
|
|
2804
|
+
...signal ? { signal } : {}
|
|
2805
|
+
});
|
|
2806
|
+
this._embeddings = embeddings.map((e, i) => ({
|
|
2807
|
+
data: e,
|
|
2808
|
+
rows: sliceSizes[i],
|
|
2809
|
+
cols: nComponents
|
|
2810
|
+
}));
|
|
2811
|
+
return this;
|
|
2812
|
+
}
|
|
2813
|
+
async fitTransform(XList, yList, relations, opts = {}) {
|
|
2814
|
+
await this.fit(XList, yList, relations, opts);
|
|
2815
|
+
return this.embeddings;
|
|
2816
|
+
}
|
|
2817
|
+
/** Append a new slice (aligned_umap.py:446-562): old slices frozen except pulls. */
|
|
2818
|
+
async update(Xnew, relation, { signal } = {}) {
|
|
2819
|
+
throwIfAborted(signal);
|
|
2820
|
+
if (!this._embeddings) throw new UmapError("model is not fitted");
|
|
2821
|
+
const o = this.options;
|
|
2822
|
+
const windowSize = o.alignmentWindowSize ?? 3;
|
|
2823
|
+
const lambda = o.alignmentRegularisation ?? .01;
|
|
2824
|
+
const nComponents = o.nComponents ?? 2;
|
|
2825
|
+
const nEpochs = o.nEpochs ?? 200;
|
|
2826
|
+
const X = toMatrix(Xnew, "Xnew");
|
|
2827
|
+
const { alignmentRegularisation: _ar, alignmentWindowSize: _aw, nEpochs: _ne,...umapOpts } = o;
|
|
2828
|
+
const mapper = new UMAP({
|
|
2829
|
+
...umapOpts,
|
|
2830
|
+
nComponents
|
|
2831
|
+
});
|
|
2832
|
+
await mapper.fit(X, null, signal ? { signal } : {});
|
|
2833
|
+
this._mappers.push(mapper);
|
|
2834
|
+
this._relations.push(relation);
|
|
2835
|
+
const graphs = this._mappers.map((m) => m.graph);
|
|
2836
|
+
const sliceSizes = graphs.map((g) => g.rows);
|
|
2837
|
+
const maxN = Math.max(...sliceSizes);
|
|
2838
|
+
const expanded = expandRelations(this._relations, windowSize, sliceSizes);
|
|
2839
|
+
const regWeights = buildNeighborhoodSimilarities(graphs, expanded, windowSize, maxN);
|
|
2840
|
+
const heads = [];
|
|
2841
|
+
const tails = [];
|
|
2842
|
+
const epochsPerSamples = [];
|
|
2843
|
+
for (let i = 0; i < graphs.length; i++) {
|
|
2844
|
+
const g = graphs[i];
|
|
2845
|
+
if (i < graphs.length - 1) {
|
|
2846
|
+
heads.push(new Int32Array(0));
|
|
2847
|
+
tails.push(new Int32Array(0));
|
|
2848
|
+
epochsPerSamples.push(new Float64Array(0));
|
|
2849
|
+
continue;
|
|
2850
|
+
}
|
|
2851
|
+
const nnz = g.indptr[g.rows];
|
|
2852
|
+
const h = new Int32Array(nnz);
|
|
2853
|
+
const t = new Int32Array(nnz);
|
|
2854
|
+
const w = new Float32Array(nnz);
|
|
2855
|
+
for (let r = 0; r < g.rows; r++) for (let p = g.indptr[r]; p < g.indptr[r + 1]; p++) {
|
|
2856
|
+
h[p] = r;
|
|
2857
|
+
t[p] = g.indices[p];
|
|
2858
|
+
w[p] = Math.fround(g.data[p]);
|
|
2859
|
+
}
|
|
2860
|
+
heads.push(h);
|
|
2861
|
+
tails.push(t);
|
|
2862
|
+
epochsPerSamples.push(makeEpochsPerSample(w, nEpochs));
|
|
2863
|
+
}
|
|
2864
|
+
const prev = this._embeddings[this._embeddings.length - 1];
|
|
2865
|
+
const rng = Pcg32.fromSeed(o.transformSeed ?? 42);
|
|
2866
|
+
const emb = new Float32Array(X.rows * nComponents);
|
|
2867
|
+
for (let v = 0; v < X.rows; v++) {
|
|
2868
|
+
let src = -1;
|
|
2869
|
+
const entries = relation instanceof Map ? relation.entries() : Object.entries(relation);
|
|
2870
|
+
for (const [kRaw, val] of entries) if (val === v) {
|
|
2871
|
+
src = Number(kRaw);
|
|
2872
|
+
break;
|
|
2873
|
+
}
|
|
2874
|
+
if (src >= 0) for (let c = 0; c < nComponents; c++) emb[v * nComponents + c] = prev.data[src * nComponents + c];
|
|
2875
|
+
else for (let c = 0; c < nComponents; c++) emb[v * nComponents + c] = rng.uniform(-10, 10);
|
|
2876
|
+
}
|
|
2877
|
+
const embeddings = [...this._embeddings.map((e) => e.data), emb];
|
|
2878
|
+
optimizeLayoutAlignedEuclidean(embeddings, nComponents, heads, tails, nEpochs, epochsPerSamples, regWeights, expanded, windowSize, maxN, rng.tauState(), {
|
|
2879
|
+
lambda,
|
|
2880
|
+
a: 1.576943460405378,
|
|
2881
|
+
b: .8950608781227859,
|
|
2882
|
+
gamma: o.repulsionStrength ?? 1,
|
|
2883
|
+
initialAlpha: o.learningRate ?? 1,
|
|
2884
|
+
negativeSampleRate: o.negativeSampleRate ?? 5,
|
|
2885
|
+
moveOther: false,
|
|
2886
|
+
...signal ? { signal } : {}
|
|
2887
|
+
});
|
|
2888
|
+
this._embeddings = embeddings.map((e, i) => ({
|
|
2889
|
+
data: e,
|
|
2890
|
+
rows: sliceSizes[i],
|
|
2891
|
+
cols: nComponents
|
|
2892
|
+
}));
|
|
2893
|
+
return this;
|
|
2894
|
+
}
|
|
2895
|
+
};
|
|
2896
|
+
/** optimize_layout_aligned_euclidean (layouts.py:877-1104; notes §19.4). */
|
|
2897
|
+
function optimizeLayoutAlignedEuclidean(embeddings, dim, heads, tails, nEpochs, epochsPerSamples, regWeights, relations, windowSize, maxN, rngState, P) {
|
|
2898
|
+
const nEmbeddings = embeddings.length;
|
|
2899
|
+
const { a, b, gamma, lambda, negativeSampleRate, moveOther } = P;
|
|
2900
|
+
let alpha = P.initialAlpha;
|
|
2901
|
+
const epochsPerNegative = epochsPerSamples.map((eps) => {
|
|
2902
|
+
const out = new Float64Array(eps.length);
|
|
2903
|
+
for (let i = 0; i < eps.length; i++) out[i] = eps[i] / negativeSampleRate;
|
|
2904
|
+
return out;
|
|
2905
|
+
});
|
|
2906
|
+
const epochOfNext = epochsPerSamples.map((eps) => eps.slice());
|
|
2907
|
+
const epochOfNextNeg = epochsPerNegative.map((e) => e.slice());
|
|
2908
|
+
const maxNEdges = Math.max(...epochsPerSamples.map((e) => e.length));
|
|
2909
|
+
const order = new Int32Array(nEmbeddings);
|
|
2910
|
+
const applyRegularisation = (m, vertex, d, coord) => {
|
|
2911
|
+
let g = 0;
|
|
2912
|
+
for (let offset = -windowSize; offset < windowSize; offset++) {
|
|
2913
|
+
const nm = m + offset;
|
|
2914
|
+
if (offset === 0 || nm < 0 || nm >= nEmbeddings) continue;
|
|
2915
|
+
const identified = relations[m][(windowSize + offset) * maxN + vertex];
|
|
2916
|
+
if (identified < 0) continue;
|
|
2917
|
+
g -= clip(lambda * Math.exp(-(Math.abs(offset) - 1)) * regWeights[m][(windowSize + offset) * maxN + vertex] * (coord - embeddings[nm][identified * dim + d]));
|
|
2918
|
+
}
|
|
2919
|
+
return g;
|
|
2920
|
+
};
|
|
2921
|
+
for (let n = 0; n < nEpochs; n++) {
|
|
2922
|
+
throwIfAborted(P.signal);
|
|
2923
|
+
for (let i2 = 0; i2 < nEmbeddings; i2++) order[i2] = i2;
|
|
2924
|
+
Pcg32.fromSeed(Math.abs(rngState[0])).shuffle(order);
|
|
2925
|
+
for (let i = 0; i < maxNEdges; i++) for (const m of order) {
|
|
2926
|
+
const eons = epochOfNext[m];
|
|
2927
|
+
if (i >= eons.length || eons[i] > n) continue;
|
|
2928
|
+
const emb = embeddings[m];
|
|
2929
|
+
const j = heads[m][i];
|
|
2930
|
+
let k = tails[m][i];
|
|
2931
|
+
const jb = j * dim;
|
|
2932
|
+
let kb = k * dim;
|
|
2933
|
+
let d2 = 0;
|
|
2934
|
+
for (let d = 0; d < dim; d++) {
|
|
2935
|
+
const diff = emb[jb + d] - emb[kb + d];
|
|
2936
|
+
d2 += diff * diff;
|
|
2937
|
+
}
|
|
2938
|
+
let gradCoeff = 0;
|
|
2939
|
+
if (d2 > 0) gradCoeff = -2 * a * b * d2 ** (b - 1) / (a * d2 ** b + 1);
|
|
2940
|
+
for (let d = 0; d < dim; d++) {
|
|
2941
|
+
let gradD = clip(gradCoeff * (emb[jb + d] - emb[kb + d]));
|
|
2942
|
+
gradD += applyRegularisation(m, j, d, emb[jb + d]);
|
|
2943
|
+
emb[jb + d] = emb[jb + d] + clip(gradD) * alpha;
|
|
2944
|
+
if (moveOther) {
|
|
2945
|
+
let otherGrad = clip(gradCoeff * (emb[kb + d] - emb[jb + d]));
|
|
2946
|
+
otherGrad += applyRegularisation(m, k, d, emb[kb + d]);
|
|
2947
|
+
emb[kb + d] = emb[kb + d] + clip(otherGrad) * alpha;
|
|
2948
|
+
}
|
|
2949
|
+
}
|
|
2950
|
+
eons[i] = eons[i] + epochsPerSamples[m][i];
|
|
2951
|
+
const epn = epochsPerNegative[m][i];
|
|
2952
|
+
let nNeg = 0;
|
|
2953
|
+
if (epn > 0) nNeg = Math.trunc((n - epochOfNextNeg[m][i]) / epn);
|
|
2954
|
+
const nVertices = emb.length / dim;
|
|
2955
|
+
for (let p = 0; p < nNeg; p++) {
|
|
2956
|
+
const r = tauRandInt(rngState);
|
|
2957
|
+
k = r % nVertices + (r % nVertices < 0 ? nVertices : 0);
|
|
2958
|
+
kb = k * dim;
|
|
2959
|
+
let dn = 0;
|
|
2960
|
+
for (let d = 0; d < dim; d++) {
|
|
2961
|
+
const diff = emb[jb + d] - emb[kb + d];
|
|
2962
|
+
dn += diff * diff;
|
|
2963
|
+
}
|
|
2964
|
+
let repCoeff = 0;
|
|
2965
|
+
if (dn > 0) repCoeff = 2 * gamma * b / ((.001 + dn) * (a * dn ** b + 1));
|
|
2966
|
+
else if (j === k) continue;
|
|
2967
|
+
for (let d = 0; d < dim; d++) {
|
|
2968
|
+
let gradD = 0;
|
|
2969
|
+
if (repCoeff > 0) gradD = clip(repCoeff * (emb[jb + d] - emb[kb + d]));
|
|
2970
|
+
gradD += applyRegularisation(m, j, d, emb[jb + d]);
|
|
2971
|
+
emb[jb + d] = emb[jb + d] + clip(gradD) * alpha;
|
|
2972
|
+
}
|
|
2973
|
+
}
|
|
2974
|
+
epochOfNextNeg[m][i] = epochOfNextNeg[m][i] + nNeg * epn;
|
|
2975
|
+
}
|
|
2976
|
+
alpha = P.initialAlpha * (1 - n / nEpochs);
|
|
2977
|
+
}
|
|
2978
|
+
}
|
|
2979
|
+
|
|
2980
|
+
//#endregion
|
|
2981
|
+
//#region src/index.ts
|
|
2982
|
+
/**
|
|
2983
|
+
* umap-web — standalone TypeScript UMAP for browsers and Node.js.
|
|
2984
|
+
*
|
|
2985
|
+
* Public entry point. See `umap-web/core` for the low-level functional API.
|
|
2986
|
+
*/
|
|
2987
|
+
const VERSION = "0.1.0";
|
|
2988
|
+
|
|
2989
|
+
//#endregion
|
|
2990
|
+
export { ANGULAR_METRICS, AlignedUMAP, DISCONNECTION_DISTANCES, MIN_K_DIST_SCALE, NAMED_METRICS, Pcg32, SMALL_DATA_CUTOFF, SMOOTH_K_TOLERANCE, UMAP, UmapBackendError, UmapConvergenceWarning, UmapError, UmapValidationError, VERSION, canonicalMetricName, computeMembershipStrengths, connectedComponents, csrBinaryOp, csrCopy, csrEliminateZeros, csrForEach, csrFromCoo, csrMapValues, csrMaxAbsDiff, csrNnz, csrRowNormalizeMax, csrTranspose, defaultNEpochs, degreeInvSqrt, exactKnnDense, exactKnnSparse, finalInitRescale, findABParams, floorMod, fuzzySimplicialSet, initializeEmbedding, isNamedMetric, knnFromPrecomputed, lobpcgSmallest, makeEpochsPerSample, metrics_exports as metrics, mgsOrthonormalize, nearestNeighborsDense, nearestNeighborsSparse, nnDescent, noisyScaleCoords, pcaScores, pruneGraphForEpochs, resolveMetric, searchKnn, smoothKnnDist, spectralLayout, symmetricEigen, tauRand, tauRandInt, toMatrix, umapNnDescentParams };
|
|
2991
|
+
//# sourceMappingURL=index.js.map
|