Metadata-Version: 2.4
Name: mlx-vis
Version: 0.2.0
Summary: Pure MLX implementations of UMAP, t-SNE, PaCMAP, TriMap, DREAMS, and NNDescent for Apple Silicon. Metal GPU acceleration for both computation and video rendering. No scipy, no sklearn, no matplotlib.
Author-email: Han Xiao <han.xiao@jina.ai>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/hanxiao/mlx-vis
Project-URL: Repository, https://github.com/hanxiao/mlx-vis
Project-URL: Issues, https://github.com/hanxiao/mlx-vis/issues
Keywords: mlx,apple-silicon,umap,tsne,pacmap,nndescent,dimensionality-reduction,visualization,metal,gpu
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Visualization
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: mlx>=0.20.0
Requires-Dist: numpy>=1.24.0

# mlx-vis

Pure MLX implementations of UMAP, t-SNE, PaCMAP, TriMap, DREAMS, and NNDescent for Apple Silicon. Metal GPU acceleration for both computation and video rendering. No scipy, no sklearn, no matplotlib.

![Fashion-MNIST 70K on M3 Ultra. Top: dark theme, bottom: light theme. Left to right: UMAP, t-SNE, PaCMAP, TriMap, DREAMS.](comparison.png)

## Install

```bash
uv pip install mlx-vis
```

From source:

```bash
git clone https://github.com/hanxiao/mlx-vis.git
cd mlx-vis
uv pip install .
```

Requires `mlx >= 0.20.0` and `numpy >= 1.24.0`.

## Usage

```python
import numpy as np
from mlx_vis import UMAP, TSNE, PaCMAP, TriMap, DREAMS, NNDescent

X = np.random.randn(10000, 128).astype(np.float32)

# UMAP
Y = UMAP(n_components=2, n_neighbors=15).fit_transform(X)

# t-SNE
Y = TSNE(n_components=2, perplexity=30).fit_transform(X)

# PaCMAP
Y = PaCMAP(n_components=2, n_neighbors=10).fit_transform(X)

# TriMap
Y = TriMap(n_components=2, n_iters=400).fit_transform(X)

# DREAMS (t-SNE + PCA regularization)
Y = DREAMS(n_components=2, lam=0.15).fit_transform(X)

# NNDescent (approximate k-NN graph)
indices, distances = NNDescent(k=15).build(X)
```

Submodule imports also work:

```python
from mlx_vis.umap import UMAP
from mlx_vis.tsne import TSNE
from mlx_vis.pacmap import PaCMAP
from mlx_vis.trimap import TriMap
from mlx_vis.dreams import DREAMS
from mlx_vis.nndescent import NNDescent
```

## Methods

| Method | Class | Main API | Output |
|--------|-------|----------|--------|
| UMAP | `UMAP(n_components, n_neighbors, min_dist, ...)` | `fit_transform(X)` | `np.ndarray (n, d)` |
| t-SNE | `TSNE(n_components, perplexity, ...)` | `fit_transform(X)` | `np.ndarray (n, d)` |
| PaCMAP | `PaCMAP(n_components, n_neighbors, ...)` | `fit_transform(X)` | `np.ndarray (n, d)` |
| TriMap | `TriMap(n_components, n_iters, ...)` | `fit_transform(X)` | `np.ndarray (n, d)` |
| DREAMS | `DREAMS(n_components, lam, ...)` | `fit_transform(X)` | `np.ndarray (n, d)` |
| NNDescent | `NNDescent(k, n_iters, ...)` | `build(X)` | `(indices, distances)` |

## Visualization

All rendering runs on Metal GPU via MLX: coordinate mapping, circle-splatting, and color blending are fully vectorized MLX operations. Raw frames are piped to ffmpeg for PNG/video encoding. Zero matplotlib.

### Static plots

```python
from mlx_vis import UMAP, scatter_gpu
import numpy as np

X = np.random.randn(10000, 128).astype(np.float32)
labels = np.random.randint(0, 5, 10000)
Y = UMAP(n_components=2).fit_transform(X)

scatter_gpu(Y, labels=labels, theme="dark", save="plot.png")
```

### Animation

Video frames are rendered on GPU and piped to ffmpeg with `h264_videotoolbox` hardware encoding. **500 frames of 15K points in 1.5 seconds** on M3 Ultra.

**UMAP:**

https://github.com/user-attachments/assets/547b8ce2-17d4-4172-be59-ba83eafd1785

**t-SNE:**

https://github.com/user-attachments/assets/b8a4840b-7e71-4992-a26b-8332666af52a

**PaCMAP:**

https://github.com/user-attachments/assets/563c6a58-48e8-435d-b99d-5eafbb11be27

**TriMap:**

https://github.com/user-attachments/assets/eea51acf-b8da-4da1-9b23-6322bf300275

**DREAMS:**

https://github.com/user-attachments/assets/e6a3bb52-95f4-46d0-9a8d-f7714e85a18f

**Benchmark on Fashion-MNIST 70,000 x 784, M3 Ultra:**

| | UMAP | t-SNE | PaCMAP | TriMap | DREAMS |
|---|---|---|---|---|---|
| Iterations | 500 | 500 | 450 | 500 | 500 |
| Embedding | 3.5s | 3.9s | 2.4s | 2.8s | 4.0s |
| GPU render (800 frames) | 2.1s | 1.9s | 1.8s | 1.9s | 1.9s |
| Total | 5.6s | 5.8s | 4.2s | 4.7s | 5.9s |

```python
from mlx_vis import UMAP, animate_gpu
import numpy as np, time

X = np.random.randn(10000, 128).astype(np.float32)
labels = np.random.randint(0, 5, 10000)

snaps, times = [], []
t0 = time.time()
def cb(epoch, Y_np):
    snaps.append(Y_np.copy())
    times.append(time.time() - t0)

Y = UMAP(n_components=2, n_epochs=200).fit_transform(X, epoch_callback=cb)

animate_gpu(snaps, labels=labels, timestamps=times,
            method_name="umap-mlx", fps=120, theme="dark",
            save="animation.mp4")
```

Full Fashion-MNIST example:

```bash
python -m mlx_vis.examples.fashion_mnist --method umap --theme dark
python -m mlx_vis.examples.fashion_mnist --method all
```

## License

Apache-2.0
