Chelombus

Billion-scale molecular clustering and visualization on commodity hardware.

Chelombus enables interactive exploration of ultra-large chemical datasets (up to billions of molecules) using Product Quantization and nested TMAPs. Process the entire Enamine REAL database (9.6B molecules) on a single workstation.

Live Demo: https://chelombus.gdb.tools

Overview

Chelombus implements the "Nested TMAP" framework for visualizing billion-sized molecular datasets:

SMILES → MQN Fingerprints → PQ Encoding → PQk-means Clustering → Nested TMAPs

Key Features:

• Scalability: Stream billions of molecules without loading everything into memory
• Efficiency: Compress 42-dimensional MQN vectors to 6-byte PQ codes (28x compression)
• GPU acceleration: Optional CUDA support for PQ encoding and cluster assignment (~25x speedup)
• Visualization: Navigate from global overview to individual molecules in two clicks
• Accessibility: Runs on commodity hardware (tested: AMD Ryzen 7, 64GB RAM)

Installation

From PyPI (recommended)

pip install chelombus

From Source

git clone https://github.com/afloresep/chelombus.git
cd chelombus
pip install -e .

Platform Notes

Apple Silicon (M1/M2/M3): The pqkmeans library is not currently supported on Apple Silicon Macs. My plan is to rewrite pqkmeans with Silicon and GPU support but that's for a future release... For now, clustering functionality requires an x86_64 system.

Windows: Not supported natively. pqkmeans has no Windows wheels (building it needs MSVC and CMake), and the GPU kernels need Triton, which is Linux only. Use WSL2 instead.

Scripts on macOS (and Windows): FingerprintFromSmiles runs in parallel worker processes, which re-import the calling script on these systems. Put a script's code under if __name__ == "__main__":. Notebooks are not affected.

GPU Acceleration

Every step in the pipeline supports optional GPU acceleration via the device parameter: encoder training, PQ encoding, cluster training, and label assignment. When a CUDA GPU is available, device='auto' (the default) uses the GPU transparently; otherwise it falls back to CPU.

Requirements: torch and triton (on Linux, pip install torch installs both).

encoder = PQEncoder(k=256, m=6, iterations=20)
encoder.fit(training_fps, device='auto')          # GPU KMeans fitting
pq_codes = encoder.transform(fingerprints)        # GPU batch assignment

clusterer = PQKMeans(encoder, k=100000)
clusterer.fit(pq_codes)                           # GPU Triton assign + CPU centroid update
labels = clusterer.predict(pq_codes)              # GPU Triton kernel

# Or force a specific device
labels_cpu = clusterer.predict(pq_codes, device='cpu')
labels_gpu = clusterer.predict(pq_codes, device='gpu')

GPU benchmarks on 1B Enamine REAL molecules (RTX 4070 Ti SUPER 16GB, K=100,000):

Stage	Description	Time
Encoder training	Train codebook on 50M MQN fingerprints	1.8 min
PQ encoding	Encode 1B MQN fingerprints → 1B PQ codes (6-dim)	3.7 min
Cluster training	Train PQKMeans on 1B PQ codes (5 iters, tol=0)	2.3 hrs
Label assignment	Assign 1B PQ codes to 100K clusters	26.2 min
Total		2.9 hrs

Cluster training with the default tol=1e-3 converges earlier (changed centers dropped from 11.8% → 0.2% over 5 iterations), reducing training time further. Extrapolated to 10B molecules: ~1.2 days.

To reproduce (requires 1B MQN fingerprints as .npy chunks):

python scripts/benchmark_1B_pipeline.py --chunks /path/to/chunks --n 1000000000

The GPU implementation uses a custom Triton kernel for cluster assignment that tiles over centers with an online argmin, never materializing the N x K distance matrix. The kernel supports any number of subvectors M via compile-time unrolling (tl.static_range). VRAM usage is ~(M+4) bytes/point, so even an 8 GB GPU can process hundreds of millions of points per batch.

Quick Start

from chelombus import DataStreamer, FingerprintCalculator, PQEncoder, PQKMeans

