manifolds-rs 0.2.2

Embedding methods implemented in Rust: (parametric) UMAP, tSNE, PHATE, Diffusion Map and PacMAP.
Documentation

CI Crates.io License: MIT

manifolds-rs

High-performance manifold learning and dimensionality reduction algorithms implemented in Rust. Contains as for now:

  • UMAP
    • Has different optimisers: SGD (traditional), Adam and a parallelised version of Adam for very fast fitting.
    • Optional GPU-accelerated kNN search via cubecl.
  • Parametric UMAP (optional feature)
  • tSNE
    • Barnes Hut tSNE (With a O(n log n) complexity).
    • Fast Fourier Transform-accelerated Interpolation-based t-SNE (Flt-SNE) (optional feature; with a O(n) complexity for large datasets).
    • Optional GPU-accelerated kNN search.
  • PHATE
  • PaCMAP
  • Diffusion Maps
    • Classical diffusion maps (Coifman & Lafon, 2006) with anisotropic (alpha) normalisation for density correction.
    • Optional landmark-based approximation with Nystroem extension for larger datasets.

Description

Rust implementations of various methods to project data onto two dimensions, i.e, learn low dimensional manifolds from the data. The current crate contains the big classic UMAP and tSNE (with the Barnes-Hut implementation and optionally the FFT-acceleration version). These are typically used methods for visualising high-dimensional biological data, but not without controversy. Moreover, the crate also provides via the Burn DL framework optionally parametric UMAP that can be optionally be used via the prospective feature flag. Since release 0.1.8, we also have PHATE. With 0.1.9, PaCMAP has been also implemented. More recently, classical diffusion maps have been added as well. Changelog can be found here)

Features

  • UMAP algorithm: Complete implementation of the UMAP dimensionality reduction algorithm with several optimisations: SGD, Adam and a parallelised version of ADAM for increased optimisation speed.
  • tSNE algorithm: Implementation of the Barnes-Hut accelerated version and the FFT-accelerated version (optional).
  • PHATE: Implementation of Potential of Heat-diffusion for Affinity-based Trajectory Embedding with different landmark methods.
  • Diffusion Maps: Classical diffusion maps with anisotropic normalisation (alpha in [0, 1] controlling density correction from the normalised graph Laplacian to the Laplace-Beltrami operator), Von Neumann entropy-based diffusion time selection, and an optional landmark variant with Nystroem extension for larger datasets.
  • GPU acceleration (optional feature gpu): The most expensive part of UMAP and tSNE for large datasets is the nearest neighbour search. With the gpu feature enabled, kNN search runs on the GPU via cubecl, with backends for Vulkan, Metal, DirectX 12 (through wgpu) and CUDA. Graph construction and optimisation remain on CPU.
  • Multiple ANN backends via ann-search-rs:
    • Exhaustive - If you want precise results and have a small data set in which the approximate nearest neighbour index building is actually slower.
    • KmKnn - An exact nearest neighbour search algorithm, leveraging k-means clustering under the hood for speed. If you need exact results.
    • BallTree - A small, fast index for smaller data sets with lower dimensions.
    • Annoy (Approximate Nearest Neighbours Oh Yeah) - Good for medium low- dimensionality datasets.
    • NNDescent (Nearest Neighbour Descent) - good for larger datasets with higher dimensionality.
    • HNSW (Hierarchical Navigable Small World) - good for (very) larger datasets with higher dimensionality.
    • IVF (inversted file index) - Another fast and optimised index.
  • GPU ANN backends (with gpu feature):
    • Exhaustive GPU - Brute-force kNN on the GPU; deterministic and accurate.
    • IVF GPU - Inverted-file index on the GPU; good default for larger data.
    • NNDescent GPU - CAGRA-style graph construction on the GPU.
  • Distance metrics:
    • Euclidean
    • Cosine
    • Maybe more to come over time ... ?
  • Multiple initialisations:
    • Graph Laplacian eigenvector-based initialisation using Lanczos iteration
    • Random initialisation
    • PCA-based initialisation
  • Customisable parameters: Full control over fuzzy simplicial set construction, graph symmetrisation, and optimisation parameters for tSNE, UMAP and PHATE.
  • High performance: Parallel processing with Rayon, efficient sparse matrix operations, cache-friendly structures and optimised SGD and Adam optimisers for UMAP (for the latter also a parallelised version...) and fast optimisers for tSNE and also PHATE.
  • Synthetic datasets: Some synthetic datasets are available for testing and experimentation: Swiss role, clustered data and a trajectory-like structure.

