diff --git a/.gitignore b/.gitignore index f6da783b..9c933fcd 100644 --- a/.gitignore +++ b/.gitignore @@ -73,6 +73,8 @@ cluster-mol-ids*.json config*.json timings*.json max-rss.txt +*.smi +!chembl-33-natural-products-subset.smi # Pickle *.pkl diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index 6088e1f4..4affb3df 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -1482,9 +1482,12 @@ def global_clustering( if not self.is_init: raise ValueError("The model has not been fitted yet.") centroids = np.vstack(self.get_centroids(packed=False)) - labels = self._centrals_global_clustering( + predictor = self._global_clustering_predictor( centroids, n_clusters, method=method, input_is_packed=False, **method_kwargs ) + # Add 1 to start labels from 1 instead of 0, so 0 can be used as sentinel + # value + labels = predictor.fit_predict(centroids) + 1 num_centroids = len(centroids) self._n_global_clusters = ( n_clusters if num_centroids > n_clusters else num_centroids @@ -1493,7 +1496,7 @@ def global_clustering( return self @staticmethod - def _centrals_global_clustering( + def _global_clustering_predictor( centrals: NDArray[np.uint8], n_clusters: int, *, @@ -1502,7 +1505,7 @@ def _centrals_global_clustering( n_features: int | None = None, # TODO: Type correctly **method_kwargs: tp.Any, - ) -> NDArray[np.int64]: + ) -> tp.Any: r""":meta private:""" if method not in {"agglomerative", "kmeans", "kmeans-normalized"}: raise ValueError(f"Unknown method {method}") @@ -1532,11 +1535,7 @@ def _centrals_global_clustering( else: raise ValueError("method must be one of 'kmeans' or 'agglomerative'") - # Add 1 to start labels from 1 instead of 0, so 0 can be used as sentinel - # value - # This is the bottleneck for building this index - # K-means is feasible, agglomerative is extremely expensive - return predictor.fit_predict(centrals) + 1 + return predictor # There are 4 cases here: diff --git a/bblean/cli.py b/bblean/cli.py index f33e5962..082903ea 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -410,6 +410,10 @@ def _plot_pops( bool, Option("--show/--no-show", hidden=True), ] = True, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""Population plot of the clustering results""" from bblean._console import get_console @@ -434,6 +438,7 @@ def _plot_pops( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -527,6 +532,10 @@ def _plot_umap( rich_help_panel="Advanced", ), ] = None, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""UMAP visualization of the clustering results""" from bblean._console import get_console @@ -559,6 +568,7 @@ def _plot_umap( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -627,6 +637,10 @@ def _plot_pca( str | None, Option("--filename"), ] = None, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""PCA visualization of the clustering results""" from bblean._console import get_console @@ -651,6 +665,7 @@ def _plot_pca( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -782,6 +797,10 @@ def _plot_tsne( bool, Option("--show/--no-show", hidden=True), ] = True, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""t-SNE visualization of the clustering results""" from bblean._console import get_console @@ -818,6 +837,7 @@ def _plot_tsne( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -852,7 +872,7 @@ def _plot_summary( Path | None, Option( "-s", - "--smiles-path", + "--smiles", show_default=False, help="Optional smiles path, if passed a scaffold analysis is performed", ), @@ -902,6 +922,10 @@ def _plot_summary( bool, Option("--show/--no-show", hidden=True), ] = True, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""Summary plot of the clustering results""" from bblean._console import get_console @@ -928,6 +952,7 @@ def _plot_summary( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -994,6 +1019,15 @@ def _run( help="Dir to dump the output files", ), ] = None, + smiles_path: Annotated[ + Path | None, + Option( + "-s", + "--smiles", + show_default=False, + help="Optional smiles path, to store smiles with results", + ), + ] = None, overwrite: Annotated[bool, Option(help="Allow overwriting output files")] = False, branching_factor: Annotated[ int, @@ -1129,6 +1163,18 @@ def _run( bool, Option("-v/-V", "--verbose/--no-verbose"), ] = True, + build_idx: Annotated[ + bool, + Option("--build-idx/--no-build-idx"), + ] = False, + idx_method: Annotated[ + str, + Option("--idx-method", hidden=True), + ] = "kmeans", + idx_n_clusters: Annotated[ + int | None, + Option("--n-clusters", hidden=True), + ] = None, ) -> None: r"""Run standard, serial BitBIRCH clustering over `*.npy` fingerprint files""" # TODO: Remove code duplication with multiround @@ -1261,12 +1307,31 @@ def _run( # Symlink or copy fingerprint files input_fps_dir = (out_dir / "input-fps").resolve() input_fps_dir.mkdir() + + input_smiles_dir = (out_dir / "input-smiles").resolve() + smiles_files = [] + if smiles_path: + smiles_files = ( + [smiles_path] + if not smiles_path.is_dir() + else sorted(smiles_path.glob("*.smi")) + ) + if smiles_files: + input_smiles_dir.mkdir() if copy_inputs: for file in input_files: shutil.copy(file, input_fps_dir / file.name) + for file in smiles_files: + shutil.copy(file, input_smiles_dir / file.name) else: for file in input_files: (input_fps_dir / file.name).symlink_to(file.resolve()) + for file in smiles_files: + (input_smiles_dir / file.name).symlink_to(file.resolve()) + + # Build the index with defaults + if build_idx: + _build_idx(out_dir, method=idx_method, n_clusters=idx_n_clusters) # TODO: Currently sometimes after a round is triggered *more* files are output, since @@ -1465,6 +1530,18 @@ def _multiround( bool, Option("--cleanup/--no-cleanup", hidden=True), ] = True, + build_idx: Annotated[ + bool, + Option("--build-idx/--no-build-idx"), + ] = False, + idx_method: Annotated[ + str, + Option("--idx-method", hidden=True), + ] = "kmeans", + idx_n_clusters: Annotated[ + int | None, + Option("--n-clusters", hidden=True), + ] = None, ) -> None: r"""Run multi-round BitBIRCH clustering, optionally parallelize over `*.npy` files""" # noqa:E501 from bblean._console import get_console @@ -1556,6 +1633,10 @@ def _multiround( for file in input_files: (input_fps_dir / file.name).symlink_to(file.resolve()) + # Build the index with defaults + if build_idx: + _build_idx(out_dir, method=idx_method, n_clusters=idx_n_clusters) + @app.command("fps-info", rich_help_panel="Fingerprints") def _fps_info( @@ -1942,6 +2023,208 @@ def _split_fps( ) +@app.command("query-idx", hidden=True) +def _query_idx( + idx_path: Annotated[ + Path, + Argument(help="Path to the index dir, or its parent dir with index files"), + ], + query: Annotated[ + str, + Option("-s", "--smiles"), + ] = "", + k: Annotated[ + int, + Option("-k", "--num"), + ] = 10, + threshold: Annotated[ + float, + Option("-t", "--threshold"), + ] = 0.0, + num_probe: Annotated[ + int, + Option("-p", "--probe"), + ] = 1, + n_features: Annotated[ + int, + Option("-n", "--n-features"), + ] = DEFAULTS.n_features, + fp_kind: Annotated[ + str, + Option("-f", "--fp-kind"), + ] = DEFAULTS.fp_kind, +) -> None: + from bblean._console import get_console + from bblean.fingerprints import fps_from_smiles + from bblean.ivf import IVFIndex + + console = get_console() + index = IVFIndex.from_dir(idx_path) + query_fp = fps_from_smiles([query], kind=fp_kind, n_features=n_features, pack=True)[ + 0 + ] + results_list = index.search(query_fp, n_probe=num_probe, threshold=threshold, k=k) + + for r in results_list: + console.print(r) + + +def _get_centroids_and_clusters_paths(clusters_path: Path) -> tuple[Path, Path]: + if clusters_path.is_dir(): + centroids_path = clusters_path / "cluster-centroids-packed.pkl" + clusters_path = clusters_path / "clusters.pkl" + # TODO: If this file doesn't exist, it will have to be created on-the-fly + if not centroids_path.exists(): + raise ValueError( + "Centroids must be saved for building IVF or global clustering." + " This limitation may be lifted in the future" + ) + else: + raise ValueError( + "clusters path must be a dir for building IVF or global clustering." + " This limitation may be lifted in the future" + ) + return clusters_path, centroids_path + + +@app.command("build-idx", hidden=True) +def _build_idx( + clusters_path: Annotated[ + Path, + Argument(help="Path to the clusters file, or a dir with a clusters.pkl file"), + ], + fps_path: Annotated[ + Path | None, + Option( + "-f", + "--fps-path", + help="Path to fingerprint file, or directory with fingerprint files", + show_default=False, + ), + ] = None, + smiles_path: Annotated[ + Path | None, + Option( + "-s", + "--smiles", + show_default=False, + help="Optional smiles path, if used smiles are returned in search results", + ), + ] = None, + method: Annotated[ + str, + Option("--method"), + ] = "kmeans", + n_clusters: Annotated[ + int | None, + Option("--n-clusters"), + ] = None, + direct_reassignment: Annotated[ + bool, + Option("-d/-D", "--direct-reassignment/--no-direct-reassignment"), + ] = False, + verbose: Annotated[ + bool, + Option("--verbose/--no-verbose"), + ] = True, +) -> None: + import numpy as np + from bblean.utils import _has_files_or_valid_symlinks + from bblean.ivf import IVFIndex + + from bblean._console import get_console + + console = get_console(silent=not verbose) + clusters_path, centroids_path = _get_centroids_and_clusters_paths(clusters_path) + with open(clusters_path, mode="rb") as f: + cluster_members = pickle.load(f) + + with open(centroids_path, mode="rb") as f: + centroids_packed = np.vstack(pickle.load(f)) + + inferred_fps_path = clusters_path.parent / "input-fps" + symlink_fps = False + if fps_path is None: + fps_path = inferred_fps_path + elif inferred_fps_path.is_dir(): + raise ValueError("Fingerprints are already present in cluster dir") + else: + symlink_fps = True + + if fps_path.is_dir() and _has_files_or_valid_symlinks(fps_path): + fps_files = sorted(fps_path.glob("*.npy")) + elif fps_path.is_file(): + fps_files = [fps_path] + fps_path = fps_path.parent + else: + raise ValueError("Fingerprints are required to build the index") + + if len(fps_files) > 1: + with console.status( + "[italic]Merging fingerprint files for IVF index...[/italic]", + spinner="dots", + ): + fps_path = fps_path.parent / "merged-fps" + _merge_fps(fps_path.parent / "input-fps", fps_path) + fps_files = sorted(fps_path.glob("*.npy")) + assert len(fps_files) == 1 + symlink_fps = False + + inferred_smiles_path = clusters_path.parent / "input-smiles" + if smiles_path is not None: + if inferred_smiles_path.is_dir(): + raise ValueError("Smiles already present in cluster dir") + if smiles_path.is_dir(): + smiles_files = sorted(smiles_path.glob("*.smi")) + else: + smiles_files = [smiles_path] + smiles_path = smiles_path.parent + else: + if not ( + inferred_smiles_path.is_dir() + and _has_files_or_valid_symlinks(inferred_smiles_path) + ): + msg = "Smiles won't be returned when searching index, please use --smiles" + warnings.warn(msg) + + fps = np.load(fps_files[0]) # packed + kwargs = {"random_state": 42} if method.startswith("kmeans") else {} + + with console.status("[italic]Building IVF index...[/italic]", spinner="dots"): + # Build with no smiles, since we don't need to search anyting yet + index = IVFIndex.from_bitbirch_clusters( + cluster_members, + centroids_packed, + fps, + (), + method, + n_clusters, + input_is_packed=True, + sort=True, + direct_reassignment=direct_reassignment, + **kwargs, + ) + + with console.status("[italic]Saving global clusters...[/italic]", spinner="dots"): + global_cluster_medoids_path = ( + clusters_path.parent / "global-cluster-medoids-packed.npy" + ) + np.save( + global_cluster_medoids_path, + index._medoids_packed, + ) + global_clusters_path = clusters_path.parent / "global-clusters.pkl" + with open(global_clusters_path, mode="wb") as f: + pickle.dump(index._members, f) + if symlink_fps: + for file in fps_files: + (fps_path / file.name).symlink_to(file.resolve()) + if smiles_path is not None: + for file in smiles_files: + (smiles_path / file.name).symlink_to(file.resolve()) + console.print("Successfully built IVF index") + + @app.command("fps-shuffle", rich_help_panel="Fingerprints") def _shuffle_fps( in_path: Annotated[ diff --git a/bblean/ivf.py b/bblean/ivf.py new file mode 100644 index 00000000..043eb9b1 --- /dev/null +++ b/bblean/ivf.py @@ -0,0 +1,217 @@ +r"""IVF (Inverted File) search index implementation using BitBIRCH clustering. + +IVF is efficient search index for chemical fingerprints that uses BitBIRCH to partition +the cluster space. Searches in this space use approximate nearest neighbors (ANN). +""" + +import pickle +from pathlib import Path +import typing_extensions as tpx +import dataclasses +import math +import typing as tp +import numpy as np +from numpy.typing import NDArray + +from bblean.bitbirch import BitBirch +from bblean.smiles import load_smiles +from bblean.similarity import jt_sim_packed +from bblean.fingerprints import pack_fingerprints, unpack_fingerprints + + +@dataclasses.dataclass +class SearchResult: + index: int + similarity: float + smi: str | None + + def __repr__(self) -> str: + sim = self.similarity + if sim > 1e-4: + sim_str = f"{sim:.4f}" + elif sim > 0: + sim_str = f"{sim:.4e}" + elif sim == 0: + sim_str = "0" + else: + raise RuntimeError("Negative similarity found") + out = f"SearchResult(index={self.index}, similarity={sim_str}" + if self.smi is not None: + return f"{out}, smi='{self.smi}')" + return f"{out})" + + +class IVFIndex: + r""" + Inverted File (IVF) index for efficient similarity search of chemical fingerprints. + + The index uses BitBIRCH clustering to partition fingerprints into clusters, + then at query time, only the most relevant clusters are searched, providing + a significant speedup over exhaustive search. + """ + + def __init__( + self, + medoids_packed: NDArray[np.uint8], + members: tp.Sequence[list[int]], + fps: NDArray[np.uint8], + smiles: tp.Sequence[str] | NDArray[np.str_] = (), + input_is_packed: bool = True, + n_features: int | None = None, + ): + # Build directly from global clusters + self._medoids_packed = medoids_packed + self._members = list(members) + fps = fps.astype(np.uint8, copy=False) + if not input_is_packed: + fps = pack_fingerprints(fps) + self._packed_fps = fps + self._smiles = np.asarray(smiles, dtype=np.str_) + + @classmethod + def from_dir(cls, idx_path: Path) -> tpx.Self: + global_cluster_medoids_path = idx_path / "global-cluster-medoids-packed.npy" + global_clusters_path = idx_path / "global-clusters.pkl" + if (idx_path / "merged-fps").is_dir(): + fps_paths = sorted((idx_path / "merged-fps").glob("*.npy")) + else: + fps_paths = sorted((idx_path / "input-fps").glob("*.npy")) + if len(fps_paths) > 1: + raise ValueError("Currently only a single fp file is supported") + fps_path = fps_paths[0] + smiles_files = sorted((idx_path / "input-smiles").glob("*.smi")) + with open(global_clusters_path, "rb") as f: + members = pickle.load(f) + medoids_packed = np.load(global_cluster_medoids_path) + fps = np.load(fps_path) + return cls(medoids_packed, members, fps, load_smiles(smiles_files)) + + @classmethod + def from_bitbirch_clusters( + cls, + members: tp.Sequence[list[int]], + centrals: NDArray[np.uint8], + fps: NDArray[np.uint8], + smiles: tp.Sequence[str] | NDArray[np.str_] = (), + method: str = "kmeans", + n_clusters: int | None = None, + input_is_packed: bool = True, + n_features: int | None = None, + sort: bool = True, + direct_reassignment: bool = False, + **method_kwargs: tp.Any, + ) -> tpx.Self: + """Build the IVF index from bitbirch clusters""" + n_samples = fps.shape[0] + if n_clusters is None: + n_clusters = max(int(math.sqrt(n_samples)), 1) + if n_clusters is not None and n_clusters <= 0: + raise ValueError("n_clusters must be a positive integer or None") + + fps = fps.astype(np.uint8, copy=False) + if input_is_packed: + fps = unpack_fingerprints(fps, n_features) + centrals = unpack_fingerprints(centrals, n_features) + + predictor = BitBirch._global_clustering_predictor( + centrals, n_clusters, method=method, **method_kwargs + ) + # This is the bottleneck for building this index + # K-means is feasible, agglomerative is extremely expensive + # In this case there is no need to add 1 either in either case + predictor.fit(centrals) + if direct_reassignment: + # Direct reassignment reassigns the fingerprints directly using the + # predictor instead of indirectly reassigning using the central labels + if method.endswith("-normalized"): + labels = predictor.predict( + fps / np.linalg.norm(fps, axis=1, keepdims=True) + ) + else: + labels = predictor.predict(fps) + mol_ids = [(labels == i).nonzero()[0].tolist() for i in range(n_clusters)] + else: + labels = predictor.predict(centrals) + num_centrals = len(centrals) + n_clusters = n_clusters if num_centrals > n_clusters else num_centrals + mol_ids = BitBirch._new_ids_from_labels(members, labels, n_clusters) + + if sort: + mol_ids.sort(key=lambda x: len(x), reverse=True) + _, medoids = BitBirch._unpacked_medoids_from_members( + fps, mol_ids, input_is_packed=False + ) + return cls(pack_fingerprints(medoids), mol_ids, pack_fingerprints(fps), smiles) + + def _find_candidate_idxs( + self, query_fp_packed: NDArray[np.uint8], n_probe: int + ) -> NDArray[np.int64]: + """ + Find the n_probe nearest clusters to the query fingerprint. + + Args: + query_fp: Query fingerprint (numpy array or RDKit ExplicitBitVect) + n_probe: Number of clusters to return + + Returns: + list of cluster IDs, sorted by similarity to query + """ + similarities = jt_sim_packed(self._medoids_packed, query_fp_packed) + # Get indices of top n_probe most similar medoids, limiting to avail clusters + n_probe = min(n_probe, len(self._medoids_packed)) + top_indices = np.argsort(similarities)[-n_probe:] + candidates = [] + for idx in top_indices: + candidates.extend(self._members[idx]) + return np.array(candidates) + + def search( + self, + query_fp: NDArray[np.uint8], + k: int = 10, + n_probe: int = 1, + threshold: float = 0.0, + input_is_packed: bool = True, + n_features: int | None = None, + ) -> list[SearchResult]: + """ + Search for the k most similar fingerprints to the query. + + Args: + query_fp: Query fingerprint (numpy array or RDKit ExplicitBitVect) + k: Number of results to return + n_probe: Number of clusters to search + threshold: Minimum similarity threshold (0.0 means no threshold) + n_features: provided for API consistency only, does nothing. + + Returns: + list of SearchResult, in sorted order, each with: + - index: Index of the fingerprint + - similarity: Tanimoto similarity to query + - smi: SMILES string (if available, else None) + """ + if k <= 0: + raise ValueError("k must be > 0") + if n_probe <= 0: + raise ValueError("n_probe must be > 0") + + if not input_is_packed: + query_fp = pack_fingerprints(query_fp) + + candidates = self._find_candidate_idxs(query_fp, n_probe) + similarities = jt_sim_packed(self._packed_fps[candidates], query_fp) + + # Apply threshold filter + if threshold > 0.0: + is_selected = similarities >= threshold + similarities = similarities[is_selected] + candidates = candidates[is_selected] + + # Sort by similarity (descending) and prepare results + sorted_indices = np.argsort(similarities)[::-1][:k] + results = [] + for idx in sorted_indices: + fp_idx = candidates[idx].item() + smi = self._smiles[fp_idx].strip() if self._smiles.size > 0 else None + results.append(SearchResult(fp_idx, similarities[idx].item(), smi)) + return results diff --git a/bblean/plotting.py b/bblean/plotting.py index 18381c38..a8703015 100644 --- a/bblean/plotting.py +++ b/bblean/plotting.py @@ -442,9 +442,12 @@ def _dispatch_visualization( verbose: bool = True, save: bool = True, show: bool = True, + use_global: bool = False, ) -> None: if clusters_path.is_dir(): - clusters_path = clusters_path / "clusters.pkl" + clusters_path = clusters_path / ( + "global-clusters.pkl" if use_global else "clusters.pkl" + ) with open(clusters_path, mode="rb") as f: clusters = pickle.load(f) if fps_path is None: diff --git a/tests/legacy_fns.py b/tests/legacy_fns.py new file mode 100644 index 00000000..f5f2509c --- /dev/null +++ b/tests/legacy_fns.py @@ -0,0 +1,24 @@ +# type: ignore +import numpy as np + + +def calculate_comp_sim(data): + """Returns vector of complementary similarities""" + n_objects = len(data) - 1 + c_total = np.sum(data, axis=0) + comp_matrix = c_total - data + a = comp_matrix * (comp_matrix - 1) / 2 + comp_sims = np.sum(a, axis=1) / np.sum( + (a + comp_matrix * (n_objects - comp_matrix)), axis=1 + ) + return comp_sims + + +def calculate_medoid(data): + """Returns index of medoid""" + return data[np.argmin(calculate_comp_sim(data))] + + +def calculate_medoid_idx(data): + """Returns index of medoid""" + return np.argmin(calculate_comp_sim(data)) diff --git a/tests/test_similarity.py b/tests/test_similarity.py index 9e3144c2..0578b719 100644 --- a/tests/test_similarity.py +++ b/tests/test_similarity.py @@ -3,6 +3,7 @@ import pytest from inline_snapshot import snapshot +from legacy_fns import calculate_medoid # type: ignore # TODO: Fix the tests with pytest-subtests so that both the _py_similarity and the # _cpp_similarity are tested independently @@ -294,3 +295,4 @@ def test_jt_isim_medoid() -> None: idx, m = bblean.similarity.jt_isim_medoid(fps) assert idx == snapshot(26) assert m.tolist() == snapshot([1, 1, 0, 1, 1, 1, 1, 1]) + assert calculate_medoid(fps).tolist() == [1, 1, 0, 1, 1, 1, 1, 1]