trimap 1.2.0

News!

[Aug 18, 2026] A GPU-parallel PyTorch implementation is now available as trimap.TorchTRIMAP, with support for CUDA, MPS, and CPU.

[Mar 16, 2022] An example colab using TriMap JAX implementation is now available at https://github.com/eamid/examples/blob/master/TriMap.ipynb. We analyze the results on S-curve, MNIST, Fashion MNIST, etc. using t-SNE, UMAP, TriMap, and PCA.

[Feb 17, 2022] A JAX implementation is now available at https://github.com/google-research/google-research/tree/master/trimap. More updates are coming soon!

How to use TriMap

TriMap has a transformer API similar to other sklearn libraries. To use TriMap with the default parameters, simply do:

import trimap
from sklearn.datasets import load_digits

digits = load_digits()

embedding = trimap.TRIMAP().fit_transform(digits.data)

To find the embedding using a precomputed pairwise distance matrix D, pass D as input and set use_dist_matrix to True:

embedding = trimap.TRIMAP(use_dist_matrix=True).fit_transform(D)

You can also pass the precomputed k-nearest neighbors and their corresponding distances as a tuple (knn_nbrs, knn_distances). Note that the rows must be in order, starting from point 0 to n-1. This feature also requires X to compute the embedding

embedding = trimap.TRIMAP(knn_tuple=(knn_nbrs, knn_distances)).fit_transform(X)

To calculate the global score, do:

gs = trimap.TRIMAP(verbose=False).global_score(digits.data, embedding)
print("global score %2.2f" % gs)

Parameters

The list of parameters is given blow:

• n_dims: Number of dimensions of the embedding (default = 2)

• n_inliers: Number of nearest neighbors for forming the nearest neighbor triplets (default = 12).

• n_outliers: Number of outliers for forming the nearest neighbor triplets (default = 4).

• n_random: Number of random triplets per point (default = 3).

• distance: Distance measure (‘euclidean’ (default), ‘manhattan’, ‘angular’ (or ‘cosine’), ‘hamming’)

• weight_temp: Temperature of the logarithm applied to the weights. Larger temperatures generate more compact embeddings. weight_temp=0. corresponds to no transformation (default=0.5).

• weight_adj (deprecated): The value of gamma for the log-transformation (default = 500.0).

• lr: Learning rate (default = 0.1).

• n_iters: Number of iterations (default = 400).

The other parameters include:

• knn_tuple: Use the precomputed nearest-neighbors information in form of a tuple (knn_nbrs, knn_distances) (default = None)

• use_dist_matrix: Use the precomputed pairwise distance matrix (default = False)

• apply_pca: Reduce the number of dimensions of the data to 100 if necessary before applying the nearest-neighbor search (default = True).

• opt_method: Optimization method {‘sd’ (steepest descent), ‘momentum’ (GD with momentum), ‘dbd’ (delta-bar-delta, default)}.

• verbose: Print the progress report (default = False).

• return_seq: Store the intermediate results and return the results in a tensor (default = False).

An example of adjusting these parameters:

import trimap
from sklearn.datasets import load_digits

digits = load_digits()

embedding = trimap.TRIMAP(n_inliers=20,
                          n_outliers=10,
                          n_random=10).fit_transform(digits.data)

The nearest-neighbor calculation in the legacy implementation is performed using ANNOY.

GPU-parallel PyTorch version

TorchTRIMAP is a tensor-native implementation of the complete pipeline. It keeps preprocessing, nearest-neighbor results, triplets, weights, the loss, gradient, optimizer state, and embedding on the selected accelerator. The native nearest-neighbor search is exact and blocked, so it does not allocate an n x n distance matrix.

A single import trimap exposes both estimator APIs:

import trimap

legacy_model = trimap.TRIMAP()
torch_model = trimap.TorchTRIMAP(device="cuda")

trimap.TRIMAP is the legacy NumPy/Numba implementation, while trimap.TorchTRIMAP runs the tensor-native pipeline. Both provide fit_transform. A NumPy array can be passed directly to the PyTorch version; it is moved to the selected device before preprocessing:

import numpy as np
import trimap
from sklearn.datasets import load_digits

X = np.asarray(load_digits().data, dtype=np.float32)
model = trimap.TorchTRIMAP(
    device="cuda",
    knn_backend="auto",
    n_iters=400,
    random_state=42,
    triplet_batch_size=1_000_000,
)

embedding = model.fit_transform(X)
embedding_numpy = embedding.detach().cpu().numpy()

The returned embedding is a PyTorch tensor on the input/selected device. Use embedding.detach().cpu().numpy() only when a NumPy consumer needs it. The pipeline runs preprocessing/PCA, nearest-neighbor search, triplet and weight construction, initialization, and gradient-descent optimization on that device.

For a large CUDA data set such as Covertype, after removing its target column:

model = trimap.TorchTRIMAP(
    device="cuda",
    knn_backend="cuvs-cagra",
    gradient="explicit",
    n_iters=400,
    random_state=42,
    triplet_batch_size=30_000_000,
)
embedding = model.fit_transform(X[:, :54])

python benchmarks/benchmark_legacy_parity.py --device mps --n 10000

Examples

The following are some of the results on real-world datasets. The values of nearest-neighbor accuracy and global score are shown as a pair (NN, GS) on top of each figure. For more results, please refer to our paper.

USPS Handwritten Digits (n = 11,000, d = 256)

20 News Groups (n = 18,846, d = 100)

Tabula Muris (n = 53,760, d = 23,433)

MNIST Handwritten Digits (n = 70,000, d = 784)

Fashion MNIST (n = 70,000, d = 784)

TV News (n = 129,685, d = 100)

Runtime of t-SNE, LargeVis, UMAP, and TriMap in the hh:mm:ss format on a single machine with 2.6 GHz Intel Core i5 CPU and 16 GB of memory is given in the following table. We limit the runtime of each method to 12 hours. Also, UMAP runs out of memory on datasets larger than ~4M points.

Installing

Requirements:

• numpy

• scikit-learn

• numba

• annoy

Installing annoy

If you are having trouble with installing annoy on macOS using the command:

pip3 install annoy

you can alternatively try:

pip3 install git+https://github.com/sutao/annoy.git@master

Install Options

If you have all the requirements installed, you can use pip:

sudo pip install trimap

Please regularly check for updates and make sure you are using the most recent version. If you have TriMap installed and would like to upgrade to the newer version, you can use the command:

sudo pip install --upgrade --force-reinstall trimap

An alternative is to install the dependencies manually using anaconda and using pip to install TriMap:

conda install numpy
conda install scikit-learn
conda install numba
conda install annoy
pip install trimap

For a manual install get this package:

wget https://github.com/eamid/trimap/archive/master.zip
unzip master.zip
rm master.zip
cd trimap-master

Install the requirements

sudo pip install -r requirements.txt

or

conda install scikit-learn numba annoy

Install the package

python setup.py install

Support and Contribution

This implementation is still a work in progress. Any comments/suggestions/bug-reports are highly appreciated. Please feel free contact me at: eamid@ucsc.edu. If you would like to contribute to the code, please fork the project and send me a pull request.

Citation

If you use TriMap in your publications, please cite our current reference on arXiv:

@article{2019TRIMAP,
     author = {{Amid}, Ehsan and {Warmuth}, Manfred K.},
     title = "{TriMap: Large-scale Dimensionality Reduction Using Triplets}",
     journal = {arXiv preprint arXiv:1910.00204},
     archivePrefix = "arXiv",
     eprint = {1910.00204},
     year = 2019,
}

License

Please see the LICENSE file.