@wlearn/dimred

Dimensionality reduction authored in C, executed by the caller's Polygrad runtime. One algorithm implementation serves Python, Node and the browser; there is no embedded second Polygrad core and no JSON graph author.

Implemented:

• PCA: centered block power iteration with reorthogonalization, fitted loadings, sample variances and a residual-based convergence check. Up to 16 components. Rank-deficient components can be zero; signs and bases within repeated eigenspaces are not unique. This is an iterative solver, not SVD.
• Exact t-SNE: calibration, affinities, scheduled updates, saved-model continuation. Quadratic storage and work; 4–2,000 points.
• Sparse-affinity t-SNE (opt-in): calibrated kNN affinities, symmetric attraction and blocked exact repulsion. Accepts precomputed approximate neighbours; does not store dense N×N affinities. No 2,000-point guard.
• UMAP, experimental synchronous variant: exact blocked Euclidean kNN, local-connectivity calibration and fuzzy union, weighted attractive forces and five sampled negatives per edge. See the algorithm contract below.
• TriMap: scaled-distance triplet selection, tempered weights and full-batch lookahead delta-bar-delta optimization. Exact neighbours and a seeded integer sampler replace Annoy/NumPy RNG. See the contract below.

Version 0.1.0 is the first release candidate. It supports Polygrad 0.7.x. Native portability has been tested on Linux, not macOS/Windows.

Package API

JavaScript uses @wlearn/dimred; Python installs wlearn-dimred and exposes the same classes through wlearn_dimred. The wlearn.dimred facade requires wlearn >=0.3.1 (pip install 'wlearn[dimred]>=0.3.1'). The Polygrad dependency range is >=0.7.0 <0.8.0; 0.6 lacks the required extension loader. Build authors against the matching Polygrad v0.7.0 headers.

const { PCA, UMAP } = require('@wlearn/dimred')
async function main() {
  const X = Array.from({length: 24}, (_, i) => [i, Math.sin(i), Math.cos(i)])
  const pca = await PCA.create({ nComponents: 2, polygrad: {core: 'wasm', device: 'cpu'} })
  try {
    const scores = await pca.fitTransform(X) // also accepts a DenseMatrix
    const projected = await pca.transform(X.slice(0, 3))
    const bytes = pca.save() // WLRN, synchronous; Python can load it
    console.log(scores.rows, projected.rows, bytes.length)
  } finally { pca.dispose() }
}
main().catch(console.error)

from wlearn.dimred import PCA, UMAP
import numpy as np

X = np.array([[i, np.sin(i), np.cos(i)] for i in range(24)], dtype=np.float32)
pca = PCA(n_components=2)
try:
    scores = pca.fit_transform(X)
    projected = pca.transform(X[:3])
    pca.save("pca.wlrn")
finally:
    pca.dispose()

PCA, TSNE, UMAP, and TriMap provide getParams/setParams, fit/fitTransform, isFitted, capabilities, save, load, and dispose. Python uses snake_case and embedding_; JS exposes a copied embedding DenseMatrix. JS fitting and PCA transforms are asynchronous on all backends; Python is synchronous. Hyperparameters use camelCase in JS and snake_case in Python, with identical names in WLRN metadata. maxIter/max_iter controls the iteration budget. TSNE defaults to method: 'exact'; opt into 'sparse' explicitly. Nonlinear methods currently support two output dimensions only.

PCA can be an intermediate wlearn Pipeline transformer. Its JS components matrix is [features, components]; Python components_ follows sklearn's [components, features]. Both use the same C transform and WLRN basis layout. The nonlinear methods are transductive: save/load preserves the fitted embedding and parameters, not an unseen-point transform or optimizer checkpoint. They deliberately have no transform, predict, or supervised score and are not AutoML candidates. Advanced Model drivers below retain checkpoint support.

Importing dimred registers its four wlearn.dimred.*@1 WLRN loaders with @wlearn/core / wlearn. In tree-shaken JS bundles, explicitly call the exported registerWlearn() when using only core load(); a bare side-effect import may be removed. Load a PCA bundle with PCA.load(bytes, {polygrad: pg}) or PCA.load(path, polygrad=runtime) to reuse a runtime. Supplying polygrad runtime options instead creates and disposes a runtime per operation. Sharing a runtime preserves compiler caches; callers must serialize its access. Failed refits preserve the previous fitted state. Changing parameters clears it.

Dimred is opt-in, not eagerly imported by the SDK or wlearn[all]. wlearn[dimred] and the wlearn.dimred facade need the updated wlearn package; use wlearn_dimred directly until that is published. No unscoped dimred compatibility package is added.