Installation

Add this to your Cargo.toml:

[dependencies]
manifolds-rs = "*"

If you want to enable parametric UMAP, please use:

[dependencies]
manifolds-rs = { version = "*", features = [ "parametric" ] }

If you want to enable the FFT-accelerated version of tSNE, please use:

[dependencies]
manifolds-rs = { version = "*", features = [ "fft_tsne" ] }

If you want to enable GPU-accelerated kNN search, please use:

[dependencies]
manifolds-rs = { version = "*", features = [ "gpu" ] }

Feature flags can be combined, e.g. features = [ "gpu", "fft_tsne" ].

Notes

Please use version 0.1.3 and higher. These ones are not extensively tested against real data.

Note on GPU support: kNN runs on GPU through wgpu (Vulkan, Metal, DX12) or CUDA via cubecl. Computations on the GPU side are performed in f32; this is a limitation of WGSL (the wgpu shader language) which has no f64, and also reflects the fact that f64 throughput on consumer GPUs is typically 1/32 to 1/64 of f32. If you need double precision, stick to the CPU path. GPU results are not bit-reproducible across runs (parallel reductions do not have a fixed accumulation order), but structural quality is consistent.

Usage

R package

This crate powers manifoldsR, an R package leveraging the incredible speed that Rust offers.

UMAP Example

Below are examples of how to use UMAP.

use manifolds_rs::prelude::*;

// Generate synthetic clustered data
let (data, labels) = generate_clustered_data(
    1000,  // n_samples
    50,    // dimensionality
    5,     // n_clusters
    42,    // seed
);

// Configure UMAP parameters
let params = UmapParams::default_2d(
    Some(2),     // n_dim (output dimensions)
    Some(15),    // k (number of neighbours)
    Some(0.1),   // min_dist
    Some(1.0),   // spread
);

// Run UMAP
let embedding = umap(
    data.as_ref(),
    None,        // precomputed kNN (None = compute internally)
    &params,
    42,          // seed
    true,        // verbose
);

// embedding[0] contains x-coordinates
// embedding[1] contains y-coordinates

t-SNE Example

Below are examples of how to use t-SNE.

use manifolds_rs::prelude::*;

// Generate synthetic clustered data
let (data, labels) = generate_clustered_data(
    1000,  // n_samples
    50,    // dimensionality
    5,     // n_clusters
    42,    // seed
);

// Configure t-SNE parameters
let params = TsneParams::new(
    Some(2),      // n_dim (output dimensions)
    Some(30.0),   // perplexity
    Some(1e-4),   // init_range
    Some(200.0),  // learning_rate
    Some(1000),   // n_epochs
    None,         // ann_type (None = default "hnsw")
    Some(0.5),    // theta (Barnes-Hut angle)
    Some(3),      // n_interp_points (FFT interpolation grid points)
);

// Run t-SNE (Barnes-Hut)
let embedding = tsne(
    data.as_ref(),
    None,        // precomputed kNN (None = compute internally)
    &params,
    "bh",        // approximation type: "bh" or "fft" (requires fft_tsne feature)
    42,          // seed
    true,        // verbose
);

// embedding[0] contains x-coordinates
// embedding[1] contains y-coordinates

Using Precomputed k-NN

Both algorithms support precomputed k-nearest neighbour graphs for efficiency when running multiple embeddings:

use manifolds_rs::prelude::*;

let (data, _) = generate_clustered_data(500, 50, 5, 42);

