wlearn-dimred

Dimensionality reduction for Python: PCA, t-SNE, UMAP and TriMap, running on the standalone C backend by default. Optional CPU/GPU acceleration uses Polygrad. Part of wlearn.

Install

The C backend shares fixed math routines with the JavaScript/Wasm backend for reproducible seeded fits within a package version. Polygrad uses different algorithms/arithmetic and need not produce the same coordinates.

pip install wlearn-dimred
# or, together with wlearn:
pip install 'wlearn[dimred]'

Requires Python 3.10 or newer. The package is distributed as source and compiles a small C extension during installation, so a C compiler is needed. Polygrad is optional: pip install 'wlearn-dimred[polygrad]' enables its backend.

Quick start

import numpy as np
from wlearn_dimred import PCA, UMAP

# Three clusters in 4 dimensions
rng = np.random.default_rng(0)
centers = np.array([[0, 0, 0, 0], [5, 0, 0, 0], [0, 5, 0, 0]], dtype=np.float32)
X = (centers[np.arange(150) % 3] + rng.normal(scale=0.5, size=(150, 4))).astype(np.float32)

pca = PCA(n_components=2)
umap = UMAP(n_neighbors=15)
try:
    scores = pca.fit_transform(X)   # shape (150, 2)
    layout = umap.fit_transform(X)  # shape (150, 2)

    pca.save("pca.wlrn")
    restored = PCA.load("pca.wlrn")
    print(scores.shape, layout.shape, restored.transform(X[:3]).shape)
    restored.dispose()
finally:
    pca.dispose()
    umap.dispose()

With wlearn 0.3.1 or newer, the same classes are also available as from wlearn.dimred import PCA, UMAP.

The C backend does not compile kernels. The optional Polygrad backend compiles on first use and can reuse cached kernels on subsequent fits.

Methods

Every class accepts backend='c' | 'polygrad' (default 'c') and polygrad. Supplying polygrad selects it unless backend is explicit. C PCA has no 16-component limit; max_iter/tolerance configure only Polygrad PCA. C t-SNE defaults to 'barnes_hut'; Polygrad to 'exact'. C UMAP's default learning rate is 1, versus 0.1 for Polygrad. Unknown parameters raise a ValidationError. PCA accepts up to min(features, samples - 1) components on C and at most 16 on Polygrad. Sparse Polygrad t-SNE accepts block_size (256); Polygrad t-SNE and UMAP accept segment_sum (False). C sparse/Barnes-Hut t-SNE uses at most 64 neighbours and requires perplexity <= min(64, samples - 1).

API

All classes:

• fit(X, y=None, on_step=None) and fit_transform(X, y=None, on_step=None). X is a 2-D array; y is ignored.
• embedding_ is the fitted layout; is_fitted reports state.
• get_params() and set_params(**params) read and change parameters. Changing parameters clears the fitted state.
• save(path=None) returns the model as bytes and optionally writes them. load(source, polygrad=None) accepts bytes or a path.
• dispose() releases native resources. Call it when you are done.
• capabilities describes what the estimator supports.

PCA also has transform(X), components_ (components x features, as in scikit-learn), mean_ and explained_variance_.

For t-SNE, UMAP and TriMap, on_step(iteration, embedding) is called every 10 iterations with the current coordinates. Return False to cancel; the fit then raises InterruptedError on C. On Polygrad, it raises RuntimeError for exact t-SNE and InterruptedError for the other methods. C also reports the final iteration when it is not a multiple of 10.

Choosing a runtime

polygrad is either a dict of polygrad.create() options or an existing runtime:

import numpy as np
import polygrad
from wlearn_dimred import UMAP

# Two groups of points in 3 dimensions
X = np.random.default_rng(1).normal(size=(60, 3)).astype(np.float32)
X[30:] += 4

with polygrad.create(device='CPU') as pg:
    umap = UMAP(polygrad=pg, max_iter=100)
    try:
        print(umap.fit_transform(X).shape)
    finally:
        umap.dispose()