make package stages C sources, generated bridge wrappers and public Polygrad headers into each package. It embeds small Wasm graph authors, not another Polygrad runtime. npm installation attempts an optional native build; without a C compiler/Node headers, use polygrad: {core: 'wasm', device: 'cpu'}. Browsers use the same package through a bundler with device: 'webgpu' or 'cpu'. Python wheels compile native authors; source builds require a C compiler. Incompatible Polygrad versions are rejected. Package build/size and smoke tests must be repeated for each release platform; Windows native authors are not implemented.

From the source checkout, make test-api test-types checks estimators and Python↔JS bundles; make test-api-browser checks the packaged browser API. make test-package builds an npm archive and an sdist/wheel, compiles the extracted native authors and tests installed artifacts with the local dependency checkouts. It needs Python's build, wheel, NumPy and setuptools77+; it does not claim registry dependency resolution before the pending Polygrad release.

Build and try

Build the sibling Polygrad native/Wasm/browser artifacts first. Author generation requires Clang and Python 3.10+; native builds need a C compiler and Node headers; Wasm builds need Emscripten. Override POLYGRAD_PATH and NODE_INCLUDE if needed.

make all POLYGRAD_PATH=/home/anton/projects/polygrad/polygrad
make demo-c-api POLYGRAD_PATH=/home/anton/projects/polygrad/polygrad

Open http://127.0.0.1:8081 and select PCA, t-SNE, synchronous UMAP or TriMap, and Wasm CPU or WebGPU. Everything runs locally. WebGPU failure is reported, not silently replaced with another backend. PORT=8082 changes port. Start with 128 points: construction/compile cost has not been tuned.

The current grouped UMAP update is experimental: at 256 points it improved Wasm fit time but regressed WebGPU, and both still trail umap-js. It is not a backend-independent speedup; see the local PLAN for measurements and next work.

An optional direct incoming reduction is available with segment_sum=True (Python) or segmentSum: true (JS), together with Polygrad Runtime(kernels=True) or create({kernels: true}). It uses the same C author and ordinary Tensor operations, retaining gradients and portable Models. Without the kernel policy it remains correct but falls back to a dense masked reduction. The default remains the grouped scan; enabling segments changes summation order, not sampling or the optimizer.

Matched Chrome runs on 256-point digits/Swiss datasets (300 epochs, seed42, median of three warm whole fits) reduced Wasm time by 10% and WebGPU time by 30–31%. Digits: Wasm787→708ms, WebGPU1941→1351ms; umap-js213ms. WebGPU dispatches per step fell from39 to26. This still trails umap-js, whose sequential optimizer is different, and does not establish large-data scaling. Reproduce the optional checks with make test-projections-python PROJECTION_ARGS=--segments and SEGMENT_SUM=1 make test-projections-node; CORE=wasm selects Node Wasm. BENCH_SEGMENTS=1 make bench-npm compares both Dimred reductions and umap-js in one browser session.

Python, with dimred/py and Polygrad on PYTHONPATH:

import numpy as np
from polygrad import create
from wlearn_dimred.graphs import pca, umap

X = np.random.default_rng(42).normal(size=(128, 8)).astype(np.float32)
with create(device='CPU', logical='always') as pg:
    fitted = pca(pg, X, components=2)
    Y = umap(pg, X, k=15, initial=fitted['embedding'])
# New points: (new_X - fitted['mean']) @ fitted['components'].

Browser (load js/graphs.js after Polygrad):

const pg = await polygrad.create({device: 'webgpu', logical: 'always'})
const emb = projectionBindings(pg, async name =>
  pg.loadExtension(await (await fetch('/' + name + '.wasm')).arrayBuffer()))
try {
  const X = Float32Array.from({length: 128 * 8}, (_, i) => Math.sin(i * .73))
  const fitted = await emb.pca(X, 128, 8)
  const Y = await emb.umap(X, 128, 8, {initial: fitted.embedding})
} finally { await pg.dispose() }

Node uses the same binding function with pg.loadExtension(require(path)) for build/{pca,knn,umap}/native.node; its Wasm path loads the corresponding author.wasm. test/projections.cjs is a complete runnable example. pca_model/pcaModel and umap_model/umapModel expose ordinary Models for save/load and explicit entrypoint calls. The t-SNE lifecycle is documented in the demo README.

Sparse-affinity t-SNE

Choose tsne_sparse in Python or tsneSparse in JS. The interactive demo also offers Sparse-affinity t-SNE. Dense exact t-SNE remains unchanged.