// Compute k-NN once
let nn_params = NearestNeighbourParams::default();
let (knn_indices, knn_dist) = run_ann_search(
    data.as_ref(),
    15,              // k
    "hnsw".to_string(),
    &nn_params,
    42,              // seed
    true             // verbosity
);

// Use precomputed k-NN for UMAP
let params = UmapParams::default_2d(None, Some(15), None, None);
let embedding = umap(
    data.as_ref(),
    Some((knn_indices.clone(), knn_dist.clone())),
    &params,
    42,
    false,
);

GPU-Accelerated UMAP and tSNE (requires gpu feature)

With the gpu feature, the nearest neighbour search runs on the GPU while graph construction and optimisation stay on CPU. This is typically the bottleneck for larger datasets (say, 50k+ samples).

use manifolds_rs::prelude::*;
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};

// Generate synthetic clustered data (use f32 for GPU compatibility)
let (data, _) = generate_clustered_data(
    10_000,  // n_samples
    50,      // dimensionality
    5,       // n_clusters
    42,      // seed
);

// Configure GPU UMAP parameters
let params = UmapParamsGpu::default_2d(
    Some(2),     // n_dim
    Some(15),    // k
    Some(0.1),   // min_dist
    Some(1.0),   // spread
);

// Run GPU UMAP. ann_type defaults to "ivf_gpu"; alternatives are
// "exhaustive_gpu" (deterministic, brute force) and "nndescent_gpu".
let device = WgpuDevice::default();
let embedding = umap_gpu::<f32, WgpuRuntime>(
    data.as_ref(),
    None,        // precomputed kNN
    &params,
    device,
    42,          // seed
    true,        // verbose
);

GPU t-SNE works analogously:

use manifolds_rs::prelude::*;
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};

let (data, _) = generate_clustered_data(10_000, 50, 5, 42);

let params = TsneParamsGpu::new(
    Some(2),                           // n_dim
    Some(30.0),                        // perplexity
    Some(1e-4),                        // init_range
    Some(200.0),                       // learning rate
    Some(1000),                        // n_epochs
    Some("ivf_gpu".to_string()),       // ann_type
    Some(0.5),                         // theta
    Some(3),                           // n_interp_points
);

let device = WgpuDevice::default();
let embedding = tsne_gpu::<f32, WgpuRuntime>(
    data.as_ref(),
    None,
    &params,
    "bh",        // "bh" or "fft" (fft requires fft_tsne feature)
    device,
    42,
    true,
);

To use CUDA instead of wgpu, swap WgpuRuntime/WgpuDevice for CudaRuntime/CudaDevice from cubecl::cuda. On Linux CI, Vulkan via mesa-vulkan-drivers (lavapipe) is the simplest path for headless testing.

Parametric UMAP Example (requires parametric feature)

Parametric UMAP learns a neural network encoder that can transform new data points:

use manifolds_rs::prelude::*;
use burn::backend::ndarray::{NdArray, NdArrayDevice};
use burn::backend::Autodiff;

type Backend = Autodiff<NdArray<f64>>;

// Generate synthetic clustered data
let (data, labels) = generate_clustered_data(
    1000,  // n_samples
    50,    // dimensionality
    5,     // n_clusters
    42,    // seed
);

// Configure parametric UMAP
let fit_params = TrainParametricParams::from_min_dist_spread(
    0.1,       // min_dist
    1.0,       // spread
    0.0,       // correlation_weight
    None,      // negative_sample_rate
    Some(16),  // batch_size
    Some(100), // n_epochs
    None,      // learning_rate
);

let params = ParametricUmapParams::new(
    Some(2),              // n_dim (output dimensions)
    Some(15),             // n_neighbours
    Some("hnsw".into()),  // ann_type
    Some(vec![128, 64]),  // hidden_layers (neural network architecture)
    None,                 // nn_params
    None,                 // umap_graph_params
    Some(fit_params),     // training parameters
);

// Set up device
let device = NdArrayDevice::Cpu;

// Train parametric UMAP
let embedding = parametric_umap::<f64, Backend>(
    data.as_ref(),
    None,        // precomputed kNN (None = compute internally)
    &params,
    &device,
    42,          // seed
    true,        // verbose
);

