@elarsaks/umap-wasm 0.4.4 → 0.4.6

package/dist/matrix.d.ts

@@ -2,7 +2,7 @@ export declare class SparseMatrix {
     private entries;
     readonly nRows: number;
     readonly nCols: number;
-    constructor(rows: number[], cols: number[], values: number[], dims: number[]);
+    constructor(rows: ArrayLike<number>, cols: ArrayLike<number>, values: ArrayLike<number>, dims: number[]);
     private makeKey;
     private checkDims;
     set(row: number, col: number, value: number): void;

package/dist/nn_descent.d.ts

@@ -2,7 +2,7 @@ import * as heap from './heap.js';
 import * as matrix from './matrix.js';
 import * as tree from './tree.js';
 import { RandomFn, Vectors, DistanceFn } from './umap.js';
-export declare function makeNNDescent(distanceFn: DistanceFn, random: RandomFn, useWasm?: boolean): (data: Vectors, leafArray: Vectors, nNeighbors: number, nIters?: number, maxCandidates?: number, delta?: number, rho?: number, rpTreeInit?: boolean) => {
+export declare function makeNNDescent(distanceFn: DistanceFn, random: RandomFn, useWasm?: boolean): (data: Vectors, leafArray: ArrayLike<number>[], nNeighbors: number, nIters?: number, maxCandidates?: number, delta?: number, rho?: number, rpTreeInit?: boolean) => {
     indices: number[][];
     weights: number[][];
 };

package/dist/nn_descent.js

@@ -11,10 +11,39 @@ export function makeNNDescent(distanceFn, random, useWasm = false) {
             }
             const distanceMetric = distanceFn.name === 'cosine' ? 'cosine' : 'euclidean';
             const seed = Math.floor(random() * 0xFFFFFFFF);
-            const result = wasmBridge.nnDescentWasm(data, leafArray, nNeighbors, nIters, maxCandidates, delta, rho, rpTreeInit, distanceMetric, seed);
+            const nSamples = data.length;
+            const dim = data[0].length;
+            const flatData = new Float64Array(nSamples * dim);
+            for (let i = 0; i < nSamples; i++) {
+                for (let j = 0; j < dim; j++) {
+                    flatData[i * dim + j] = data[i][j];
+                }
+            }
+            const nLeaves = leafArray.length;
+            const leafSize = nLeaves > 0 ? leafArray[0].length : 0;
+            const flatLeafArray = new Int32Array(nLeaves * leafSize);
+            for (let i = 0; i < nLeaves; i++) {
+                for (let j = 0; j < leafSize; j++) {
+                    flatLeafArray[i * leafSize + j] = leafArray[i][j];
+                }
+            }
+            const result = wasmBridge.nnDescentWasmFlat(flatData, nSamples, dim, flatLeafArray, nLeaves, leafSize, nNeighbors, nIters, maxCandidates, delta, rho, rpTreeInit, distanceMetric, seed);
+            const indices = [];
+            const distances = [];
+            const offset1 = nSamples * nNeighbors;
+            for (let i = 0; i < nSamples; i++) {
+                const rowIndices = [];
+                const rowDistances = [];
+                for (let j = 0; j < nNeighbors; j++) {
+                    rowIndices.push(result[i * nNeighbors + j]);
+                    rowDistances.push(result[offset1 + i * nNeighbors + j]);
+                }
+                indices.push(rowIndices);
+                distances.push(rowDistances);
+            }
             return {
-                indices: result[0],
-                weights: result[1],
+                indices,
+                weights: distances,
             };
         }
         const nVertices = data.length;

package/dist/tree.d.ts

@@ -1,16 +1,38 @@
 import { RandomFn, Vector, Vectors } from './umap.js';
 import { WasmFlatTree } from './wasmBridge.js';
 export declare class FlatTree {
-    hyperplanes: number[][];
-    offsets: number[];
-    children: number[][];
-    indices: number[][];
+    hyperplanes: ArrayLike<number>[];
+    offsets: ArrayLike<number>;
+    children: ArrayLike<number>[];
+    indices: ArrayLike<number>[];
     private wasmTree?;
-    constructor(hyperplanes: number[][], offsets: number[], children: number[][], indices: number[][]);
+    private hyperplanesFlat?;
+    private offsetsFlat?;
+    private childrenFlat?;
+    private indicesFlat?;
+    private dim?;
+    private nNodes?;
+    private nLeaves?;
+    private leafSize?;
+    constructor(hyperplanes: ArrayLike<number>[], offsets: ArrayLike<number>, children: ArrayLike<number>[], indices: ArrayLike<number>[]);
     static fromWasm(wasmTree: WasmFlatTree): FlatTree;
     getWasmTree(): WasmFlatTree | undefined;
+    getFlatHyperplanes(): Float64Array | undefined;
+    getFlatOffsets(): Float64Array | undefined;
+    getFlatChildren(): Int32Array | undefined;
+    getFlatLeafMeta(): {
+        indices: Int32Array;
+        nLeaves: number;
+        leafSize: number;
+    } | undefined;
+    getDim(): number | undefined;
     dispose(): void;
 }
 export declare function makeForest(data: Vectors, nNeighbors: number, nTrees: number, random: RandomFn, useWasm?: boolean): FlatTree[];
-export declare function makeLeafArray(rpForest: FlatTree[]): number[][];
-export declare function searchFlatTree(point: Vector, tree: FlatTree, random: RandomFn): number[];
+export declare function makeLeafArray(rpForest: FlatTree[]): ArrayLike<number>[];
+export declare function makeLeafArrayFlat(rpForest: FlatTree[]): {
+    flatLeafArray: Int32Array;
+    nLeaves: number;
+    leafSize: number;
+};
+export declare function searchFlatTree(point: Vector, tree: FlatTree, random: RandomFn): ArrayLike<number>;

package/dist/tree.js

@@ -1,5 +1,5 @@
 import * as utils from './utils.js';
-import { isWasmAvailable, buildRpTreeWasm, searchFlatTreeWasm, wasmTreeToJs } from './wasmBridge.js';
+import { isWasmAvailable, buildRpTreeWasmFlat, searchFlatTreeWasm } from './wasmBridge.js';
 export class FlatTree {
     constructor(hyperplanes, offsets, children, indices) {
         this.hyperplanes = hyperplanes;
@@ -8,14 +8,77 @@ export class FlatTree {
         this.indices = indices;
     }
     static fromWasm(wasmTree) {
-        const jsData = wasmTreeToJs(wasmTree);
-        const tree = new FlatTree(jsData.hyperplanes, jsData.offsets, jsData.children, jsData.indices);
+        const tree = new FlatTree([], new Float64Array(0), [], []);
+        const childrenFlat = wasmTree.children();
+        let maxLeafIdx = -1;
+        for (let i = 0; i < childrenFlat.length; i++) {
+            const v = childrenFlat[i];
+            if (v <= 0) {
+                const leafIdx = -v;
+                if (leafIdx > maxLeafIdx)
+                    maxLeafIdx = leafIdx;
+            }
+        }
+        const nLeaves = maxLeafIdx + 1;
+        const indicesFlat = wasmTree.indices();
+        const leafSize = nLeaves > 0 ? Math.floor(indicesFlat.length / nLeaves) : 0;
+        tree.hyperplanesFlat = wasmTree.hyperplanes();
+        tree.offsetsFlat = wasmTree.offsets();
+        tree.childrenFlat = childrenFlat;
+        tree.indicesFlat = indicesFlat;
+        tree.dim = wasmTree.dim();
+        tree.nNodes = wasmTree.n_nodes();
+        tree.nLeaves = nLeaves;
+        tree.leafSize = leafSize;
+        const hyperplanes = new Array(tree.nNodes);
+        for (let i = 0; i < tree.nNodes; i++) {
+            const start = i * tree.dim;
+            hyperplanes[i] = tree.hyperplanesFlat.subarray(start, start + tree.dim);
+        }
+        const children = new Array(tree.nNodes);
+        for (let i = 0; i < tree.nNodes; i++) {
+            const start = i * 2;
+            children[i] = tree.childrenFlat.subarray(start, start + 2);
+        }
+        const indices = new Array(tree.nLeaves);
+        for (let i = 0; i < tree.nLeaves; i++) {
+            const start = i * tree.leafSize;
+            indices[i] = tree.indicesFlat.subarray(start, start + tree.leafSize);
+        }
+        tree.hyperplanes = hyperplanes;
+        tree.offsets = tree.offsetsFlat;
+        tree.children = children;
+        tree.indices = indices;
         tree.wasmTree = wasmTree;
         return tree;
     }
     getWasmTree() {
         return this.wasmTree;
     }
+    getFlatHyperplanes() {
+        return this.hyperplanesFlat;
+    }
+    getFlatOffsets() {
+        return this.offsetsFlat;
+    }
+    getFlatChildren() {
+        return this.childrenFlat;
+    }
+    getFlatLeafMeta() {
+        if (this.indicesFlat &&
+            this.nLeaves !== undefined &&
+            this.leafSize !== undefined) {
+            return {
+                indices: this.indicesFlat,
+                nLeaves: this.nLeaves,
+                leafSize: this.leafSize,
+            };
+        }
+        return undefined;
+    }
+    getDim() {
+        return this.dim;
+    }
     dispose() {
         if (this.wasmTree) {
             this.wasmTree.free();
@@ -41,9 +104,15 @@ function makeForestWasm(data, leafSize, nTrees, random) {
     const nSamples = data.length;
     const dim = data[0].length;
     const forest = [];
+    const flatData = new Float64Array(nSamples * dim);
+    for (let i = 0; i < nSamples; i++) {
+        for (let j = 0; j < dim; j++) {
+            flatData[i * dim + j] = data[i][j];
+        }
+    }
     for (let i = 0; i < nTrees; i++) {
         const seed = Math.floor(random() * 0xFFFFFFFF);
-        const wasmTree = buildRpTreeWasm(data, nSamples, dim, leafSize, seed);
+        const wasmTree = buildRpTreeWasmFlat(flatData, nSamples, dim, leafSize, seed);
         forest.push(FlatTree.fromWasm(wasmTree));
     }
     return forest;
@@ -183,7 +252,23 @@ export function makeLeafArray(rpForest) {
     if (rpForest.length > 0) {
         const output = [];
         for (let tree of rpForest) {
-            output.push(...tree.indices);
+            if (tree.indices.length > 0) {
+                output.push(...tree.indices);
+                continue;
+            }
+            const flatMeta = tree.getFlatLeafMeta();
+            if (!flatMeta) {
+                continue;
+            }
+            const { indices, nLeaves, leafSize } = flatMeta;
+            for (let leaf = 0; leaf < nLeaves; leaf++) {
+                const start = leaf * leafSize;
+                const row = new Array(leafSize);
+                for (let i = 0; i < leafSize; i++) {
+                    row[i] = indices[start + i];
+                }
+                output.push(row);
+            }
         }
         return output;
     }
@@ -191,6 +276,44 @@ export function makeLeafArray(rpForest) {
         return [[-1]];
     }
 }
+export function makeLeafArrayFlat(rpForest) {
+    if (rpForest.length === 0) {
+        return { flatLeafArray: new Int32Array([-1]), nLeaves: 1, leafSize: 1 };
+    }
+    let leafSize = -1;
+    let totalLeaves = 0;
+    for (const tree of rpForest) {
+        const flatMeta = tree.getFlatLeafMeta();
+        if (flatMeta) {
+            if (leafSize === -1)
+                leafSize = flatMeta.leafSize;
+            totalLeaves += flatMeta.nLeaves;
+        }
+        else {
+            if (leafSize === -1)
+                leafSize = tree.indices[0]?.length ?? 0;
+            totalLeaves += tree.indices.length;
+        }
+    }
+    const flatLeafArray = new Int32Array(totalLeaves * leafSize);
+    let offset = 0;
+    for (const tree of rpForest) {
+        const flatMeta = tree.getFlatLeafMeta();
+        if (flatMeta) {
+            flatLeafArray.set(flatMeta.indices, offset);
+            offset += flatMeta.indices.length;
+            continue;
+        }
+        for (let i = 0; i < tree.indices.length; i++) {
+            const row = tree.indices[i];
+            for (let j = 0; j < leafSize; j++) {
+                flatLeafArray[offset + i * leafSize + j] = row[j];
+            }
+        }
+        offset += tree.indices.length * leafSize;
+    }
+    return { flatLeafArray, nLeaves: totalLeaves, leafSize };
+}
 function selectSide(hyperplane, offset, point, random) {
     let margin = offset;
     for (let d = 0; d < point.length; d++) {
@@ -213,6 +336,37 @@ export function searchFlatTree(point, tree, random) {
         const seed = Math.floor(random() * 0xFFFFFFFF);
         return searchFlatTreeWasm(wasmTree, point, seed);
     }
+    const childrenFlat = tree.getFlatChildren();
+    const hyperplanesFlat = tree.getFlatHyperplanes();
+    const offsetsFlat = tree.getFlatOffsets();
+    const dim = tree.getDim();
+    const flatMeta = tree.getFlatLeafMeta();
+    if (childrenFlat &&
+        hyperplanesFlat &&
+        offsetsFlat &&
+        dim !== undefined &&
+        flatMeta) {
+        let node = 0;
+        while (childrenFlat[node * 2] > 0) {
+            const offset = offsetsFlat[node];
+            const base = node * dim;
+            const side = selectSide(hyperplanesFlat.subarray(base, base + dim), offset, point, random);
+            if (side === 0) {
+                node = childrenFlat[node * 2];
+            }
+            else {
+                node = childrenFlat[node * 2 + 1];
+            }
+        }
+        const leafIdx = -childrenFlat[node * 2];
+        const { indices, leafSize } = flatMeta;
+        const start = leafIdx * leafSize;
+        const result = new Array(leafSize);
+        for (let i = 0; i < leafSize; i++) {
+            result[i] = indices[start + i];
+        }
+        return result;
+    }
     let node = 0;
     while (tree.children[node][0] > 0) {
         const side = selectSide(tree.hyperplanes[node], tree.offsets[node], point, random);

package/dist/umap.js

@@ -237,6 +237,34 @@ export class UMAP {
         const nTrees = 5 + Math.floor(round(X.length ** 0.5 / 20.0));
         const nIters = Math.max(5, Math.floor(Math.round(log2(X.length))));
         this.rpForest = tree.makeForest(X, nNeighbors, nTrees, this.random, this.useWasmTree);
+        if (this.useWasmNNDescent && wasmBridge.isWasmAvailable()) {
+            const nSamples = X.length;
+            const dim = X[0].length;
+            const flatData = new Float64Array(nSamples * dim);
+            for (let i = 0; i < nSamples; i++) {
+                for (let j = 0; j < dim; j++) {
+                    flatData[i * dim + j] = X[i][j];
+                }
+            }
+            const { flatLeafArray, nLeaves, leafSize } = tree.makeLeafArrayFlat(this.rpForest);
+            const distanceMetric = distanceFn.name === 'cosine' ? 'cosine' : 'euclidean';
+            const seed = Math.floor(this.random() * 0xFFFFFFFF);
+            const result = wasmBridge.nnDescentWasmFlat(flatData, nSamples, dim, flatLeafArray, nLeaves, leafSize, nNeighbors, nIters, 50, 0.001, 0.5, true, distanceMetric, seed);
+            const indices = [];
+            const weights = [];
+            const offset1 = nSamples * nNeighbors;
+            for (let i = 0; i < nSamples; i++) {
+                const rowIndices = [];
+                const rowWeights = [];
+                for (let j = 0; j < nNeighbors; j++) {
+                    rowIndices.push(result[i * nNeighbors + j]);
+                    rowWeights.push(result[offset1 + i * nNeighbors + j]);
+                }
+                indices.push(rowIndices);
+                weights.push(rowWeights);
+            }
+            return { knnIndices: indices, knnDistances: weights };
+        }
         const leafArray = tree.makeLeafArray(this.rpForest);
         const { indices, weights } = metricNNDescent(X, leafArray, nNeighbors, nIters);
         return { knnIndices: indices, knnDistances: weights };
@@ -255,10 +283,7 @@ export class UMAP {
             const b = wasmBridge.sparseMultiplyScalarWasm(a, setOpMixRatio);
             const c = wasmBridge.sparseMultiplyScalarWasm(prodMatrix, 1.0 - setOpMixRatio);
             const resultWasm = wasmBridge.sparseAddWasm(b, c);
-            const entries = wasmBridge.wasmSparseMatrixGetAll(resultWasm);
-            const jsRows = entries.map(e => e.row);
-            const jsCols = entries.map(e => e.col);
-            const jsVals = entries.map(e => e.value);
+            const { rows: jsRows, cols: jsCols, values: jsVals } = wasmBridge.wasmSparseMatrixGetAllTyped(resultWasm);
             return new matrix.SparseMatrix(jsRows, jsCols, jsVals, size);
         }
         if (this.useWasmMatrix && wasmBridge.isWasmAvailable()) {
@@ -270,10 +295,7 @@ export class UMAP {
             const b = wasmBridge.sparseMultiplyScalarWasm(a, setOpMixRatio);
             const c = wasmBridge.sparseMultiplyScalarWasm(prodMatrix, 1.0 - setOpMixRatio);
             const resultWasm = wasmBridge.sparseAddWasm(b, c);
-            const entries = wasmBridge.wasmSparseMatrixGetAll(resultWasm);
-            const jsRows = entries.map(e => e.row);
-            const jsCols = entries.map(e => e.col);
-            const jsVals = entries.map(e => e.value);
+            const { rows: jsRows, cols: jsCols, values: jsVals } = wasmBridge.wasmSparseMatrixGetAllTyped(resultWasm);
             return new matrix.SparseMatrix(jsRows, jsCols, jsVals, size);
         }
         const sparseMatrix = new matrix.SparseMatrix(rows, cols, vals, size);

package/dist/wasmBridge.d.ts

@@ -12,6 +12,7 @@ export interface WasmFlatTree {
     free(): void;
 }
 export declare function buildRpTreeWasm(data: number[][], nSamples: number, dim: number, leafSize: number, seed: number): WasmFlatTree;
+export declare function buildRpTreeWasmFlat(flatData: Float64Array, nSamples: number, dim: number, leafSize: number, seed: number): WasmFlatTree;
 export declare function searchFlatTreeWasm(tree: WasmFlatTree, point: number[], seed: number): number[];
 export declare function wasmTreeToJs(wasmTree: WasmFlatTree): {
     hyperplanes: number[][];
@@ -55,7 +56,13 @@ export declare function wasmSparseMatrixGetAll(matrix: WasmSparseMatrix): {
     row: number;
     col: number;
 }[];
-export declare function nnDescentWasm(data: number[][], leafArray: number[][], nNeighbors: number, nIters?: number, maxCandidates?: number, delta?: number, rho?: number, rpTreeInit?: boolean, distanceMetric?: string, seed?: number): number[][][];
+export declare function wasmSparseMatrixGetAllTyped(matrix: WasmSparseMatrix): {
+    rows: Int32Array;
+    cols: Int32Array;
+    values: Float64Array;
+};
+export declare function nnDescentWasm(data: number[][], leafArray: ArrayLike<number>[], nNeighbors: number, nIters?: number, maxCandidates?: number, delta?: number, rho?: number, rpTreeInit?: boolean, distanceMetric?: string, seed?: number): number[][][];
+export declare function nnDescentWasmFlat(flatData: Float64Array, nSamples: number, dim: number, flatLeafArray: Int32Array, nLeaves: number, leafSize: number, nNeighbors: number, nIters?: number, maxCandidates?: number, delta?: number, rho?: number, rpTreeInit?: boolean, distanceMetric?: string, seed?: number): number[];
 export interface WasmOptimizerState {
     head_embedding: Float64Array;
     current_epoch: number;

package/dist/wasmBridge.js

@@ -64,6 +64,11 @@ export function buildRpTreeWasm(data, nSamples, dim, leafSize, seed) {
             flatData[i * dim + j] = data[i][j];
         }
     }
+    return buildRpTreeWasmFlat(flatData, nSamples, dim, leafSize, seed);
+}
+export function buildRpTreeWasmFlat(flatData, nSamples, dim, leafSize, seed) {
+    if (!wasmModule)
+        throw new Error('WASM module not initialized');
     return wasmModule.build_rp_tree(flatData, nSamples, dim, leafSize, BigInt(seed));
 }
 export function searchFlatTreeWasm(tree, point, seed) {
@@ -179,17 +184,22 @@ export function sparseGetCSRWasm(matrix) {
     return { indices, values, indptr };
 }
 export function wasmSparseMatrixToArray(matrix) {
-    const flat = Array.from(matrix.to_array());
+    const flat = matrix.to_array();
     const nRows = matrix.n_rows;
     const nCols = matrix.n_cols;
     const result = [];
     for (let i = 0; i < nRows; i++) {
-        result.push(flat.slice(i * nCols, (i + 1) * nCols));
+        const row = new Array(nCols);
+        const start = i * nCols;
+        for (let j = 0; j < nCols; j++) {
+            row[j] = flat[start + j];
+        }
+        result.push(row);
     }
     return result;
 }
 export function wasmSparseMatrixGetAll(matrix) {
-    const flat = Array.from(matrix.get_all_ordered());
+    const flat = matrix.get_all_ordered();
     const entries = [];
     for (let i = 0; i < flat.length; i += 3) {
         entries.push({
@@ -200,6 +210,21 @@ export function wasmSparseMatrixGetAll(matrix) {
     }
     return entries;
 }
+export function wasmSparseMatrixGetAllTyped(matrix) {
+    const flat = matrix.get_all_ordered();
+    const count = Math.floor(flat.length / 3);
+    const rows = new Int32Array(count);
+    const cols = new Int32Array(count);
+    const values = new Float64Array(count);
+    let out = 0;
+    for (let i = 0; i < flat.length; i += 3) {
+        rows[out] = flat[i];
+        cols[out] = flat[i + 1];
+        values[out] = flat[i + 2];
+        out += 1;
+    }
+    return { rows, cols, values };
+}
 export function nnDescentWasm(data, leafArray, nNeighbors, nIters = 10, maxCandidates = 50, delta = 0.001, rho = 0.5, rpTreeInit = true, distanceMetric = 'euclidean', seed = 42) {
     if (!wasmModule)
         throw new Error('WASM module not initialized');
@@ -219,7 +244,7 @@ export function nnDescentWasm(data, leafArray, nNeighbors, nIters = 10, maxCandi
             flatLeafArray[i * leafSize + j] = leafArray[i][j];
         }
     }
-    const result = wasmModule.nn_descent(flatData, nSamples, dim, flatLeafArray, nLeaves, leafSize, nNeighbors, nIters, maxCandidates, delta, rho, rpTreeInit, distanceMetric, BigInt(seed));
+    const result = nnDescentWasmFlat(flatData, nSamples, dim, flatLeafArray, nLeaves, leafSize, nNeighbors, nIters, maxCandidates, delta, rho, rpTreeInit, distanceMetric, seed);
     const indices = [];
     const distances = [];
     const flags = [];
@@ -240,6 +265,11 @@ export function nnDescentWasm(data, leafArray, nNeighbors, nIters = 10, maxCandi
     }
     return [indices, distances, flags];
 }
+export function nnDescentWasmFlat(flatData, nSamples, dim, flatLeafArray, nLeaves, leafSize, nNeighbors, nIters = 10, maxCandidates = 50, delta = 0.001, rho = 0.5, rpTreeInit = true, distanceMetric = 'euclidean', seed = 42) {
+    if (!wasmModule)
+        throw new Error('WASM module not initialized');
+    return wasmModule.nn_descent(flatData, nSamples, dim, flatLeafArray, nLeaves, leafSize, nNeighbors, nIters, maxCandidates, delta, rho, rpTreeInit, distanceMetric, BigInt(seed));
+}
 export function createOptimizerState(head, tail, headEmbedding, tailEmbedding, epochsPerSample, epochsPerNegativeSample, moveOther, initialAlpha, gamma, a, b, dim, nEpochs, nVertices) {
     if (!wasmModule)
         throw new Error('WASM module not initialized');