from wlearn_dimred.graphs import tsne_sparse

# Existing Polygrad runtime; X is a finite float32 N×D array.
Y = tsne_sparse(pg, X, perplexity=20, k=60, block_size=256)
# Or supply exact/approximate neighbours from another search implementation:
Y = tsne_sparse(pg, X, perplexity=20, neighbours=(indices, distances))

const Y = await emb.tsneSparse(X, n, d, {
  perplexity: 20, k: 60, blockSize: 256,
  // Optional: neighbours: {indices, distances}
})

Neighbour arrays have shape N×k, exclude self, and contain unique indices per row. Distances are Euclidean, not squared, finite and nonnegative. The built-in exact blocked neighbour search currently supports k≤64; precomputed graphs may use larger k. Perplexity must be ≤k; ties that make it unattainable are rejected. Calibration must converge before optimization begins.

This approximates the high-dimensional affinities by retaining neighbours. It does not replace t-SNE repulsion with negative sampling. Repulsion still evaluates every pair, in blocks: quadratic time remains. With k=N−1 the force matches the dense implementation within floating-point tolerance. There is no Barnes–Hut, FFT/interpolation or built-in approximate neighbour search yet. Storage savings require k much smaller than N; k=N−1 intentionally loses them.

State consists of neighbour arrays, conditional weights, optimizer state and incoming-edge groups. Pairwise repulsion workspace is N×blockSize, rather than N×N. Compiler temporaries and cached buffers also use memory; this is not a peak-memory guarantee. Edge sorting uses bounded int32-key chunks so large graphs do not fall back to quadratic duplicate-index recovery.

segment_sum=True / segmentSum:true, with kernels=True / kernels:true on the runtime, reuses Polygrad's optional incoming interval-reduction kernel. The default is the portable grouped scan. Both use the same objective; sums may differ slightly in order. Without the runtime kernel policy, the interval form is correct but can be very slow.

tsne_sparse_model / tsneSparseModel exposes a portable Model. Its protocol is initialize → calibrate (check status) → prepare, then for each iteration begin → repulse × ceil(N/blockSize) → step. Keep Y unchanged during repulsion. An incomplete or repeated step does not update Y. inspect returns the gradient for the accumulated repulsion; use it after all blocks. Model save/load retains partial accumulators as well as optimizer state. Defaults are 500 iterations, exaggeration 12 until iteration 250, momentum .5 then .8, and learning rate 50 with the explicit factor 4 in the gradient. Supply initial for a controlled comparison; otherwise the driver uses PCA coordinates.

Validation: make test-tsne-sparse-python, make test-tsne-sparse-node (CORE=wasm optional), and make test-tsne-sparse-browser. TSNE_ARGS=--segments or SEGMENT_SUM=1 enables the optional reduction in Python/Node checks. The browser check runs both policies. make bench-tsne-sparse BENCH_OUTPUT=/absolute/fresh/path runs 100/1,000/10,000-point browser fits; BENCH_SEGMENTS=1 selects the interval reduction. These benchmarks explicitly use external SciPy neighbours, report their time separately, validate the initial gradient/grouping, and score a fixed subset of at most 1,000 points. They are not end-to-end neighbour-search comparisons with another t-SNE library.

UMAP algorithm contract

This is not trajectory parity with umap-learn. The calibration and pair-force formulas follow umap-learn 0.5.7. Differences are deliberate and visible:

• k counts neighbours excluding self (reference n_neighbors is k+1). Euclidean metric, local connectivity 1, fuzzy union only.
• Exact kNN uses query blocks, not an N×N distance matrix. Search remains quadratic in pair count. The graph stores O(Nk) neighbour data; reverse-edge assembly compares O(Nk²) entries. This is not approximate neighbour search.
• Each step evaluates all weighted positive edges at the same layout snapshot. Mutual edges have half their union weight per direction. Negative sampling randomly orients each undirected edge, updates only its sampled anchor, and clips each pair force to [-4,4]. Integer hashing makes sampling reproducible from seed and step; there is no host RNG call in the optimization loop.
• Batch updates use learning_rate / k, decaying linearly over the requested epochs. Defaults are 300 epochs and learning rate .1. These are not interchangeable with umap-learn's sequential SGD settings.
• a=1.5769434603, b=.8950608779 correspond to reference min_dist=.1, spread=1. Low-level Models accept other positive fitted a,b; fitting custom min_dist/spread curves is not implemented. Initialization defaults to PCA, not spectral embedding. A caller can supply finite N×2 coordinates.
• on_step/onStep gets coordinates every ten steps and may return false to cancel. JS also accepts an AbortSignal during optimization. Neighbour construction/calibration cancellation is not implemented.