// embedding[0] contains x-coordinates
// embedding[1] contains y-coordinates

PHATE Example

PHATE is well-suited for data with continuous structure and branching trajectories, such as single-cell differentiation data.

use manifolds_rs::prelude::*;

// Generate a synthetic branching trajectory
let branches = generate_example_branches(&TrajectoryTopology::DeepBifurcation);
let (data, branch_assignments) = generate_trajectory(
    1000,        // n_samples
    &branches,   // branch topology
    50,          // dimensionality
    0.5,         // noise
    42,          // seed
);

// Configure PHATE parameters
let params = PhateParams::new(
    Some(2),     // n_dim (output dimensions)
    Some(5),     // k (number of neighbours)
    None,        // ann_type (None = default "hnsw")
    None,        // decay (None = default 40.0)
    None,        // bandwidth_scale (None = default 1.0)
    None,        // graph_symmetry (None = default "average")
    None,        // t_max (None = auto)
    None,        // gamma (None = default 1.0)
    None,        // n_landmarks (None = full operator)
    None,        // landmark_method (None = default "spectral")
    None,        // n_svd
    None,        // t_custom
    None,        // mds_method (None = default "sgd_dense")
    None,        // mds_iter
    None,        // randomised (None = default true)
);

// Run PHATE
let embedding = phate(
    data.as_ref(),
    None,        // precomputed kNN (None = compute internally)
    params,      // note: consumed by value, not borrowed
    42,          // seed
    true,        // verbose
);

// embedding[0] contains x-coordinates
// embedding[1] contains y-coordinates

PaCMAP Example

PaCMAP preserves both local and global structure via three pair types (near, mid-near, and further pairs) and a phased optimisation schedule.

use manifolds_rs::prelude::*;

// Generate synthetic clustered data
let (data, labels) = generate_clustered_data(
    1000,  // n_samples
    50,    // dimensionality
    5,     // n_clusters
    42,    // seed
);

// Configure PaCMAP parameters
let params = PacmapParams::default();

// Run PaCMAP
let embedding = pacmap(
    data.as_ref(),
    None,    // precomputed kNN (None = compute internally)
    &params,
    42,      // seed
    true,    // verbose
);

// embedding[0] contains x-coordinates
// embedding[1] contains y-coordinates

Diffusion Maps Example

Classical diffusion maps embed the data via the top eigenvectors of a row-stochastic diffusion operator, optionally with anisotropic normalisation to correct for non-uniform sampling density. Setting alpha_norm = 1.0 recovers the Laplace-Beltrami operator; 0.5 gives the Fokker-Planck operator; 0.0 gives the unnormalised graph Laplacian.

use manifolds_rs::prelude::*;

// Generate synthetic clustered data
let (data, _) = generate_clustered_data(
    1000,  // n_samples
    50,    // dimensionality
    5,     // n_clusters
    42,    // seed
);

// Configure diffusion maps parameters
let params = DiffusionMapsParams::new(
    Some(2),     // n_dim (output dimensions)
    Some(5),     // k (number of neighbours)
    None,        // ann_type (None = default "kmknn")
    None,        // bandwidth_scale (None = default 1.0)
    None,        // thresh (None = default 1e-4)
    None,        // graph_symmetry (None = default "add")
    None,        // alpha_norm (None = default 1.0, Laplace-Beltrami)
    None,        // t_max (None = default)
    None,        // t_custom (None = auto-select via VNE knee)
    None,        // n_landmarks (None = full operator)
    None,        // landmark_method (None = default "spectral")
    None,        // n_svd
);

// Run diffusion maps
let embedding = diffusion_maps(
    data.as_ref(),
    None,        // precomputed kNN (None = compute internally)
    params,      // note: consumed by value, not borrowed
    42,          // seed
    true,        // verbose
);

// embedding[0] contains the first non-trivial diffusion component
// embedding[1] contains the second non-trivial diffusion component

Licence

MIT Licence

Copyright (c) 2025 Gregor Alexander Lueg

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.