Metadata-Version: 2.4
Name: umap-metal
Version: 0.1.0
Summary: UMAP accelerated by Apple MLX on Mac MPS (Metal Performance Shaders)
License-Expression: MIT
Project-URL: Homepage, https://github.com/takaho/umap_metal
Project-URL: Bug Tracker, https://github.com/takaho/umap_metal/issues
Keywords: umap,dimensionality-reduction,mlx,apple-silicon,mps,metal,machine-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Visualization
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: mlx>=0.20
Requires-Dist: numpy>=1.24
Requires-Dist: scipy>=1.10
Requires-Dist: scikit-learn>=1.2
Requires-Dist: numba>=0.57
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: matplotlib; extra == "dev"
Requires-Dist: umap-learn; extra == "dev"
Requires-Dist: pynndescent; extra == "dev"

# umap-metal

**UMAP accelerated by Apple MLX on Mac MPS (Metal Performance Shaders)**

`umap-metal` is a drop-in replacement for [`umap-learn`](https://github.com/lmcinnes/umap) that runs all major compute phases on Apple Silicon GPUs via [MLX](https://github.com/ml-explore/mlx). It is designed to be interface-compatible with `umap-learn.UMAP` so existing code requires only a one-line import change.

---

## Requirements

| | |
|---|---|
| Hardware | Mac with Apple Silicon (M1 / M2 / M3 / M4 family) |
| macOS | 12.3 Ventura or later |
| Python | 3.10+ |

---

## Installation

```bash
pip install umap-metal
```

> **Note** `mlx` is an Apple-only package and will only install on macOS with Apple Silicon.

---

## Quick Start

The API mirrors `umap-learn`:

```python
from umap_metal import UMAP
import numpy as np

X = np.random.standard_normal((10_000, 50)).astype("float32")

reducer = UMAP(n_neighbors=15, n_components=2, random_state=42)
embedding = reducer.fit_transform(X)   # shape (10_000, 2)
```

### Drop-in replacement

```python
# Before
from umap import UMAP

# After — nothing else changes
from umap_metal import UMAP
```

### scikit-learn Pipeline compatibility

`UMAP` extends `sklearn.base.BaseEstimator` and `TransformerMixin`:

```python
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from umap_metal import UMAP

pipe = Pipeline([
    ("scaler", StandardScaler()),
    ("umap",   UMAP(n_neighbors=15, n_components=2)),
])
embedding = pipe.fit_transform(X)
```

---

## Parameters

| Parameter | Default | Description |
|---|---|---|
| `n_neighbors` | `15` | Size of local neighborhood |
| `n_components` | `2` | Embedding dimensionality |
| `metric` | `"euclidean"` | `"euclidean"` or `"cosine"` |
| `n_epochs` | `None` | SGD epochs (auto: 200 for large, 500 for small) |
| `learning_rate` | `1.0` | Initial learning rate |
| `init` | `"spectral"` | `"spectral"` or `"random"` |
| `min_dist` | `0.1` | Minimum distance between embedded points |
| `spread` | `1.0` | Scale of the embedding |
| `n_negative_samples` | `5` | Negative samples per positive edge |
| `random_state` | `None` | `int`, `RandomState`, or `None` |
| `verbose` | `False` | Print per-step timing |
| `chunk_size` | `4096` | Chunk size for exact MLX kNN (small n) |
| `force_exact_knn` | `False` | Force O(n²) exact kNN regardless of n |
| `use_mlx_layout` | `True` | Use MLX GPU layout optimization |
| `nndescent_iters` | `5` | NN-Descent refinement iterations (large n) |
| `nndescent_trees` | `10` | RP-trees for NN-Descent initialization (large n) |

---

## How It Works

`umap-metal` accelerates all three compute-heavy phases of UMAP on Apple MPS:

### Phase 1 — k-Nearest Neighbor Search

| n | Method | Complexity |
|---|---|---|
| ≤ 15,000 | **MLX exact GEMM** — chunked `‖xᵢ − xⱼ‖²` via `xᵢxⱼᵀ` | O(n²·d) |
| > 15,000 | **MLX NN-Descent** — RP-Forest init + neighbor-of-neighbor refinement | O(n·k²·d·iters) |

The MLX NN-Descent builds binary random-projection trees (vectorized with `np.einsum` + `np.bincount`) to generate high-recall initial candidates, then refines with forward and reverse neighbor propagation. No Numba JIT is used, eliminating the 4–6 s cold-start penalty of `pynndescent`.

### Phase 2 — Fuzzy Graph Construction

Smooth kNN distances (σ, ρ per point via binary search) and membership strength computation remain on CPU with Numba JIT.

### Phase 3 — Layout Optimization (SGD)

A **node-centric** reformulation replaces the edge-centric SGD scatter pattern:

```
# Edge-centric (requires scatter_add — GPU-unfriendly)
for edge (i, j):
    force_i += attract(i, j)
    force_j -= attract(i, j)

# Node-centric (pure gather + reduce — GPU-optimal)
neighbors = emb[knn_indices[i, :]]          # (N, k, D) gather
attr_force = reduce(attract(i, neighbors))   # (N, D) no scatter needed
rep_force  = reduce(repel(i, random_neg))    # (N, D)
emb[i]    += lr * (attr_force + rep_force)
```

This achieves a consistent **4–4.5× speedup** on the layout step alone, and requires no atomic operations or scatter kernels.

---

## Benchmark

Tested on Apple M2 Max (96 GB unified memory), `n_neighbors=15`, `dim=50`, `epochs=200`.  
`umap-learn` timings are after one Numba warm-up call (subsequent-call performance).

```
       n    umap-metal    umap-learn    speedup    kNN method
─────────────────────────────────────────────────────────────
   5,000        0.61 s        5.48 s      9.0×    MLX exact
  10,000        1.32 s        2.95 s      2.2×    MLX exact
  20,000        2.75 s        5.94 s      2.2×    MLX NN-Descent
  50,000        6.37 s       15.33 s      2.4×    MLX NN-Descent
 100,000       12.80 s       32.42 s      2.5×    MLX NN-Descent
```

Step-level breakdown for `n = 50,000`:

```
  kNN (MLX NN-Descent)   4.33 s   68%
  Layout (MLX SGD)       1.42 s   22%
  Graph construction     0.45 s    7%
  Spectral init          0.17 s    3%
  ─────────────────────────────────
  Total                  6.37 s
```

> The 9× speedup at n=5,000 reflects umap-learn's first-call Numba compilation cost.  
> From n=10,000 onward, comparisons are against warm (post-compilation) umap-learn.

### kNN quality (recall@15 vs exact brute-force)

| Method | n=5,000 | n=10,000 |
|---|---|---|
| MLX exact | 100% | 100% |
| MLX NN-Descent | 77% | 68% |
| pynndescent (umap-learn) | 81% | 73% |

---

## Verbose Output

```python
reducer = UMAP(n_neighbors=15, n_epochs=200, verbose=True)
reducer.fit_transform(X)
```

```
[umap-metal] n=50000  k=15  epochs=200  kNN=MLX NN-Descent
  kNN (MLX NN-Descent): 4.33s
  Graph construction: 0.45s
  Spectral init: 0.17s
  a=1.5769  b=0.8951  layout=MLX
  epoch 0/200  lr=1.0000
  ...
  Layout optimization: 1.42s
```

---

## Project Structure

```
umap_metal/
├── __init__.py          # exposes UMAP
├── umap_metal.py        # UMAP class (sklearn-compatible)
├── knn.py               # dispatch: MLX exact or MLX NN-Descent
├── knn_nndescent.py     # RP-Forest init + NN-Descent (MLX, no Numba)
├── graph.py             # fuzzy simplicial set construction (Numba)
├── spectral.py          # spectral embedding initialization (scipy)
├── layout.py            # SGD layout, Numba CPU fallback
└── layout_mlx.py        # node-centric SGD layout on MLX GPU
```

---

## Limitations

- **macOS / Apple Silicon only** — MLX does not run on Linux or Windows.
- `transform()` (out-of-sample projection) is not yet implemented.
- Approximate kNN (n > 15,000) achieves ~70–77% recall; embeddings are visually equivalent to `umap-learn` but may differ in fine structure.
- Very high-dimensional data (d > 200) may reduce the GPU efficiency advantage.

---

## License

MIT

---

## Citation

If you use this software, please also cite the original UMAP paper:

> McInnes, L., Healy, J., & Melville, J. (2018).  
> *UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction.*  
> arXiv:1802.03426.