• Options, for example polygrad={'device': 'CUDA'}, create a runtime for each fit or transform and dispose it afterwards.
• An existing runtime is borrowed and never disposed by dimred. It keeps compiled kernels between fits. Do not use one runtime from two estimators at the same time.

Faster UMAP and sparse t-SNE

For larger inputs, enable Polygrad's dedicated reduction kernel with both segment_sum=True and polygrad={'kernels': True}:

from wlearn_dimred import UMAP

umap = UMAP(segment_sum=True, polygrad={'kernels': True})
print(umap.get_params()['segment_sum'])
umap.dispose()

segment_sum without kernels gives correct results but is slow. On 10,000 points, this made UMAP optimization 5 to 8 times faster on WebGPU and WebAssembly. Results can differ slightly because the summation order changes.

Precomputed neighbours

wlearn_dimred.graphs.tsne_sparse runs sparse t-SNE on neighbours from your own search, for example an approximate nearest-neighbour index. indices and distances have shape (N, k), exclude each point itself, and hold Euclidean (not squared) distances:

import numpy as np
import polygrad
from wlearn_dimred.graphs import tsne_sparse

X = np.random.default_rng(2).normal(size=(80, 3)).astype(np.float32)
X[40:] += 4
d = np.linalg.norm(X[:, None] - X[None], axis=-1)
np.fill_diagonal(d, np.inf)
indices = np.argsort(d, axis=1)[:, :30].astype(np.int32)
distances = np.take_along_axis(d, indices, axis=1).astype(np.float32)

with polygrad.create(device='CPU', logical='always') as pg:
    Y = tsne_sparse(pg, X, perplexity=10, neighbours=(indices, distances))
print(Y.shape)

Saving and loading

save() writes the wlearn .wlrn format. The npm package @wlearn/dimred reads the same files. Importing wlearn_dimred registers its loaders, so wlearn.load() also works.

t-SNE, UMAP and TriMap save the fitted layout and parameters. They cannot place new points.

Notes and limitations

• t-SNE, UMAP and TriMap return 2-D layouts and have no transform or predict. They are not used as AutoML candidates. PCA can be a step in a wlearn Pipeline.
• C PCA uses double-precision centered covariance/Gram and a direct eigensolver. Inputs and portable arrays remain float32; it is not a direct SVD. The following iterative restrictions concern only Polygrad PCA.
• Polygrad PCA uses float32 covariance/Jacobi or subspace iteration, not an exact SVD. Up to 64 features it uses the full covariance. Larger inputs start with subspace iteration; up to 1,024 features a full solve is available as a fallback within the same max_iter budget. The budget counts subspace steps or full Jacobi sweeps. PCA checks the residual against the original data and raises if the requested tolerance is not reached; very small tolerances may be unattainable in float32. Component signs and repeated-eigenvalue bases are not unique. Nonlinear methods can use an approximate PCA initializer; their max_iter controls layout optimization, not PCA convergence.
• Exact t-SNE (method='exact') supports 4 to 2,000 points. method='sparse' has no fixed 2,000-point limit, but still compares all pairs, so time grows quadratically and available memory limits the input size.
• C 'barnes_hut' t-SNE approximates repulsion using a quadtree. 'sparse' retains exact repulsion on both backends. Polygrad has no Barnes-Hut path.
• C neighbor search is exact for small/low-dimensional inputs; above 2,000 points and 16 features it uses approximate neighbor refinement. Polygrad compares all pairs. n_neighbors is at most 64.
• C UMAP uses sequential scheduled updates. Polygrad updates all edges at once. TriMap uses the selected backend's neighbour search instead of Annoy. Layouts are reproducible for a fixed seed but differ from umap-learn and trimap. All three nonlinear methods start from PCA coordinates.
• Tested on Linux. Windows is not supported.
• The X86 kernel-argument limit affecting wider PCA and TriMap fits is fixed in Polygrad 0.7.1, the minimum supported optional runtime.

Compatibility

Requires wlearn 0.3 or newer. Only the optional accelerated backend requires polygrad >=0.7.1,<0.8.0. The default backend and fitted artifacts do not require it.

License

Apache-2.0. See NOTICE for third-party attributions.