Cluster X into K clusters using FAISS IVF. Returns ------- centroids : np.ndarray, shape (K, dim), float32 cluster_ids : np.ndarray, shape (n,), uint32
(
X: np.ndarray,
K: int,
metric_str: str = "l2",
num_threads: int = 0,
)
| 38 | # ────────────────────────────────────────────── |
| 39 | |
| 40 | def cluster_data( |
| 41 | X: np.ndarray, |
| 42 | K: int, |
| 43 | metric_str: str = "l2", |
| 44 | num_threads: int = 0, |
| 45 | ) -> tuple[np.ndarray, np.ndarray]: |
| 46 | """ |
| 47 | Cluster X into K clusters using FAISS IVF. |
| 48 | |
| 49 | Returns |
| 50 | ------- |
| 51 | centroids : np.ndarray, shape (K, dim), float32 |
| 52 | cluster_ids : np.ndarray, shape (n,), uint32 |
| 53 | """ |
| 54 | if num_threads > 0: |
| 55 | faiss.omp_set_num_threads(num_threads) |
| 56 | |
| 57 | dim = X.shape[1] |
| 58 | |
| 59 | if metric_str == "ip": |
| 60 | metric = faiss.METRIC_INNER_PRODUCT |
| 61 | print("Clustering metric: InnerProduct") |
| 62 | else: |
| 63 | metric = faiss.METRIC_L2 |
| 64 | print("Clustering metric: L2") |
| 65 | |
| 66 | index = faiss.index_factory(dim, f"IVF{K},Flat", metric) |
| 67 | index.verbose = True |
| 68 | |
| 69 | t0 = time() |
| 70 | index.train(X) |
| 71 | print(f"IVF training time: {time() - t0:.2f}s") |
| 72 | |
| 73 | centroids = index.quantizer.reconstruct_n(0, index.nlist) # (K, dim) float32 |
| 74 | _, cluster_ids_2d = index.quantizer.search(X, 1) # (n, 1) int64 |
| 75 | cluster_ids = cluster_ids_2d.flatten().astype(np.uint32) # (n,) uint32 |
| 76 | |
| 77 | return centroids, cluster_ids |