fast-umap

UMAP (Uniform Manifold Approximation and Projection) in Rust

This repository contains a Rust implementation of UMAP (Uniform Manifold Approximation and Projection), a dimensionality reduction algorithm that preserves both the local and global structure of data. UMAP is widely used for visualizing high-dimensional data in 2D or 3D space.

This implementation leverages the burn machine learning framework, which provides automatic differentiation and GPU support, allowing you to train and apply UMAP models on high-dimensional datasets efficiently.

Features

Dimensionality Reduction: Reduces high-dimensional data to a lower-dimensional space (e.g., 2D or 3D) for visualization or further analysis.
Customizable UMAP Model: The model architecture can be configured with different numbers of input features, hidden layer sizes, and output dimensions.
GPU Support: Powered by the burn framework with support for training on CPU and GPU using the wgpu backend.
Flexible Data Handling: Functions for converting between vectors and tensors, generating synthetic data, and more.

Roadmap

Add k-neighbors element and add precalculated KNN tree to do faster calculations on high dimensional data.
Add MNIST dataset example
Add different distance functions, not just euclidean
Create testbench to compare different hyper parameters (patience vs n_features vs epochs etc.)

Installation

Clone the repository:

git clone https://github.com/eugenehp/fast-umap.git
cd fast-umap

Build the project:
```
cargo build --release
```

Example Usage

1. Fitting a UMAP Model

You can fit a UMAP model to your data using the UMAP::fit function. Here's how to do it:

use burn::backend::Autodiff;
use burn::backend::wgpu::{Wgpu, WgpuDevice};
use fast_umap::prelude::*;

fn main() {
    // Example data (list of samples, each with a list of features)
    let data: Vec<Vec<f64>> = vec![
        vec![1.0, 2.0, 3.0],
        vec![4.0, 5.0, 6.0],
        vec![7.0, 8.0, 9.0],
        // Add more samples...
    ];

    // Fit the UMAP model
    let model = umap(data);

    // You can now use the model to transform new data
    let transformed = model.transform(data);

    // Print the transformed data (low-dimensional representation)
    for sample in transformed {
        println!("{:?}", sample);
    }
}

2. Transforming Data

Once the UMAP model is trained, you can transform new high-dimensional data into its low-dimensional representation:

let transformed_data = model.transform(new_data);

This function will take new_data in the form of Vec<Vec<f64>> and return its 2D or 3D representation, suitable for visualization.

3. Generating Test Data

You can generate synthetic test data to experiment with the UMAP model using the generate_test_data function:

let data = generate_test_data(100, 50); // 100 samples, each with 50 features

4. Visualizing Data

After transforming the data to a 2D or 3D space, you can use external charting libraries (e.g., plotters in Rust or matplotlib in Python) to visualize the results.

Model Configuration

The UMAP model configuration is customizable through the UMAPModelConfigBuilder. You can set the following parameters:

input_size: Number of input features (i.e., the dimensionality of the data).
hidden_size: The number of neurons in the hidden layers.
output_size: The target number of dimensions (typically 2 or 3 for visualization).

Here's how to configure and build the model:

let model_config = UMAPModelConfigBuilder::default()
    .input_size(50)      // Input features: 50 dimensions
    .hidden_size(100)    // Hidden layer size: 100 neurons
    .output_size(2)      // Output size: 2 (for 2D visualization)
    .build()
    .unwrap();

Training the UMAP Model

You can train the UMAP model on your dataset by calling the fit method on the UMAP struct. The training process will optimize the model's weights to reduce the data's dimensionality.

Training configuration parameters include:

epochs: The number of epochs to train the model.
batch_size: The number of samples per training batch.
learning_rate: The learning rate for gradient descent.
beta1, beta2: Optimizer hyperparameters for Adam optimization.

For example:

let model = UMAP::<Autodiff<Wgpu>>::fit(data, WgpuDevice::default());

Examples

Simple

cargo run --example simple

Sample code:

use fast_umap::prelude::*;
use rand::Rng;

fn main() {
    // Number of samples in the dataset
    let num_samples = 100;

    // Number of features (dimensions) for each sample
    let num_features = 3;

    // Create a random number generator for generating random values
    let mut rng = rand::thread_rng();

    // Generate a dataset of random values with `num_samples` rows and `num_features` columns
    let data: Vec<Vec<f64>> = (0..num_samples * num_features)
        .map(|_| rng.gen::<f64>()) // Random number generation for each feature
        .collect::<Vec<f64>>() // Collect all random values into a vector
        .chunks_exact(num_features) // Chunk the vector into rows of length `num_features`
        .map(|chunk| chunk.to_vec()) // Convert each chunk into a Vec<f64>
        .collect(); // Collect the rows into a vector of vectors

    // Fit the UMAP model to the data and reduce the data to a lower-dimensional space (default: 2D)
    let umap = umap(data.clone());

    // Transform the data using the trained UMAP model to reduce its dimensions
    let reduced_dimensions_vector = umap.transform(data.clone());

    // Visualize the reduced dimensions as a vector
    chart_vector(reduced_dimensions_vector, None);

    // Optionally, you can also visualize the reduced dimensions as a tensor
    // let reduced_dimensions_tensor = umap.transform_to_tensor(data.clone());
    // print_tensor_with_title("reduced_dimensions", &reduced_dimensions_tensor);
    // chart_tensor(reduced_dimensions_tensor, None);
}

Generates this plot:

plot

Advanced

cargo run --example advanced