https://github.com/motiwari/BanditPAM
BanditPAM C++ implementation and Python package
https://github.com/motiwari/BanditPAM
clustering machine-learning python
Last synced: about 1 month ago
JSON representation
BanditPAM C++ implementation and Python package
- Host: GitHub
- URL: https://github.com/motiwari/BanditPAM
- Owner: motiwari
- License: mit
- Created: 2020-10-22T01:01:33.000Z (over 5 years ago)
- Default Branch: main
- Last Pushed: 2025-08-25T15:48:28.000Z (7 months ago)
- Last Synced: 2026-02-15T06:02:50.533Z (about 2 months ago)
- Topics: clustering, machine-learning, python
- Language: C++
- Homepage:
- Size: 73.9 MB
- Stars: 657
- Watchers: 8
- Forks: 48
- Open Issues: 71
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# BanditPAM: Almost Linear-Time $k$-Medoids Clustering
[](https://github.com/motiwari/BanditPAM/actions/workflows/linux_run_tests.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/linux_build_wheels.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/mac_arm64_run_tests.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/mac_arm64_cmake_build.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/mac_intel_run_tests.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/mac_intel_cmake_build.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/mac_build_wheels.yml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/pages/pages-build-deployment)
[](https://github.com/motiwari/BanditPAM/actions/workflows/R-CMD-check.yaml)
[](https://github.com/motiwari/BanditPAM/actions/workflows/run_style_checks.yml)
This repo contains a high-performance implementation of BanditPAM from [BanditPAM: Almost Linear-Time k-Medoids Clustering](https://proceedings.neurips.cc/paper/2020/file/73b817090081cef1bca77232f4532c5d-Paper.pdf) and [BanditPAM++: Faster k-medoids Clustering](https://proceedings.neurips.cc/paper_files/paper/2023/file/e885e5bc6e13b9dd8f80bc5482b1fa2f-Paper-Conference.pdf). The code can be called directly from Python, R, or C++.
If you use this software, please cite:
Mo Tiwari, Martin Jinye Zhang, James Mayclin, Sebastian Thrun, Chris Piech, Ilan Shomorony. "BanditPAM: Almost Linear Time *k*-medoids Clustering via Multi-Armed Bandits" Advances in Neural Information Processing Systems (NeurIPS) 2020.
Mo Tiwari, Ryan Kang*, Donghyun Lee*, Sebastian Thrun, Chris Piech, Ilan Shomorony, Martin Jinye Zhang. "BanditPAM++: Faster k-medoids Clustering" Advances in Neural Information Processing Systems (NeurIPS) 2023.
```python
@inproceedings{tiwari2020banditpam,
title={BanditPAM: Almost Linear Time $k$-medoids Clustering via Multi-Armed Bandits},
author={Tiwari, Mo and Zhang, Martin J and Mayclin, James and Thrun, Sebastian and Piech, Chris and Shomorony, Ilan},
booktitle={Advances in Neural Information Processing Systems},
pages={368--374},
year={2020}
}
@inproceedings{tiwari2023banditpam++,
title={BanditPAM++: Faster $k$-medoids Clustering},
author={Tiwari, Mo and Kang, Ryan and Lee, Donghyun and Thrun, Sebastian and Shomorony, Ilan and Zhang, Martin J},
journal={Advances in Neural Information Processing Systems},
volume={36},
pages={73371--73382},
year={2023}
}
```
# Requirements
# TL;DR run `python -m pip install banditpam` or `install.packages(banditpam)` and jump to the [examples](https://github.com/motiwari/BanditPAM#example-1-synthetic-data-from-a-gaussian-mixture-model).
If you have any difficulties, please see the [platform-specific guides](https://github.com/motiwari/BanditPAM#platform-specific-installation-guides) and file a Github issue if you have additional trouble.
## Further Reading
* [NeurIPS 2020 Paper](https://proceedings.neurips.cc/paper/2020/file/73b817090081cef1bca77232f4532c5d-Paper.pdf)
* [NeurIPS 2023 Paper](https://proceedings.neurips.cc/paper_files/paper/2023/file/e885e5bc6e13b9dd8f80bc5482b1fa2f-Paper-Conference.pdf)
* [3-minute summary video](https://slideslive.com/38936275/banditpam-almost-linear-time-kmedoids-clustering-via-multiarmed-bandits)
* [Blog post](https://ai.stanford.edu/blog/banditpam/)
* [Code](https://github.com/motiwari/BanditPAM)
* [PyPI](https://pypi.org/project/banditpam/)
* [Documentation](https://banditpam.readthedocs.io/en/latest)
# Python Quickstart
## Install the repo and its dependencies:
This can be done either through PyPI (recommended)
```python
/BanditPAM/: python -m pip install -r requirements.txt
/BanditPAM/: python -m pip install banditpam
```
OR through the source code via
```python
/BanditPAM/: git submodule update --init --recursive
/BanditPAM/: cd headers/carma
/BanditPAM/: mkdir build && cd build && cmake -DCARMA_INSTALL_LIB=ON .. && sudo cmake --build . --config Release --target install
/BanditPAM/: cd ../../..
/BanditPAM/: python -m pip install -r requirements.txt
/BanditPAM/: sudo python -m pip install .
```
### Example 1: Synthetic data from a Gaussian Mixture Model
```python
from banditpam import KMedoids
import numpy as np
import matplotlib.pyplot as plt
# Generate data from a Gaussian Mixture Model with the given means:
np.random.seed(0)
n_per_cluster = 40
means = np.array([[0,0], [-5,5], [5,5]])
X = np.vstack([np.random.randn(n_per_cluster, 2) + mu for mu in means])
# Fit the data with BanditPAM:
kmed = KMedoids(n_medoids=3, algorithm="BanditPAM")
kmed.fit(X, 'L2')
print(kmed.average_loss) # prints 1.2482391595840454
print(kmed.labels) # prints cluster assignments [0] * 40 + [1] * 40 + [2] * 40
# Visualize the data and the medoids:
for p_idx, point in enumerate(X):
if p_idx in map(int, kmed.medoids):
plt.scatter(X[p_idx, 0], X[p_idx, 1], color='red', s = 40)
else:
plt.scatter(X[p_idx, 0], X[p_idx, 1], color='blue', s = 10)
plt.show()
```
### Example 2: MNIST and its medoids visualized via t-SNE
```python
# Start in the repository root directory, i.e. '/BanditPAM/'.
from banditpam import KMedoids
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
# Load the 1000-point subset of MNIST and calculate its t-SNE embeddings for visualization:
X = pd.read_csv('data/MNIST_1k.csv', sep=' ', header=None).to_numpy()
X_tsne = TSNE(n_components=2).fit_transform(X)
# Fit the data with BanditPAM:
kmed = KMedoids(n_medoids=10, algorithm="BanditPAM")
kmed.fit(X, 'L2')
# Visualize the data and the medoids via t-SNE:
for p_idx, point in enumerate(X):
if p_idx in map(int, kmed.medoids):
plt.scatter(X_tsne[p_idx, 0], X_tsne[p_idx, 1], color='red', s = 40)
else:
plt.scatter(X_tsne[p_idx, 0], X_tsne[p_idx, 1], color='blue', s = 5)
plt.show()
```
# R Examples
Please see [here](https://github.com/motiwari/BanditPAM/tree/main/R_package/banditpam).
## Documentation
Documentation for BanditPAM can be found on [read the docs](https://banditpam.readthedocs.io/).
## Building the C++ executable from source
Please note that it is *NOT* necessary to build the C++ executable from source to use the Python code above. However, if you would like to use the C++ executable directly, follow the instructions below.
### Option 1: Building with Docker
We highly recommend building using Docker. One can download and install Docker by following instructions at the [Docker install page](https://docs.docker.com/get-docker/). Once you have Docker installed and the Docker Daemon is running, run the following commands:
```
/BanditPAM/scripts/docker$ chmod +x env_setup.sh
/BanditPAM/scripts/docker$ ./env_setup.sh
/BanditPAM/scripts/docker$ ./run_docker.sh
```
which will start a Docker instance with the necessary dependencies. Then:
```
/BanditPAM$ mkdir build && cd build
/BanditPAM/build$ cmake .. && make
```
This will create an executable named `BanditPAM` in `BanditPAM/build/src`.
### Option 2: Installing requirements and building directly
Building this repository requires four external requirements:
* [CMake](https://cmake.org/download/) >= 3.17
* [Armadillo](http://arma.sourceforge.net/download.html) >= 10.5.3
* [OpenMP](https://www.openmp.org/resources/openmp-compilers-tools/) >= 2.5 (OpenMP is supported by default on most Linux platforms, and can be downloaded through [homebrew](https://brew.sh/) on MacOS)
* [CARMA](https://github.com/RUrlus/carma) >= 0.6.2
If installing these requirements from source, one can generally use the following procedure to install each requirement from the library's root folder (with `armadillo` used as an example here):
```
/armadillo$ mkdir build && cd build
/armadillo/build$ cmake .. && make && sudo make install
```
Note that `CARMA` has different installation instructions; see [its instructions](https://github.com/RUrlus/carma#installation).
#### Platform-specific installation guides
Further installation information for [MacOS](docs/install_mac.md), [Linux](docs/install_linux.md), and [Windows](docs/install_windows.md) is available in the [docs folder](docs). Ensure all the requirements above are installed and then run:
```
/BanditPAM$ mkdir build && cd build
/BanditPAM/build$ cmake .. && make
```
This will create an executable named `BanditPAM` in `BanditPAM/build/src`.
## C++ Usage
Once the executable has been built, it can be invoked with:
```
/BanditPAM/build/src/BanditPAM -f [path/to/input.csv] -k [number of clusters]
```
* `-f` is mandatory and specifies the path to the dataset
* `-k` is mandatory and specifies the number of clusters with which to fit the data
For example, if you ran `./env_setup.sh` and downloaded the MNIST dataset, you could run:
```
/BanditPAM/build/src/BanditPAM -f ../data/MNIST_1k.csv -k 10
```
The expected output in the command line will be:
```
Medoids: 694,168,306,714,324,959,527,251,800,737
```
## Implementing a custom distance metric
One of the advantages of $k$-medoids is that it works with arbitrary distance metrics; in fact, your "metric" need not even be a real metric -- it can be negative, asymmetric, and/or not satisfy the triangle inequality or homogeneity. Any pairwise dissimilarity function works with $k$-medoids.
This also allows for clustering of "exotic" objects like trees, graphs, natural language, and more -- settings where running $k$-means wouldn't even make sense. We talk about one such setting in the [full paper](https://proceedings.neurips.cc/paper/2020/file/73b817090081cef1bca77232f4532c5d-Paper.pdf).
The package currently supports a number of distance metrics, including all $L_p$ losses and cosine distance.
If you're willing to write a little C++, you only need to add a few lines to [kmedoids_algorithm.cpp](https://github.com/motiwari/BanditPAM/blob/main/src/kmedoids_algorithm.cpp#L560-L615) and [kmedoids_algorithm.hpp](https://github.com/motiwari/BanditPAM/blob/main/headers/kmedoids_algorithm.hpp#L136-L142) to implement your distance metric / pairwise dissimilarity!
Then, be sure to re-install the repository with a `python -m pip install .` (note the trailing `.`).
The maintainers of this repository are working on permitting arbitrary dissimilarity metrics that users write in Python, as well; see [#4](https://github.com/motiwari/BanditPAM/issues/4).
## Testing
To run the full suite of tests, run in the root directory:
```
/BanditPAM$ python -m unittest discover -s tests
```
Alternatively, to run a "smaller" set of tests, from the main repo folder run `python tests/test_smaller.py` or `python tests/test_larger.py` to run a set of longer, more intensive tests.
## Credits
Mo Tiwari wrote the original Python implementation of BanditPAM and many features of the C++ implementation. Mo and Adarsh Kumarappan now maintain the implementations.
James Mayclin developed the initial C++ implementation of BanditPAM.
The original BanditPAM paper was published by Mo Tiwari, Martin Jinye Zhang, James Mayclin, Sebastian Thrun, Chris Piech, and Ilan Shomorony.
We would like to thank Jerry Quinn, David Durst, Geet Sethi, and Max Horton for helpful guidance regarding the C++ implementation.