# 1. Stream SMILES in chunks
streamer = DataStreamer()
smiles_gen = streamer.parse_input('molecules.smi', chunksize=100000)

# 2. Calculate MQN fingerprints
fp_calc = FingerprintCalculator()
for smiles_chunk in smiles_gen:
    fingerprints = fp_calc.FingerprintFromSmiles(smiles_chunk, fp='mqn')
    # Save fingerprints...

# 3. Train PQ encoder on sample
encoder = PQEncoder(k=256, m=6, iterations=20)
encoder.fit(training_fingerprints)

# 4. Transform all fingerprints to PQ codes
pq_codes = encoder.transform(fingerprints)

# 5. Cluster with PQk-means
clusterer = PQKMeans(encoder, k=100000)
labels = clusterer.fit_predict(pq_codes)

Project Structure

chelombus/
├── chelombus/
│   ├── encoder/          # Product Quantization encoder
│   ├── clustering/       # PQk-means wrapper
│   ├── streamer/         # Memory-efficient data streaming
│   └── utils/            # Fingerprints, visualization, helpers
├── scripts/              # Pipeline scripts
│   └── pipeline/         # ← end-to-end example: SMILES → clusters → TMAP website
├── examples/             # Tutorial notebooks
└── tests/                # Unit tests

Want to run the whole thing end-to-end? See scripts/pipeline/. It is a self-contained, numbered, copy-paste example that takes a plain SMILES file all the way to a browsable nested-TMAP website, with a one-command driver (run_all.sh) and a plain-English README.

Choosing k (Number of Clusters)

The scripts/select_k.py script sweeps over k values on a subsample to help pick the right number of clusters. It supports checkpointing, if interrupted, rerun the same command and it resumes from where it left off.

python scripts/select_k.py \
    --pq-codes data/pq_codes.npy \
    --encoder models/encoder.joblib \
    --n-subsample 10000000 \
    --k-values 10000 25000 50000 100000 200000 \
    --iterations 10 \
    --output results/k_selection.csv \
    --plot results/k_selection.png

Results on 100M Enamine REAL molecules (RTX 4070 Ti SUPER 16GB, AMD Ryzen 7 64GB RAM):

k	Avg Distance	Empty Clusters	Median Cluster Size	Fit Time (GPU)	Fit Time (CPU)
10,000	3.67	7.1%	9,061	1.7 min	1.3 h
25,000	2.74	13.5%	3,680	3.3 min	3.1 h
50,000	2.17	19.5%	1,879	6.0 min	6.2 h
100,000	1.69	26.7%	960	9.1 min	12.6 h
200,000	1.30	34.7%	492	17.6 min	26.4 h

Guidelines:

• k = 50,000 is a good default — under 20% empty clusters, median size ~1,900, and the avg distance improvement starts plateauing beyond this point.
• k = 100,000 if you need tighter clusters and can tolerate ~27% empty clusters.
• Beyond 200K, over a third of clusters are empty — diminishing returns.

Documentation

• Full docs: https://chelombus.gdb.tools
• Tutorial: See examples/tutorial.ipynb for a hands-on introduction
• Large-scale example: See examples/enamine_1B_clustering.ipynb
• API Reference: See docs/api.md or the hosted docs

Testing

# Run all tests
pytest tests/

# Run specific test file
pytest tests/test_encoder.py -v

Citation

If you use Chelombus in your research, please cite:

@article{flores2026nested,
  title={Nested TMAPs to Visualize Billions of Molecules},
  author={Flores Sepulveda, Alejandro and Reymond, Jean-Louis},
  journal={Journal of Chemical Information and Modeling},
  volume={66},
  number={10},
  pages={5595--5601},
  year={2026},
  publisher={ACS Publications}
}

Contributing

Contributions are welcome! Please:

1. Fork the repository
2. Create a feature branch: git checkout -b feature/my-feature
3. Write tests for new functionality
4. Submit a pull request

License

MIT License. See LICENSE for details.

Acknowledgments

• PQk-means by Matsui et al.
• TMAP by Probst & Reymond
• RDKit for cheminformatics functionality
• Swiss National Science Foundation (grant no. 200020_178998)