The sparse graph does not prove linear-memory kernel execution: Polygrad's gather/autograd lowering can introduce larger intermediates. Large-data memory, throughput and real-embedding quality need measurements before making a scaling claim. No unseen-point UMAP transform, densMAP, supervision, alternate metrics, approximate kNN is implemented.

TriMap contract

The objective, half-gradient convention, lookahead momentum and delta-bar-delta gains follow the original implementation. Defaults:12inliers,4outliers,3random triplets per point,400steps,learning rate.1; 2-D Euclidean output. Scale is the mean distance to neighbours3–5. Inliers are ranked by scaled similarity, not simply raw distance. Outliers exclude self and all higher-ranked inliers, using a bounded complement-rank mapping. Random triplets are ordered by similarity. Weights use temperature.5 and the original generate_triplets path's .1 multiplier on random-triplet weights. Its separate known-kNN path omits that multiplier; we deliberately follow the main path.

Differences: exact blocked neighbours instead of Annoy, deterministic uint32 hash sampling instead of NumPy RNG, iterative PCA initialization scaled by.01, FP32 graph reductions instead of the reference's Numba loops. No high-dimensional PCA preprocessing, alternate metric or optimizer choices. Triplet indices can differ from the reference; given triplets/weights/layouts, gradient and optimizer steps are tested against the reference. Final coordinates are not promised equal.

Python trimap/trimap_model and JS trimap/trimapModel mirror UMAP's lifecycle and progress/cancellation hooks. All numerical composition remains in C. CPU, CUDA, Node native/Wasm and Chrome Wasm/WebGPU checks pass; X86 optimizer checks pass but the full-fit test currently fails in its47-neighbour kNN execution: a kernel has31arguments, beyond the X86 call dispatcher's16-argument limit. This limitation is not bypassed by silently switching backend.

npm comparison

For current t-SNE versus literal tsne@1.0.1 (not tsne-js):

TSNE_NPM_SOURCE=/absolute/path/to/tsne/tsne.js make bench-tsne-compare \
  BENCH_OUTPUT=/absolute/fresh/results BENCH_SIZES=100,1000,10000 \
  PYTHON=/path/to/python

This runs500iterations on the same seeded Swiss-roll inputs and initial coordinates, including Dimred's browser neighbor search. It records setup, first-step/JIT and optimization times, finite-output/cleanup checks, full-data neighbor recall and subset trustworthiness. Dense mode stops at its existing 2000-point limit; sparse mode covers10k. npm retains its float64 arithmetic, per-step cost computation and optimizer schedule; this is an application-level comparison, not identical work. Requires headed Chrome/WebGPU (use xvfb-run -a if there is no working display). Reference package installation is separate.

bench/README.md describes reproducible whole-fit comparisons with umap-js1.4.0 and ml-pca4.1.1, including independent layout quality. These are small-data tests of current implementations, not equivalent-workload or large-dataset scaling claims. Reference packages are benchmark-only dependencies.

Source layout and tests

src/tsne.c, src/pca.c, src/umap.c and src/trimap.c own their algorithms; src/neighbors.c owns shared neighbour construction. Private construction helpers live in src/graph.h; src/dimred.h declares the public author entrypoints. They call the existing C API. Polygrad generates the native/Wasm bridge; frontend files currently allocate, bind and drive Models. Build artifacts live under build/. No per-operation callback table is maintained. The superseded JSON and curated-callback prototypes are no longer build targets.

For checks, install NumPy, SciPy, scikit-learn and umap-learn==0.5.7 in the chosen Python environment. These are test references, not browser dependencies.

make test-python PYTHON=/path/to/python
make test-node PYTHON=/path/to/python
make test-projections-python PYTHON=/path/to/python DEVICE=CUDA
make test-projections-node PYTHON=/path/to/python DEVICE=cuda
make test-projections-node PYTHON=/path/to/python CORE=wasm
make test-browser PYTHON=/path/to/python

Checks compare PCA subspaces/variance to sklearn, neighbours to SciPy, UMAP weights to the reference Python calibration and forces to independent NumPy expressions. They also test checkpoint continuation, final synthetic quality and cleanup. The Python reference function is intentional: on this environment, Numba's compiled calibration returns infinite bandwidths for some finite rows, while the unchanged Python function converges. This is not a claim of final coordinate equality with the reference or across devices.