Skip to content

Commit 0bf5bd4

Browse files
committed
refactor(apps/deditor): replace umap with wasm implementation
1 parent 5623768 commit 0bf5bd4

26 files changed

Lines changed: 1416 additions & 32 deletions

NOTICE.md

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# NOTICES
2+
3+
This repository incorporates material as listed below or described in the code.
4+
5+
## umap-wasm
6+
7+
**Source**: https://github.com/apple/embedding-atlas/tree/main/packages/umap-wasm
8+
9+
**License**:
10+
11+
```
12+
The MIT License (MIT)
13+
14+
Copyright (c) 2025 Apple Inc.
15+
16+
Permission is hereby granted, free of charge, to any person obtaining a copy
17+
of this software and associated documentation files (the "Software"), to deal
18+
in the Software without restriction, including without limitation the rights
19+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
20+
copies of the Software, and to permit persons to whom the Software is
21+
furnished to do so, subject to the following conditions:
22+
23+
The above copyright notice and this permission notice shall be included in
24+
all copies or substantial portions of the Software.
25+
26+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
27+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
28+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
29+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
30+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
31+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
32+
THE SOFTWARE.
33+
```

apps/deditor/package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
"devDependencies": {
6565
"@deditor-app/drizzle-orm-icons": "workspace:^",
6666
"@deditor-app/shared": "workspace:^",
67+
"@deditor-app/umap-wasm": "workspace:^",
6768
"@iconify-json/ph": "^1.2.2",
6869
"@iconify-json/simple-icons": "^1.2.40",
6970
"@iconify-json/svg-spinners": "^1.2.2",
@@ -89,7 +90,6 @@
8990
"tailwind-merge": "^3.3.1",
9091
"three": "^0.177.0",
9192
"tw-animate-css": "^1.3.4",
92-
"umap-js": "^1.4.0",
9393
"unplugin-vue-macros": "^2.14.5",
9494
"unplugin-vue-markdown": "^28.3.1",
9595
"unplugin-vue-router": "^0.12.0",

apps/deditor/src/renderer/src/stores/visualizer.ts

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,14 @@ import type * as THREE from 'three'
22

33
import type { DataPointStyle, ProjectionParameters } from '@/types/visualizer'
44

5+
import { createUMAP } from '@deditor-app/umap-wasm'
56
import { PCA } from 'ml-pca'
67
import { TSNE } from 'msvana-tsne'
78
import { defineStore } from 'pinia'
8-
import { UMAP } from 'umap-js'
99
import { readonly, ref } from 'vue'
1010

1111
import { ProjectionAlgorithm } from '@/constants'
12+
import { selectNEpochs } from '@/utils'
1213
import { toVec3s } from '@/utils/three'
1314

1415
export const useVisualizerStore = defineStore('visualizer', () => {
@@ -68,18 +69,34 @@ export const useVisualizerStore = defineStore('visualizer', () => {
6869
vectors.value = source.map(v => [...v])
6970
}
7071

71-
const visualize = () => {
72+
const visualize = async () => {
7273
switch (projection.value?.type) {
7374
case ProjectionAlgorithm.UMAP: {
7475
const params = projection.value.params
75-
const umap = new UMAP({
76-
nComponents: params.dimensions,
77-
nNeighbors: params.neighbors,
78-
minDist: params.minDistance,
79-
})
76+
77+
const umap = await createUMAP(
78+
vectors.value.length, // Count
79+
vectors.value[0].length, // Input dimensions
80+
params.dimensions, // Output dimensions
81+
Float32Array.from(vectors.value.flat()), // Data
82+
{
83+
n_neighbors: params.neighbors,
84+
min_dist: params.minDistance,
85+
// TODO: Add more UI controls for the rest of the parameters
86+
},
87+
)
88+
89+
// TODO: Make use of this epoch to animate from the common PCA to a UMAP-like projection
90+
umap.run(selectNEpochs(vectors.value.length))
91+
92+
const embedding = umap.embedding
93+
const outVectors: number[][] = []
94+
for (let i = 0; i < embedding.length / params.dimensions; i++) {
95+
outVectors.push([...embedding.subarray(i * params.dimensions, (i + 1) * params.dimensions)])
96+
}
8097

8198
try {
82-
points.value = toVec3s(umap.fit(vectors.value as number[][]))
99+
points.value = toVec3s(outVectors)
83100
}
84101
catch (error) {
85102
// Sometimes this can fail if the params do not fit the data
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
export * from './three'
2+
export * from './umap'
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/**
2+
* Gets the number of epochs for optimizing the projection.
3+
*
4+
* https://github.com/PAIR-code/umap-js/blob/7767b1d0dd4b47718b5ad32d0f65dfb3b955d428/src/umap.ts#L1064-L1086
5+
*/
6+
export function selectNEpochs(count: number) {
7+
if (count <= 2500) {
8+
return 500
9+
}
10+
else if (count <= 5000) {
11+
return 400
12+
}
13+
else if (count <= 7500) {
14+
return 300
15+
}
16+
else {
17+
return 200
18+
}
19+
}

packages/umap-wasm/.clang-format

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
---
2+
# We'll use defaults from the LLVM style, but with 4 columns indentation.
3+
BasedOnStyle: LLVM
4+
IndentWidth: 4

packages/umap-wasm/.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
__pycache__
2+
node_modules/
3+
.ipynb_checkpoints/

packages/umap-wasm/LICENSE

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
The MIT License (MIT)
2+
3+
Copyright (c) 2025 Apple Inc.
4+
5+
Permission is hereby granted, free of charge, to any person obtaining a copy
6+
of this software and associated documentation files (the "Software"), to deal
7+
in the Software without restriction, including without limitation the rights
8+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9+
copies of the Software, and to permit persons to whom the Software is
10+
furnished to do so, subject to the following conditions:
11+
12+
The above copyright notice and this permission notice shall be included in
13+
all copies or substantial portions of the Software.
14+
15+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
21+
THE SOFTWARE.

packages/umap-wasm/Makefile

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
SOURCES = src/umap.cpp src/knn.cpp src/memory.cpp
2+
HEADERS = src/common.hpp src/knn.hpp src/umap.hpp src/knncolle_nndescent.hpp src/distances.hpp
3+
4+
CXXFLAGS = -std=c++20 -O3 -Wall -Wno-deprecated-declarations
5+
6+
runtime.js: $(SOURCES) $(HEADERS)
7+
emcc $(CXXFLAGS) $(SOURCES) \
8+
third_party/nndescent/src/distances.cpp \
9+
third_party/nndescent/src/dtypes.cpp \
10+
third_party/nndescent/src/nnd.cpp \
11+
third_party/nndescent/src/rp_trees.cpp \
12+
third_party/nndescent/src/utils.cpp \
13+
-I third_party/umappp/include \
14+
-I third_party/knncolle/include \
15+
-I third_party/knncolle_hnsw/include \
16+
-I third_party/CppIrlba/include \
17+
-I third_party/CppKmeans/include \
18+
-I third_party/aarand/include \
19+
-I third_party/subpar/include \
20+
-I third_party/nndescent/src \
21+
-I third_party/Eigen \
22+
-I third_party/hnswlib \
23+
-DUMAPPP_NO_PARALLEL_OPTIMIZATION \
24+
-sENVIRONMENT=web \
25+
-sMODULARIZE -sSINGLE_FILE -sALLOW_MEMORY_GROWTH \
26+
-sWASM_BIGINT=1 \
27+
-msimd128 \
28+
-o runtime.mjs
29+
mv runtime.mjs runtime.js
30+
31+
clean:
32+
rm -rf runtime.js

packages/umap-wasm/README.md

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# umap-wasm
2+
3+
This package provides a WebAssembly implementation of UMAP and approximate nearest neighbor search.
4+
5+
The UMAP implementation is [umappp](https://github.com/libscran/umappp/) by Aaron Lun.
6+
7+
The approximate nearest neighbor search is based on [knncolle](https://github.com/knncolle/knncolle), with two algorithms, including HNSW ([hnswlib](https://github.com/nmslib/hnswlib)) and [nndescent](https://github.com/brj0/nndescent) by Jon Brugger.
8+
9+
## Documentation
10+
11+
See `index.d.ts` for the full API.
12+
13+
To initialize the algorithm, use `createUMAP`:
14+
15+
```js
16+
import { createUMAP } from "umap-wasm";
17+
18+
let count = 2000;
19+
let input_dim = 100;
20+
let output_dim = 2;
21+
22+
// The data must be a Float32Array with count * input_dim elements.
23+
let data = new Float32Array(count * input_dim);
24+
// ... fill in the data
25+
26+
let options = {
27+
metric: "cosine",
28+
};
29+
30+
// Use `createUMAP` to initialize the algorithm.
31+
let umap = await createUMAP(count, input_dim, output_dim, data, options);
32+
```
33+
34+
After initialization, use the `run` method to update the embedding coordinates:
35+
36+
```js
37+
// Run the algorithm to completion.
38+
umap.run();
39+
40+
// Alternatively, you can run up to a given number of epochs.
41+
// This can be useful for animation effects.
42+
for (let i = 0; i < 100; i++) {
43+
// Run to the i-th epoch.
44+
umap.run(i);
45+
}
46+
```
47+
48+
At any time, you can get the current embedding by calling the `embedding` method.
49+
50+
```js
51+
// The result is a Float32Array with count * output_dim elements.
52+
let embedding = umap.embedding();
53+
```
54+
55+
After you are done with the instance, use the `destroy` method to release resources.
56+
57+
```js
58+
umap.destroy();
59+
```
60+
61+
## Development
62+
63+
A pre-compiled WASM file is included in the Git repository. If you'd like to build this package, proceed with the following instructions:
64+
65+
First, install [Emscripten](https://emscripten.org/).
66+
The WASM module is built with the Emscripten toolchain.
67+
68+
Download all necessary dependencies with the following commands:
69+
70+
```bash
71+
cd third_party
72+
./download_dependencies.sh
73+
```
74+
75+
Note that you don't need to build any of these packages.
76+
77+
Finally, you may run:
78+
79+
```bash
80+
make
81+
```
82+
83+
to build the WASM module. The output file should be `runtime.js`.

0 commit comments

Comments
 (0)