From 43cf4642d7dd04d8207dc96f8cb7ad693134a387 Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 27 Oct 2025 21:09:01 -0400 Subject: [PATCH 01/23] Add initial ivf index implementation --- bblean/ivf_index.py | 256 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 bblean/ivf_index.py diff --git a/bblean/ivf_index.py b/bblean/ivf_index.py new file mode 100644 index 00000000..d75acf3f --- /dev/null +++ b/bblean/ivf_index.py @@ -0,0 +1,256 @@ +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). +""" + +from collections import defaultdict +from numpy.typing import NDArray +import math +import time +import typing as tp +import numpy as np + +from bblean.bitbirch import BitBirch +from bblean.similarity import jt_sim_packed +from bblean.fingerprints import pack_fingerprints + + +class IVFIndex: + """ + 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. + + Attributes: + n_clusters (int): Number of clusters to use. If None, uses sqrt(n_samples) + threshold (float): Similarity threshold for BitBIRCH clustering + branching_factor (int): Branching factor for BitBIRCH clustering + cluster_centroids (np.ndarray): Centroids of each cluster + cluster_members (Dict[int, list[int]]): Mapping of cluster IDs to member + fingerprint indices + fingerprints (np.ndarray): Stored fingerprints for similarity search + smiles (list[str]): Optional SMILES strings corresponding to fingerprints + built (bool): Whether the index has been built + """ + + def __init__( + self, + n_clusters: int | None = None, + threshold: float = 0.7, + branching_factor: int = 50, + ): + """ + Initialize the IVF index. + + Args: + n_clusters: Number of clusters to use (required) + threshold: Similarity threshold for BitBIRCH clustering (used only for tree + building) + branching_factor: Branching factor for BitBIRCH clustering + """ + if n_clusters is not None and n_clusters <= 0: + raise ValueError("n_clusters must be a positive integer or None") + + self.n_clusters = n_clusters + self.threshold = threshold + self.branching_factor = branching_factor + + # Will be populated during build_index + self.cluster_centroids_packed = np.array([], dtype=np.uint8) # packed + self.cluster_members: dict[int, list[int]] = {} + self.fingerprints = np.array([], dtype=np.uint8) + self.smiles: list[str] = [] + self.is_built = False + + def build_index( + self, + fingerprints: NDArray[np.uint8], + smiles: tp.Sequence[str] = (), + method: str = "kmeans-normalized", + input_is_packed: bool = False, + n_features: int | None = None, + verbose: bool = False, + **method_kwargs: tp.Any, + ) -> None: + """ + Build the IVF index by clustering fingerprints using BitBIRCH. + + Args: + fingerprints: Binary fingerprints of shape (n_samples, n_features) + smiles: Optional list of SMILES strings corresponding to fingerprints + """ + n_samples = fingerprints.shape[0] + if self.n_clusters is None: + n_clusters = max(int(math.sqrt(n_samples)), 1) + else: + n_clusters = self.n_clusters + + fingerprints = fingerprints.astype(np.uint8, copy=False) + if not input_is_packed: + fingerprints = pack_fingerprints(fingerprints) + + # Store (packed) fingerprints and smiles for later use + self.fingerprints = fingerprints + self.smiles = list(smiles) + + # Always use k-clusters functionality since n_clusters is required + if verbose: + print(f"Clustering {n_samples} fps into exactly {n_clusters} clusters...") + + # Initialize BitBIRCH for clustering + birch = BitBirch( + threshold=self.threshold, branching_factor=self.branching_factor + ) + birch.fit(fingerprints) + birch.global_clustering(method=method, n_clusters=n_clusters, **method_kwargs) + + # Fetch new cluster centroids and members + bf_labels = birch._global_clustering_centroid_labels + unique_clusters = np.unique(birch._global_clustering_centroid_labels) + if verbose: + print(f"Found {len(unique_clusters)} unique clusters") + root = birch._root + assert root is not None # mypy + cluster_ls = np.zeros((len(unique_clusters), root.n_features), dtype=np.uint64) + cluster_samples = np.zeros(len(unique_clusters), dtype=np.uint64) + cluster_members: dict[int, list[int]] = defaultdict(list) + for i, bf in enumerate(birch._get_leaf_bfs()): + cluster_ls[bf_labels[i]] += bf.linear_sum + cluster_samples[bf_labels[i]] += bf.n_samples + cluster_members[bf_labels[i]].extend(bf.mol_indices) + + centroids = (cluster_ls >= cluster_samples * 0.5).view(np.uint8) + self.cluster_centroids_packed = np.packbits(centroids, axis=-1) + self.cluster_members = cluster_members + + self.is_built = True + if verbose: + print(f"IVF index built with {len(self.cluster_centroids_packed)} clusters") + + def _find_candidates( + self, query_fp_packed: NDArray[np.uint8], n_probe: int, verbose: bool = False + ) -> 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 + """ + + # Limit n_probe to available clusters + n_probe = min(n_probe, len(self.cluster_centroids_packed)) + + t1 = time.time() + similarities = jt_sim_packed(self.cluster_centroids_packed, query_fp_packed) + centroid_sim_time = time.time() - t1 + + t2 = time.time() + # Get indices of top n_probe most similar centroids + # TODO: Probably inefficient + top_indices = np.argsort(similarities)[-n_probe:][::-1] # Sort descending + # Map index to cluster ID - fix the mapping! + members = self.cluster_members + candidates = [] + for idx in top_indices: + candidates.extend(members[idx.item()]) + sort_time = time.time() - t2 + + if verbose: + print(" Cluster search details:") + print( + f" Centroid sims:" + f" {centroid_sim_time*1000:.2f}ms" + f" (vs {len(self.cluster_centroids_packed)} centroids)" + ) + print(f" Sorting/mapping/gathering: {sort_time*1000:.2f}ms") + 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 = False, + verbose: bool = False, + ) -> list[dict[str, tp.Union[int, float, str]]]: + """ + 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) + + Returns: + list of dictionaries containing search results, each with: + - 'index': Index of the fingerprint + - 'similarity': Tanimoto similarity to query + - 'smiles': SMILES string (if available) + """ + + if not self.is_built: + raise RuntimeError("Index has not been built. Call build_index first.") + if not input_is_packed: + query_fp = pack_fingerprints(query_fp) + # Find nearest clusters + t1 = time.time() + candidate_indices = self._find_candidates(query_fp, n_probe, verbose=verbose) + cluster_time = time.time() - t1 + + # Calculate similarities based on method and available formats + t3 = time.time() + # Get candidate fingerprints and convert if needed + similarities = jt_sim_packed(self.fingerprints[candidate_indices], query_fp) + sim_time = time.time() - t3 + + # Apply threshold filter + t4 = time.time() + if threshold > 0.0: + valid_idxs = (similarities > threshold).nonzero()[0].reshape(-1) + similarities = similarities[valid_idxs] + candidate_indices = candidate_indices[valid_idxs] + + # Sort by similarity (descending) + sorted_indices = np.argsort(similarities)[::-1][:k] + + # Prepare results + # TODO: Inefficient + results = [] + for idx in sorted_indices: + result = { + "index": candidate_indices[idx], + "similarity": similarities[idx], + } + + # Add SMILES if available + if self.smiles: + result["smiles"] = self.smiles[candidate_indices[idx]] + results.append(result) + post_time = time.time() - t4 + + # Print timing breakdown + total_time = cluster_time + sim_time + post_time + if verbose: + print("IVF Search timing breakdown:") + print( + f" Find clusters and gather candidates: {cluster_time*1000:.2f}ms ({cluster_time/total_time*100:.1f}%)" + ) + print( + f" Similarity calc: {sim_time*1000:.2f}ms ({sim_time/total_time*100:.1f}%)" + ) + print( + f" Post-processing: {post_time*1000:.2f}ms ({post_time/total_time*100:.1f}%)" + ) + print( + f" Total: {total_time*1000:.2f}ms, candidates: {len(candidate_indices)}" + ) + return results From ee10a011a9552fd337c6678d8a4d3010c9ff5f8a Mon Sep 17 00:00:00 2001 From: ipickering Date: Wed, 29 Oct 2025 22:30:30 -0400 Subject: [PATCH 02/23] Delete old ivf index file --- bblean/ivf_index.py | 256 -------------------------------------------- 1 file changed, 256 deletions(-) delete mode 100644 bblean/ivf_index.py diff --git a/bblean/ivf_index.py b/bblean/ivf_index.py deleted file mode 100644 index d75acf3f..00000000 --- a/bblean/ivf_index.py +++ /dev/null @@ -1,256 +0,0 @@ -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). -""" - -from collections import defaultdict -from numpy.typing import NDArray -import math -import time -import typing as tp -import numpy as np - -from bblean.bitbirch import BitBirch -from bblean.similarity import jt_sim_packed -from bblean.fingerprints import pack_fingerprints - - -class IVFIndex: - """ - 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. - - Attributes: - n_clusters (int): Number of clusters to use. If None, uses sqrt(n_samples) - threshold (float): Similarity threshold for BitBIRCH clustering - branching_factor (int): Branching factor for BitBIRCH clustering - cluster_centroids (np.ndarray): Centroids of each cluster - cluster_members (Dict[int, list[int]]): Mapping of cluster IDs to member - fingerprint indices - fingerprints (np.ndarray): Stored fingerprints for similarity search - smiles (list[str]): Optional SMILES strings corresponding to fingerprints - built (bool): Whether the index has been built - """ - - def __init__( - self, - n_clusters: int | None = None, - threshold: float = 0.7, - branching_factor: int = 50, - ): - """ - Initialize the IVF index. - - Args: - n_clusters: Number of clusters to use (required) - threshold: Similarity threshold for BitBIRCH clustering (used only for tree - building) - branching_factor: Branching factor for BitBIRCH clustering - """ - if n_clusters is not None and n_clusters <= 0: - raise ValueError("n_clusters must be a positive integer or None") - - self.n_clusters = n_clusters - self.threshold = threshold - self.branching_factor = branching_factor - - # Will be populated during build_index - self.cluster_centroids_packed = np.array([], dtype=np.uint8) # packed - self.cluster_members: dict[int, list[int]] = {} - self.fingerprints = np.array([], dtype=np.uint8) - self.smiles: list[str] = [] - self.is_built = False - - def build_index( - self, - fingerprints: NDArray[np.uint8], - smiles: tp.Sequence[str] = (), - method: str = "kmeans-normalized", - input_is_packed: bool = False, - n_features: int | None = None, - verbose: bool = False, - **method_kwargs: tp.Any, - ) -> None: - """ - Build the IVF index by clustering fingerprints using BitBIRCH. - - Args: - fingerprints: Binary fingerprints of shape (n_samples, n_features) - smiles: Optional list of SMILES strings corresponding to fingerprints - """ - n_samples = fingerprints.shape[0] - if self.n_clusters is None: - n_clusters = max(int(math.sqrt(n_samples)), 1) - else: - n_clusters = self.n_clusters - - fingerprints = fingerprints.astype(np.uint8, copy=False) - if not input_is_packed: - fingerprints = pack_fingerprints(fingerprints) - - # Store (packed) fingerprints and smiles for later use - self.fingerprints = fingerprints - self.smiles = list(smiles) - - # Always use k-clusters functionality since n_clusters is required - if verbose: - print(f"Clustering {n_samples} fps into exactly {n_clusters} clusters...") - - # Initialize BitBIRCH for clustering - birch = BitBirch( - threshold=self.threshold, branching_factor=self.branching_factor - ) - birch.fit(fingerprints) - birch.global_clustering(method=method, n_clusters=n_clusters, **method_kwargs) - - # Fetch new cluster centroids and members - bf_labels = birch._global_clustering_centroid_labels - unique_clusters = np.unique(birch._global_clustering_centroid_labels) - if verbose: - print(f"Found {len(unique_clusters)} unique clusters") - root = birch._root - assert root is not None # mypy - cluster_ls = np.zeros((len(unique_clusters), root.n_features), dtype=np.uint64) - cluster_samples = np.zeros(len(unique_clusters), dtype=np.uint64) - cluster_members: dict[int, list[int]] = defaultdict(list) - for i, bf in enumerate(birch._get_leaf_bfs()): - cluster_ls[bf_labels[i]] += bf.linear_sum - cluster_samples[bf_labels[i]] += bf.n_samples - cluster_members[bf_labels[i]].extend(bf.mol_indices) - - centroids = (cluster_ls >= cluster_samples * 0.5).view(np.uint8) - self.cluster_centroids_packed = np.packbits(centroids, axis=-1) - self.cluster_members = cluster_members - - self.is_built = True - if verbose: - print(f"IVF index built with {len(self.cluster_centroids_packed)} clusters") - - def _find_candidates( - self, query_fp_packed: NDArray[np.uint8], n_probe: int, verbose: bool = False - ) -> 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 - """ - - # Limit n_probe to available clusters - n_probe = min(n_probe, len(self.cluster_centroids_packed)) - - t1 = time.time() - similarities = jt_sim_packed(self.cluster_centroids_packed, query_fp_packed) - centroid_sim_time = time.time() - t1 - - t2 = time.time() - # Get indices of top n_probe most similar centroids - # TODO: Probably inefficient - top_indices = np.argsort(similarities)[-n_probe:][::-1] # Sort descending - # Map index to cluster ID - fix the mapping! - members = self.cluster_members - candidates = [] - for idx in top_indices: - candidates.extend(members[idx.item()]) - sort_time = time.time() - t2 - - if verbose: - print(" Cluster search details:") - print( - f" Centroid sims:" - f" {centroid_sim_time*1000:.2f}ms" - f" (vs {len(self.cluster_centroids_packed)} centroids)" - ) - print(f" Sorting/mapping/gathering: {sort_time*1000:.2f}ms") - 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 = False, - verbose: bool = False, - ) -> list[dict[str, tp.Union[int, float, str]]]: - """ - 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) - - Returns: - list of dictionaries containing search results, each with: - - 'index': Index of the fingerprint - - 'similarity': Tanimoto similarity to query - - 'smiles': SMILES string (if available) - """ - - if not self.is_built: - raise RuntimeError("Index has not been built. Call build_index first.") - if not input_is_packed: - query_fp = pack_fingerprints(query_fp) - # Find nearest clusters - t1 = time.time() - candidate_indices = self._find_candidates(query_fp, n_probe, verbose=verbose) - cluster_time = time.time() - t1 - - # Calculate similarities based on method and available formats - t3 = time.time() - # Get candidate fingerprints and convert if needed - similarities = jt_sim_packed(self.fingerprints[candidate_indices], query_fp) - sim_time = time.time() - t3 - - # Apply threshold filter - t4 = time.time() - if threshold > 0.0: - valid_idxs = (similarities > threshold).nonzero()[0].reshape(-1) - similarities = similarities[valid_idxs] - candidate_indices = candidate_indices[valid_idxs] - - # Sort by similarity (descending) - sorted_indices = np.argsort(similarities)[::-1][:k] - - # Prepare results - # TODO: Inefficient - results = [] - for idx in sorted_indices: - result = { - "index": candidate_indices[idx], - "similarity": similarities[idx], - } - - # Add SMILES if available - if self.smiles: - result["smiles"] = self.smiles[candidate_indices[idx]] - results.append(result) - post_time = time.time() - t4 - - # Print timing breakdown - total_time = cluster_time + sim_time + post_time - if verbose: - print("IVF Search timing breakdown:") - print( - f" Find clusters and gather candidates: {cluster_time*1000:.2f}ms ({cluster_time/total_time*100:.1f}%)" - ) - print( - f" Similarity calc: {sim_time*1000:.2f}ms ({sim_time/total_time*100:.1f}%)" - ) - print( - f" Post-processing: {post_time*1000:.2f}ms ({post_time/total_time*100:.1f}%)" - ) - print( - f" Total: {total_time*1000:.2f}ms, candidates: {len(candidate_indices)}" - ) - return results From 73656b6e920c56067fcf124cc5c36b2b07e694a3 Mon Sep 17 00:00:00 2001 From: ipickering Date: Wed, 29 Oct 2025 22:31:01 -0400 Subject: [PATCH 03/23] Add ivf impl --- bblean/_ivf.py | 175 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 bblean/_ivf.py diff --git a/bblean/_ivf.py b/bblean/_ivf.py new file mode 100644 index 00000000..f286cbf2 --- /dev/null +++ b/bblean/_ivf.py @@ -0,0 +1,175 @@ +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 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.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_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, + **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) + + labels = BitBirch._centrals_global_clustering( + centrals, n_clusters, method=method, **method_kwargs + ) + + 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 - 1, n_clusters) + medoids = BitBirch._unpacked_medoids_from_members(fps, mol_ids) + fps = pack_fingerprints(fps) + return cls(pack_fingerprints(medoids), mol_ids, 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] if self._smiles.size > 0 else None + results.append(SearchResult(fp_idx, similarities[idx].item(), smi)) + return results From b39118a9ad7f409f0476c7077476fd1dab7ca45c Mon Sep 17 00:00:00 2001 From: ipickering Date: Wed, 29 Oct 2025 22:51:11 -0400 Subject: [PATCH 04/23] Initial IVF implementation, library API and CLI --- bblean/_ivf.py | 4 +- bblean/bitbirch.py | 170 +++++++++++++++++++++++++++++++----- bblean/cli.py | 184 +++++++++++++++++++++++++++++++++++++++ tests/legacy_fns.py | 24 +++++ tests/test_similarity.py | 15 ++++ 5 files changed, 373 insertions(+), 24 deletions(-) create mode 100644 tests/legacy_fns.py diff --git a/bblean/_ivf.py b/bblean/_ivf.py index f286cbf2..9b874e25 100644 --- a/bblean/_ivf.py +++ b/bblean/_ivf.py @@ -34,7 +34,7 @@ def __repr__(self) -> str: 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}, smi='{self.smi}')" return f"{out})" @@ -170,6 +170,6 @@ def search( results = [] for idx in sorted_indices: fp_idx = candidates[idx].item() - smi = self._smiles[fp_idx] if self._smiles.size > 0 else None + 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/bitbirch.py b/bblean/bitbirch.py index c8be990e..5f716abc 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -70,6 +70,7 @@ from bblean.similarity import ( _jt_sim_arr_vec_packed, jt_most_dissimilar_packed, + jt_isim_medoid, centroid_from_sum, ) @@ -527,6 +528,11 @@ class _CentroidsMolIds(tp.TypedDict): mol_ids: list[list[int]] +class _MedoidsMolIds(tp.TypedDict): + medoids: NDArray[np.uint8] + mol_ids: list[list[int]] + + class BitBirch: r"""Implements the BitBIRCH clustering algorithm, 'Lean' version @@ -625,7 +631,8 @@ def __init__( self._root: _BFNode | None = None self._dummy_leaf = _BFNode(branching_factor=2, n_features=0) # TODO: Type correctly - self._global_clustering_centroid_labels: tp.Any = None + self._global_clustering_centroid_labels: NDArray[np.int64] | None = None + self._n_global_clusters = 0 # For backwards compatibility, weak-register in global state This is used to # update the merge_accept function if the global set_merge() is called @@ -891,7 +898,9 @@ def get_centroids_mol_ids( return {"centroids": centroids, "mol_ids": mol_ids} def get_centroids( - self, sort: bool = True, packed: bool = True + self, + sort: bool = True, + packed: bool = True, ) -> list[NDArray[np.uint8]]: r"""Get a list of arrays with the centroids' fingerprints""" # NOTE: This is different from the original bitbirch, here outputs are sorted by @@ -899,12 +908,94 @@ def get_centroids( attr = "packed_centroid" if packed else "unpacked_centroid" return [getattr(s, attr) for s in self._get_leaf_bfs(sort=sort)] - def get_cluster_mol_ids(self, sort: bool = True) -> list[list[int]]: + def get_medoids_mol_ids( + self, + fps: NDArray[np.uint8], + sort: bool = True, + pack: bool = True, + global_clusters: bool = False, + input_is_packed: bool = True, + n_features: int | None = None, + ) -> _MedoidsMolIds: + """Get a dict with medoids and mol indices of the leaves""" + cluster_members = self.get_cluster_mol_ids( + sort=sort, global_clusters=global_clusters + ) + + if input_is_packed: + fps = _unpack_fingerprints(fps, n_features=n_features) + cluster_medoids = self._unpacked_medoids_from_members(fps, cluster_members) + if pack: + cluster_medoids = pack_fingerprints(cluster_medoids) + return {"medoids": cluster_medoids, "mol_ids": cluster_members} + + @staticmethod + def _unpacked_medoids_from_members( + unpacked_fps: NDArray[np.uint8], cluster_members: tp.Sequence[list[int]] + ) -> NDArray[np.uint8]: + cluster_medoids = np.zeros( + (len(cluster_members), unpacked_fps.shape[1]), dtype=np.uint8 + ) + for idx, members in enumerate(cluster_members): + cluster_medoids[idx, :] = jt_isim_medoid( + unpacked_fps[members], + input_is_packed=False, + pack=False, + )[1] + return cluster_medoids + + def get_medoids( + self, + fps: NDArray[np.uint8], + sort: bool = True, + pack: bool = True, + global_clusters: bool = False, + input_is_packed: bool = True, + n_features: int | None = None, + ) -> NDArray[np.uint8]: + return self.get_medoids_mol_ids( + fps, sort, pack, global_clusters, input_is_packed, n_features + )["medoids"] + + def get_cluster_mol_ids( + self, sort: bool = True, global_clusters: bool = False + ) -> list[list[int]]: r"""Get the indices of the molecules in each cluster""" + if global_clusters: + if self._global_clustering_centroid_labels is None: + raise ValueError( + "Must perform global clustering before fetching global labels" + ) + bf_labels = ( + self._global_clustering_centroid_labels - 1 + ) # sub 1 to use as idxs + + # Collect the members of all clusters + it = (bf.mol_indices for bf in self._get_leaf_bfs(sort=sort)) + return self._new_ids_from_labels(it, bf_labels, self._n_global_clusters) + return [s.mol_indices for s in self._get_leaf_bfs(sort=sort)] + @staticmethod + def _new_ids_from_labels( + members: tp.Iterable[list[int]], + labels: NDArray[np.int64], + n_labels: int | None = None, + ) -> list[list[int]]: + r"""Get the indices of the molecules in each cluster""" + if n_labels is None: + n_labels = len(np.unique(labels)) + new_members: list[list[int]] = [[] for _ in range(n_labels)] + for i, idxs in enumerate(members): + new_members[labels[i]].extend(idxs) + return new_members + def get_assignments( - self, n_mols: int | None = None, sort: bool = True, check_valid: bool = True + self, + n_mols: int | None = None, + sort: bool = True, + check_valid: bool = True, + global_clusters: bool = False, ) -> NDArray[np.uint64]: r"""Get an array with the cluster labels associated with each fingerprint idx""" if n_mols is not None: @@ -927,7 +1018,11 @@ def get_assignments( s.mol_indices for leaf in self._get_leaves() for s in leaf._subclusters ) - if self._global_clustering_centroid_labels is not None: + if global_clusters: + if self._global_clustering_centroid_labels is None: + raise ValueError( + "Must perform global clustering before fetching global labels" + ) # Assign according to global clustering labels final_labels = self._global_clustering_centroid_labels for mol_ids, label in zip(iterator, final_labels): @@ -947,6 +1042,7 @@ def dump_assignments( path: Path | str, smiles: tp.Iterable[str] = (), sort: bool = True, + global_clusters: bool = False, check_valid: bool = True, ) -> None: r"""Dump the cluster assignments to a ``*.csv`` file""" @@ -957,7 +1053,9 @@ def dump_assignments( smiles = [smiles] smiles = np.asarray(smiles, dtype=np.str_) # Dump cluster assignments to *.csv - assignments = self.get_assignments(sort=sort, check_valid=check_valid) + assignments = self.get_assignments( + sort=sort, check_valid=check_valid, global_clusters=global_clusters + ) if smiles.size and (len(assignments) != len(smiles)): raise ValueError( f"Len of the provided smiles {len(smiles)}" @@ -1220,13 +1318,54 @@ def global_clustering( **method_kwargs: tp.Any, ) -> tpx.Self: r""":meta private:""" + if not self.is_init: + raise ValueError("The model has not been fitted yet.") + # Add 1 to start labels from 1 instead of 0, so 0 can be used as sentinel + # value + centroids = np.vstack(self.get_centroids(packed=False)) + labels = self._centrals_global_clustering( + centroids, n_clusters, method=method, input_is_packed=False, **method_kwargs + ) + num_centroids = len(centroids) + self._n_global_clusters = ( + n_clusters if num_centroids > n_clusters else num_centroids + ) + self._global_clustering_centroid_labels = labels + return self + + @staticmethod + def _centrals_global_clustering( + centrals: NDArray[np.uint8], + n_clusters: int, + *, + method: str = "kmeans", + input_is_packed: bool = True, + n_features: int | None = None, + # TODO: Type correctly + **method_kwargs: tp.Any, + ) -> NDArray[np.int64]: + r""":meta private:""" + if method not in {"agglomerative", "kmeans", "kmeans-normalized"}: + raise ValueError(f"Unknown method {method}") + # Returns the labels associated with global clustering # Lazy import because sklearn is very heavy from sklearn.cluster import KMeans, AgglomerativeClustering from sklearn.exceptions import ConvergenceWarning - if not self.is_init: - raise ValueError("The model has not been fitted yet.") + if input_is_packed: + centrals = _unpack_fingerprints(centrals, n_features) + num_centrals = len(centrals) + if num_centrals < n_clusters: + msg = ( + f"Number of subclusters found ({num_centrals}) by BitBIRCH is less " + "than ({n_clusters}). Decrease k or the threshold." + ) + warnings.warn(msg, ConvergenceWarning, stacklevel=2) + n_clusters = num_centrals + + if method == "kmeans-normalized": + centrals = centrals / np.linalg.norm(centrals, axis=1, keepdims=True) if method in ["kmeans", "kmeans-normalized"]: predictor = KMeans(n_clusters=n_clusters, **method_kwargs) elif method == "agglomerative": @@ -1234,22 +1373,9 @@ def global_clustering( else: raise ValueError("method must be one of 'kmeans' or 'agglomerative'") - centroids = np.vstack(self.get_centroids(packed=False)) - if method == "kmeans-normalized": - centroids = centroids / np.linalg.norm(centroids, axis=1, keepdims=True) - num_centroids = len(centroids) - if num_centroids < n_clusters: - msg = ( - f"Number of subclusters found ({num_centroids}) by BitBIRCH is less " - "than ({n_clusters}). Decrease k or the threshold." - ) - warnings.warn(msg, ConvergenceWarning, stacklevel=2) - n_clusters = num_centroids - # Add 1 to start labels from 1 instead of 0, so 0 can be used as sentinel # value - self._global_clustering_centroid_labels = predictor.fit_predict(centroids) + 1 - return self + return predictor.fit_predict(centrals) + 1 # There are 4 cases here: diff --git a/bblean/cli.py b/bblean/cli.py index 08004773..b5be63cc 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1,5 +1,6 @@ r"""Command line interface entrypoints""" +import numpy as np import warnings import random import typing as tp @@ -863,6 +864,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, @@ -1125,12 +1135,19 @@ 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() + input_smiles_dir.mkdir() if copy_inputs: for file in input_files: shutil.copy(file, input_fps_dir / file.name) + if smiles_path: + shutil.copy(smiles_path, input_smiles_dir / smiles_path.name) else: for file in input_files: (input_fps_dir / file.name).symlink_to(file.resolve()) + if smiles_path: + (input_smiles_dir / smiles_path.name).symlink_to(smiles_path.resolve()) # TODO: Currently sometimes after a round is triggered *more* files are output, since @@ -1782,6 +1799,173 @@ def _split_fps( ) +@app.command("query-ivf", hidden=True) +def _query_ivf( + ivf_path: Annotated[ + Path, + Argument(help="Path to the IVF file, or a dir with a ivf.pkl file"), + ], + 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, +) -> None: + from bblean._console import get_console + from bblean.fingerprints import fps_from_smiles + + console = get_console() + + if ivf_path.is_dir(): + ivf_file = ivf_path / "ivf.pkl" + else: + ivf_file = ivf_path + + with open(ivf_file, mode="rb") as f: + ivf_index = pickle.load(f) + # TODO: This is a placeholder, must be modified!! + query_fp = fps_from_smiles([query], kind="rdkit", n_features=2048, pack=True)[0] + results_list = ivf_index.search( + query_fp, n_probe=num_probe, threshold=threshold, k=k + ) + + for r in results_list: + console.print(r) + + +@app.command("build-ivf", hidden=True) +def _build_ivf( + 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-path", + 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, + verbose: Annotated[ + bool, + Option("--verbose/--no-verbose"), + ] = True, +) -> None: + from bblean.utils import _has_files_or_valid_symlinks + from bblean._ivf import IVFIndex + from bblean.smiles import load_smiles + from bblean._console import get_console + + console = get_console(silent=not verbose) + + if clusters_path.is_dir(): + centroids_path = clusters_path / "cluster-centroids-packed.pkl" + clusters_path = clusters_path / "clusters.pkl" + if not centroids_path.exists(): + raise ValueError( + "Centroids must be saved for building IVF." + " This limitation may be lifted in the future" + ) + else: + raise ValueError( + "clusters path must be a dir for building IVF" + " This limitation may be lifted in the future" + ) + + with open(clusters_path, mode="rb") as f: + cluster_members = pickle.load(f) + + # TODO: If this file doesn't exist, it will have to be created on-the-fly + with open(centroids_path, mode="rb") as f: + centroids_packed = np.vstack(pickle.load(f)) + + if fps_path is None: + input_fps_path = clusters_path.parent / "input-fps" + if input_fps_path.is_dir() and _has_files_or_valid_symlinks(input_fps_path): + fps_path = input_fps_path + else: + msg = ( + "Could not find input fingerprints. Please use --fps-path." + " Summary plot without fingerprints doesn't include isim values" + ) + warnings.warn(msg) + if smiles_path is None: + input_smiles_path = clusters_path.parent / "input-smiles" + if input_smiles_path.is_dir() and _has_files_or_valid_symlinks( + input_smiles_path + ): + smiles_paths = sorted(input_smiles_path.glob("*.smi")) + if len(smiles_paths) > 1: + raise ValueError("Currently only a single smiles file is supported") + smiles_path = smiles_paths[0] + else: + msg = ( + "Could not find input smiles. Please use --smiles-path." + " Search results won't include smiles" + ) + warnings.warn(msg) + + if fps_path is None: + raise ValueError("Fingerprints are required to build the index") + elif fps_path.is_dir(): + fps_paths = sorted(fps_path.glob("*.npy")) + else: + fps_paths = [fps_path] + if len(fps_paths) > 1: + raise ValueError( + "Currently only a single fp file is supported," + " this restriction will be lifted in the future" + ) + + fps = np.load(fps_paths[0]) # packed + with console.status("[italic]Building IVF index...[/italic]", spinner="dots"): + index = IVFIndex.from_bitbirch_clusters( + cluster_members, + centroids_packed, + fps, + load_smiles(smiles_path) if smiles_path is not None else (), + method, + n_clusters, + input_is_packed=True, + random_state=42, + ) + + with console.status("[italic]Saving IVF index...[/italic]", spinner="dots"): + with open(clusters_path.parent / "ivf.pkl", mode="wb") as f: + pickle.dump(index, f) + console.print("Successfully built IVF index") + + @app.command("fps-shuffle", rich_help_panel="Fingerprints") def _shuffle_fps( in_file: Annotated[ 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 b39193b8..4cbf9692 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_comp_sim, calculate_medoid # TODO: Fix the tests with pytest-subtests so that both the _py_similarity and the # _cpp_similarity are tested independently @@ -255,6 +256,19 @@ def test_jt_compl_isim() -> None: _ = pysim.jt_compl_isim(fps) fps = make_fake_fingerprints(10, seed=17408390758220920002, pack=False) + assert calculate_comp_sim(fps).tolist() == [ + 0.20256457907452147, + 0.24748926949201983, + 0.22550084742079876, + 0.2002884861456855, + 0.23889840001690868, + 0.2364222674813306, + 0.1986207548061027, + 0.19904732709222533, + 0.21303348506016495, + 0.2225069540267648, + ] + assert pysim.jt_compl_isim(fps).tolist() == snapshot( [ 0.20256457907452147, @@ -282,3 +296,4 @@ def test_jt_isim_medoid() -> None: idx, m = pysim.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] From 23e0924d61d7d19830fcb0b12e37a53bbbf1d9c5 Mon Sep 17 00:00:00 2001 From: ipickering Date: Wed, 29 Oct 2025 22:52:27 -0400 Subject: [PATCH 05/23] Fix mypy --- tests/test_similarity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_similarity.py b/tests/test_similarity.py index 4cbf9692..40eee770 100644 --- a/tests/test_similarity.py +++ b/tests/test_similarity.py @@ -3,7 +3,7 @@ import pytest from inline_snapshot import snapshot -from legacy_fns import calculate_comp_sim, calculate_medoid +from legacy_fns import calculate_comp_sim, calculate_medoid # type: ignore # TODO: Fix the tests with pytest-subtests so that both the _py_similarity and the # _cpp_similarity are tested independently From 9a9a3b180e6f80756357a9633c99a0a5074a77ba Mon Sep 17 00:00:00 2001 From: ipickering Date: Wed, 29 Oct 2025 22:53:50 -0400 Subject: [PATCH 06/23] Change names --- bblean/cli.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index b5be63cc..9b67f688 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1799,8 +1799,8 @@ def _split_fps( ) -@app.command("query-ivf", hidden=True) -def _query_ivf( +@app.command("query-idx", hidden=True) +def _query_idx( ivf_path: Annotated[ Path, Argument(help="Path to the IVF file, or a dir with a ivf.pkl file"), @@ -1844,8 +1844,8 @@ def _query_ivf( console.print(r) -@app.command("build-ivf", hidden=True) -def _build_ivf( +@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"), From 0baeb59e3db9636a042be8c56336d70270680ec3 Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:55:39 -0400 Subject: [PATCH 07/23] Add comment --- bblean/bitbirch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index 5f716abc..e46873ba 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -1375,6 +1375,8 @@ def _centrals_global_clustering( # 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 From adc5d5ce6fea5ebb67b31cdaf94cc549c0ca018a Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:56:14 -0400 Subject: [PATCH 08/23] Save IVF index and add --global flag to all methods --- bblean/cli.py | 104 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 70 insertions(+), 34 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index 9b67f688..d119e515 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1,6 +1,5 @@ r"""Command line interface entrypoints""" -import numpy as np import warnings import random import typing as tp @@ -133,6 +132,7 @@ 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 @@ -157,6 +157,7 @@ def _plot_pops( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -250,6 +251,7 @@ 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 @@ -282,6 +284,7 @@ def _plot_umap( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -350,6 +353,7 @@ 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 @@ -374,6 +378,7 @@ def _plot_pca( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -505,6 +510,7 @@ 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 @@ -541,6 +547,7 @@ def _plot_tsne( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -620,6 +627,7 @@ def _table_summary( int, Option("--metrics-min-size", hidden=True), ] = 1, + use_global: Annotated[bool, Option("--global/--no-global"),] = False, verbose: Annotated[ bool, Option("--verbose/--no-verbose", hidden=True), @@ -637,7 +645,9 @@ def _table_summary( # Imports may take a bit of time since sklearn is slow, so start the spinner here with console.status("[italic]Analyzing clusters...[/italic]", spinner="dots"): 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: @@ -820,6 +830,7 @@ 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 @@ -846,6 +857,7 @@ def _plot_summary( verbose=verbose, save=save, show=show, + use_global=use_global, ) @@ -1801,9 +1813,9 @@ def _split_fps( @app.command("query-idx", hidden=True) def _query_idx( - ivf_path: Annotated[ + idx_path: Annotated[ Path, - Argument(help="Path to the IVF file, or a dir with a ivf.pkl file"), + Argument(help="Path to the index dir, or its parent dir with index files"), ], query: Annotated[ str, @@ -1824,26 +1836,41 @@ def _query_idx( ) -> None: from bblean._console import get_console from bblean.fingerprints import fps_from_smiles + from bblean._ivf import IVFIndex console = get_console() - - if ivf_path.is_dir(): - ivf_file = ivf_path / "ivf.pkl" + if idx_path.is_dir() and (idx_path / "index").is_dir(): + idx_path = idx_path / "index" else: - ivf_file = ivf_path + idx_path = idx_path - with open(ivf_file, mode="rb") as f: - ivf_index = pickle.load(f) + index = IVFIndex.from_dir(idx_path) # TODO: This is a placeholder, must be modified!! query_fp = fps_from_smiles([query], kind="rdkit", n_features=2048, pack=True)[0] - results_list = ivf_index.search( - query_fp, n_probe=num_probe, threshold=threshold, k=k - ) + 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[ @@ -1881,31 +1908,17 @@ def _build_idx( 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.smiles import load_smiles from bblean._console import get_console console = get_console(silent=not verbose) - - if clusters_path.is_dir(): - centroids_path = clusters_path / "cluster-centroids-packed.pkl" - clusters_path = clusters_path / "clusters.pkl" - if not centroids_path.exists(): - raise ValueError( - "Centroids must be saved for building IVF." - " This limitation may be lifted in the future" - ) - else: - raise ValueError( - "clusters path must be a dir for building IVF" - " This limitation may be lifted in the future" - ) - + clusters_path, centroids_path = _get_centroids_and_clusters_paths(clusters_path) with open(clusters_path, mode="rb") as f: cluster_members = pickle.load(f) - # TODO: If this file doesn't exist, it will have to be created on-the-fly with open(centroids_path, mode="rb") as f: centroids_packed = np.vstack(pickle.load(f)) @@ -1946,8 +1959,11 @@ def _build_idx( "Currently only a single fp file is supported," " this restriction will be lifted in the future" ) + fps_path = fps_paths[0] + + fps = np.load(fps_path) # packed + kwargs = {"random_state": 42} if method.startswith("kmeans") else {} - fps = np.load(fps_paths[0]) # packed with console.status("[italic]Building IVF index...[/italic]", spinner="dots"): index = IVFIndex.from_bitbirch_clusters( cluster_members, @@ -1957,12 +1973,32 @@ def _build_idx( method, n_clusters, input_is_packed=True, - random_state=42, + sort=True, + **kwargs, ) - with console.status("[italic]Saving IVF index...[/italic]", spinner="dots"): - with open(clusters_path.parent / "ivf.pkl", mode="wb") as f: - pickle.dump(index, f) + 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) + idx_path = clusters_path.parent / "index" + if idx_path.exists(): + shutil.rmtree(idx_path) + idx_path.mkdir(exist_ok=True) + if smiles_path is not None: + (idx_path / "smiles.smi").symlink_to(smiles_path.resolve()) + (idx_path / "fps.npy").symlink_to(fps_path.resolve()) + (idx_path / "global-clusters.pkl").symlink_to(global_clusters_path.resolve()) + (idx_path / "global-cluster-medoids-packed.npy").symlink_to( + global_cluster_medoids_path.resolve() + ) console.print("Successfully built IVF index") From dbdf88bf9a361fec22a789176bb1d797045116c0 Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:56:30 -0400 Subject: [PATCH 09/23] Build ivf index from dir --- bblean/_ivf.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/bblean/_ivf.py b/bblean/_ivf.py index 9b874e25..5831edba 100644 --- a/bblean/_ivf.py +++ b/bblean/_ivf.py @@ -4,6 +4,8 @@ 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 @@ -12,6 +14,7 @@ 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 @@ -65,6 +68,19 @@ def __init__( 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" + fps_path = idx_path / "fps.npy" + smiles_path = idx_path / "smiles.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) + smiles = load_smiles(smiles_path) + return cls(medoids_packed, members, fps, smiles) + @classmethod def from_bitbirch_clusters( cls, @@ -76,6 +92,7 @@ def from_bitbirch_clusters( n_clusters: int | None = None, input_is_packed: bool = True, n_features: int | None = None, + sort: bool = True, **method_kwargs: tp.Any, ) -> tpx.Self: """Build the IVF index from bitbirch clusters""" @@ -89,7 +106,6 @@ def from_bitbirch_clusters( if input_is_packed: fps = unpack_fingerprints(fps, n_features) centrals = unpack_fingerprints(centrals, n_features) - labels = BitBirch._centrals_global_clustering( centrals, n_clusters, method=method, **method_kwargs ) @@ -97,6 +113,8 @@ def from_bitbirch_clusters( 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 - 1, n_clusters) + if sort: + mol_ids.sort(key=lambda x: len(x), reverse=True) medoids = BitBirch._unpacked_medoids_from_members(fps, mol_ids) fps = pack_fingerprints(fps) return cls(pack_fingerprints(medoids), mol_ids, fps, smiles) From 7e4b8285f93583de168142fbec5d59db518bfacc Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:56:43 -0400 Subject: [PATCH 10/23] Add flag to plotting --- bblean/plotting.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/bblean/plotting.py b/bblean/plotting.py index c623f5c0..10febb07 100644 --- a/bblean/plotting.py +++ b/bblean/plotting.py @@ -435,9 +435,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: From 63e0ca8183736a9d866454fd93aebb70d14c075e Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:57:45 -0400 Subject: [PATCH 11/23] Rename _ivf -> ivf --- bblean/cli.py | 4 ++-- bblean/{_ivf.py => ivf.py} | 0 2 files changed, 2 insertions(+), 2 deletions(-) rename bblean/{_ivf.py => ivf.py} (100%) diff --git a/bblean/cli.py b/bblean/cli.py index d119e515..ea58be02 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1836,7 +1836,7 @@ def _query_idx( ) -> None: from bblean._console import get_console from bblean.fingerprints import fps_from_smiles - from bblean._ivf import IVFIndex + from bblean.ivf import IVFIndex console = get_console() if idx_path.is_dir() and (idx_path / "index").is_dir(): @@ -1910,7 +1910,7 @@ def _build_idx( ) -> None: import numpy as np from bblean.utils import _has_files_or_valid_symlinks - from bblean._ivf import IVFIndex + from bblean.ivf import IVFIndex from bblean.smiles import load_smiles from bblean._console import get_console diff --git a/bblean/_ivf.py b/bblean/ivf.py similarity index 100% rename from bblean/_ivf.py rename to bblean/ivf.py From 11c0304b68e752f380ab39eb5a4f65328c0d044d Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:57:59 -0400 Subject: [PATCH 12/23] Black --- bblean/cli.py | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index ea58be02..d9b62097 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -132,7 +132,10 @@ def _plot_pops( bool, Option("--show/--no-show", hidden=True), ] = True, - use_global: Annotated[bool, Option("--global/--no-global"),] = False, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""Population plot of the clustering results""" from bblean._console import get_console @@ -251,7 +254,10 @@ def _plot_umap( rich_help_panel="Advanced", ), ] = None, - use_global: Annotated[bool, Option("--global/--no-global"),] = False, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""UMAP visualization of the clustering results""" from bblean._console import get_console @@ -353,7 +359,10 @@ def _plot_pca( str | None, Option("--filename"), ] = None, - use_global: Annotated[bool, Option("--global/--no-global"),] = False, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""PCA visualization of the clustering results""" from bblean._console import get_console @@ -510,7 +519,10 @@ def _plot_tsne( bool, Option("--show/--no-show", hidden=True), ] = True, - use_global: Annotated[bool, Option("--global/--no-global"),] = False, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""t-SNE visualization of the clustering results""" from bblean._console import get_console @@ -627,7 +639,10 @@ def _table_summary( int, Option("--metrics-min-size", hidden=True), ] = 1, - use_global: Annotated[bool, Option("--global/--no-global"),] = False, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, verbose: Annotated[ bool, Option("--verbose/--no-verbose", hidden=True), @@ -830,7 +845,10 @@ def _plot_summary( bool, Option("--show/--no-show", hidden=True), ] = True, - use_global: Annotated[bool, Option("--global/--no-global"),] = False, + use_global: Annotated[ + bool, + Option("--global/--no-global"), + ] = False, ) -> None: r"""Summary plot of the clustering results""" from bblean._console import get_console From 3882aaef5657dd0bacc34557640ad48e4dac1c2a Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 02:59:55 -0400 Subject: [PATCH 13/23] Omit smiles --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index d8a07782..3da18afb 100644 --- a/.gitignore +++ b/.gitignore @@ -72,6 +72,8 @@ cluster-mol-ids*.json config*.json timings*.json max-rss.txt +index/*.smi +input-smiles/*.smi # Pickle *.pkl From e1595e5fa8c30246ca538e0a641de2e46ddfd1bf Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 03:07:53 -0400 Subject: [PATCH 14/23] Add option to build idx after running bblean --- bblean/cli.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/bblean/cli.py b/bblean/cli.py index d9b62097..9c3984c2 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1040,6 +1040,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"), + ] = "kmeans", + idx_n_clusters: Annotated[ + int | None, + Option("--n-clusters"), + ] = None, ) -> None: r"""Run standard, serial BitBIRCH clustering over `*.npy` fingerprint files""" # TODO: Remove code duplication with multiround @@ -1179,6 +1191,10 @@ def _run( if smiles_path: (input_smiles_dir / smiles_path.name).symlink_to(smiles_path.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 # the files are divided *both* by uint8/uint16 and the batch idx. I believe this is not From b765b52758890c8246d180345fbba0d47fdbc03c Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 12:09:04 -0400 Subject: [PATCH 15/23] Omit all smiles files --- .gitignore | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 3da18afb..33e693d8 100644 --- a/.gitignore +++ b/.gitignore @@ -72,8 +72,8 @@ cluster-mol-ids*.json config*.json timings*.json max-rss.txt -index/*.smi -input-smiles/*.smi +*.smi +!chembl-33-natural-products-subset.smi # Pickle *.pkl From 4dfaeb2e1664bbe8b767199ca8fab6a45e083014 Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 30 Oct 2025 12:18:38 -0400 Subject: [PATCH 16/23] Modify smiles file name --- bblean/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index 9c3984c2..285d40a2 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1184,12 +1184,12 @@ def _run( for file in input_files: shutil.copy(file, input_fps_dir / file.name) if smiles_path: - shutil.copy(smiles_path, input_smiles_dir / smiles_path.name) + shutil.copy(smiles_path, input_smiles_dir / "smiles.smi") else: for file in input_files: (input_fps_dir / file.name).symlink_to(file.resolve()) if smiles_path: - (input_smiles_dir / smiles_path.name).symlink_to(smiles_path.resolve()) + (input_smiles_dir / "smiles.smi").symlink_to(smiles_path.resolve()) # Build the index with defaults if build_idx: From 51f2a9408b14cbb71e1a6019dc25dc04de386da0 Mon Sep 17 00:00:00 2001 From: ipickering Date: Sat, 1 Nov 2025 23:50:20 -0400 Subject: [PATCH 17/23] Add IVF to multiround --- bblean/cli.py | 143 +++++++++++++++++++++++++++++--------------------- bblean/ivf.py | 13 +++-- 2 files changed, 91 insertions(+), 65 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index 285d40a2..e5b33a60 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -586,7 +586,7 @@ def _table_summary( Path | None, Option( "-s", - "--smiles-path", + "--smiles", show_default=False, help="Optional smiles path, if passed a scaffold analysis is performed", ), @@ -795,7 +795,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", ), @@ -1046,11 +1046,11 @@ def _run( ] = False, idx_method: Annotated[ str, - Option("--idx-method"), + Option("--idx-method", hidden=True), ] = "kmeans", idx_n_clusters: Annotated[ int | None, - Option("--n-clusters"), + Option("--n-clusters", hidden=True), ] = None, ) -> None: r"""Run standard, serial BitBIRCH clustering over `*.npy` fingerprint files""" @@ -1180,16 +1180,24 @@ def _run( input_smiles_dir = (out_dir / "input-smiles").resolve() input_smiles_dir.mkdir() + + smiles_files = [] + if smiles_path: + smiles_files = ( + [smiles_path] + if not smiles_path.is_dir() + else sorted(smiles_path.glob("*.smi")) + ) if copy_inputs: for file in input_files: shutil.copy(file, input_fps_dir / file.name) - if smiles_path: - shutil.copy(smiles_path, input_smiles_dir / "smiles.smi") + 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()) - if smiles_path: - (input_smiles_dir / "smiles.smi").symlink_to(smiles_path.resolve()) + for file in smiles_files: + (input_smiles_dir / file.name).symlink_to(file.resolve()) # Build the index with defaults if build_idx: @@ -1391,6 +1399,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 @@ -1480,6 +1500,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( @@ -1873,11 +1897,6 @@ def _query_idx( from bblean.ivf import IVFIndex console = get_console() - if idx_path.is_dir() and (idx_path / "index").is_dir(): - idx_path = idx_path / "index" - else: - idx_path = idx_path - index = IVFIndex.from_dir(idx_path) # TODO: This is a placeholder, must be modified!! query_fp = fps_from_smiles([query], kind="rdkit", n_features=2048, pack=True)[0] @@ -1924,7 +1943,7 @@ def _build_idx( Path | None, Option( "-s", - "--smiles-path", + "--smiles", show_default=False, help="Optional smiles path, if used smiles are returned in search results", ), @@ -1945,7 +1964,7 @@ def _build_idx( import numpy as np from bblean.utils import _has_files_or_valid_symlinks from bblean.ivf import IVFIndex - from bblean.smiles import load_smiles + from bblean._console import get_console console = get_console(silent=not verbose) @@ -1956,54 +1975,61 @@ def _build_idx( 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: - input_fps_path = clusters_path.parent / "input-fps" - if input_fps_path.is_dir() and _has_files_or_valid_symlinks(input_fps_path): - fps_path = input_fps_path - else: - msg = ( - "Could not find input fingerprints. Please use --fps-path." - " Summary plot without fingerprints doesn't include isim values" - ) - warnings.warn(msg) - if smiles_path is None: - input_smiles_path = clusters_path.parent / "input-smiles" - if input_smiles_path.is_dir() and _has_files_or_valid_symlinks( - input_smiles_path - ): - smiles_paths = sorted(input_smiles_path.glob("*.smi")) - if len(smiles_paths) > 1: - raise ValueError("Currently only a single smiles file is supported") - smiles_path = smiles_paths[0] - else: - msg = ( - "Could not find input smiles. Please use --smiles-path." - " Search results won't include smiles" - ) - warnings.warn(msg) + 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 None: + 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") - elif fps_path.is_dir(): - fps_paths = sorted(fps_path.glob("*.npy")) + + 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: - fps_paths = [fps_path] - if len(fps_paths) > 1: - raise ValueError( - "Currently only a single fp file is supported," - " this restriction will be lifted in the future" - ) - fps_path = fps_paths[0] + 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_path) # packed + 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, - load_smiles(smiles_path) if smiles_path is not None else (), + (), method, n_clusters, input_is_packed=True, @@ -2022,17 +2048,12 @@ def _build_idx( global_clusters_path = clusters_path.parent / "global-clusters.pkl" with open(global_clusters_path, mode="wb") as f: pickle.dump(index._members, f) - idx_path = clusters_path.parent / "index" - if idx_path.exists(): - shutil.rmtree(idx_path) - idx_path.mkdir(exist_ok=True) + if symlink_fps: + for file in fps_files: + (fps_path / file.name).symlink_to(file.resolve()) if smiles_path is not None: - (idx_path / "smiles.smi").symlink_to(smiles_path.resolve()) - (idx_path / "fps.npy").symlink_to(fps_path.resolve()) - (idx_path / "global-clusters.pkl").symlink_to(global_clusters_path.resolve()) - (idx_path / "global-cluster-medoids-packed.npy").symlink_to( - global_cluster_medoids_path.resolve() - ) + for file in smiles_files: + (smiles_path / file.name).symlink_to(file.resolve()) console.print("Successfully built IVF index") diff --git a/bblean/ivf.py b/bblean/ivf.py index 5831edba..e6bcbafe 100644 --- a/bblean/ivf.py +++ b/bblean/ivf.py @@ -72,14 +72,19 @@ def __init__( 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" - fps_path = idx_path / "fps.npy" - smiles_path = idx_path / "smiles.smi" + 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) - smiles = load_smiles(smiles_path) - return cls(medoids_packed, members, fps, smiles) + return cls(medoids_packed, members, fps, load_smiles(smiles_files)) @classmethod def from_bitbirch_clusters( From 43e762804e5c2e1a1f1d02211876cd599a27f074 Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 23 Feb 2026 14:13:44 -0500 Subject: [PATCH 18/23] Update similarity --- tests/test_similarity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_similarity.py b/tests/test_similarity.py index f21a4ac8..0578b719 100644 --- a/tests/test_similarity.py +++ b/tests/test_similarity.py @@ -3,7 +3,7 @@ import pytest from inline_snapshot import snapshot -from legacy_fns import calculate_comp_sim, calculate_medoid # type: ignore +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 From 0968cf750254eb178384194686afa2a683d9179b Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 23 Feb 2026 14:46:28 -0500 Subject: [PATCH 19/23] Add option for morgan fps in the query --- bblean/cli.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index b5dba2a1..2471edb0 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -2045,6 +2045,14 @@ def _query_idx( 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 @@ -2052,8 +2060,9 @@ def _query_idx( console = get_console() index = IVFIndex.from_dir(idx_path) - # TODO: This is a placeholder, must be modified!! - query_fp = fps_from_smiles([query], kind="rdkit", n_features=2048, pack=True)[0] + 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: From ae6b80696bc67c2e0fc2109f264320ef81806346 Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 23 Feb 2026 17:16:46 -0500 Subject: [PATCH 20/23] Implement direct reassignment --- bblean/bitbirch.py | 7 ++++--- bblean/cli.py | 5 +++++ bblean/ivf.py | 25 ++++++++++++++++++++----- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index 6088e1f4..a1c0d9aa 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -1482,9 +1482,10 @@ 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 ) + labels = predictor.fit_predict(centroids) num_centroids = len(centroids) self._n_global_clusters = ( n_clusters if num_centroids > n_clusters else num_centroids @@ -1493,7 +1494,7 @@ def global_clustering( return self @staticmethod - def _centrals_global_clustering( + def _global_clustering_predictor( centrals: NDArray[np.uint8], n_clusters: int, *, @@ -1502,7 +1503,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}") diff --git a/bblean/cli.py b/bblean/cli.py index 2471edb0..66f785cf 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -2119,6 +2119,10 @@ def _build_idx( 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"), @@ -2197,6 +2201,7 @@ def _build_idx( n_clusters, input_is_packed=True, sort=True, + direct_reassignment=direct_reassignment, **kwargs, ) diff --git a/bblean/ivf.py b/bblean/ivf.py index e6bcbafe..72cdeb3b 100644 --- a/bblean/ivf.py +++ b/bblean/ivf.py @@ -98,6 +98,7 @@ def from_bitbirch_clusters( 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""" @@ -111,16 +112,30 @@ def from_bitbirch_clusters( if input_is_packed: fps = unpack_fingerprints(fps, n_features) centrals = unpack_fingerprints(centrals, n_features) - labels = BitBirch._centrals_global_clustering( + + predictor = BitBirch._global_clustering_predictor( centrals, n_clusters, method=method, **method_kwargs ) + # Direct reassignment reassigns the fingerprints directly using the predictor + # instead of indirectly reassigning using the central labels + if direct_reassignment: + predictor.fit(centrals) + 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).tolist() for i in range(n_clusters)] + else: + labels = predictor.fit_predict(centrals) + 1 + 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 - 1, n_clusters) - 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 - 1, n_clusters) if sort: mol_ids.sort(key=lambda x: len(x), reverse=True) - medoids = BitBirch._unpacked_medoids_from_members(fps, mol_ids) + _, medoids = BitBirch._unpacked_medoids_from_members(fps, mol_ids) fps = pack_fingerprints(fps) return cls(pack_fingerprints(medoids), mol_ids, fps, smiles) From 527a465e91418bb0ce58b85cec29f661eeff728b Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 23 Feb 2026 17:20:30 -0500 Subject: [PATCH 21/23] Fix predictor --- bblean/bitbirch.py | 10 ++++------ bblean/ivf.py | 13 ++++++++----- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index a1c0d9aa..4affb3df 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -1485,7 +1485,9 @@ def global_clustering( predictor = self._global_clustering_predictor( centroids, n_clusters, method=method, input_is_packed=False, **method_kwargs ) - labels = predictor.fit_predict(centroids) + # 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 @@ -1533,11 +1535,7 @@ def _global_clustering_predictor( 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/ivf.py b/bblean/ivf.py index 72cdeb3b..40c9fad6 100644 --- a/bblean/ivf.py +++ b/bblean/ivf.py @@ -116,10 +116,13 @@ def from_bitbirch_clusters( predictor = BitBirch._global_clustering_predictor( centrals, n_clusters, method=method, **method_kwargs ) - # Direct reassignment reassigns the fingerprints directly using the predictor - # instead of indirectly reassigning using the central labels + # 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: - predictor.fit(centrals) + # 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) @@ -128,10 +131,10 @@ def from_bitbirch_clusters( labels = predictor.predict(fps) mol_ids = [(labels == i).tolist() for i in range(n_clusters)] else: - labels = predictor.fit_predict(centrals) + 1 + 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 - 1, n_clusters) + mol_ids = BitBirch._new_ids_from_labels(members, labels, n_clusters) if sort: mol_ids.sort(key=lambda x: len(x), reverse=True) From 575855c6e4c16c92c18ffd2f7d1bf96e0be3cc55 Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 23 Feb 2026 17:37:53 -0500 Subject: [PATCH 22/23] Fix reassignment --- bblean/ivf.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/bblean/ivf.py b/bblean/ivf.py index 40c9fad6..043eb9b1 100644 --- a/bblean/ivf.py +++ b/bblean/ivf.py @@ -129,7 +129,7 @@ def from_bitbirch_clusters( ) else: labels = predictor.predict(fps) - mol_ids = [(labels == i).tolist() for i in range(n_clusters)] + mol_ids = [(labels == i).nonzero()[0].tolist() for i in range(n_clusters)] else: labels = predictor.predict(centrals) num_centrals = len(centrals) @@ -138,9 +138,10 @@ def from_bitbirch_clusters( if sort: mol_ids.sort(key=lambda x: len(x), reverse=True) - _, medoids = BitBirch._unpacked_medoids_from_members(fps, mol_ids) - fps = pack_fingerprints(fps) - return cls(pack_fingerprints(medoids), mol_ids, fps, smiles) + _, 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 From c1e33127e84b67473066bde8c5b329423f8fdb3d Mon Sep 17 00:00:00 2001 From: ipickering Date: Mon, 23 Feb 2026 17:45:37 -0500 Subject: [PATCH 23/23] Fix input smiles dir --- bblean/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bblean/cli.py b/bblean/cli.py index 66f785cf..082903ea 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1309,8 +1309,6 @@ def _run( input_fps_dir.mkdir() input_smiles_dir = (out_dir / "input-smiles").resolve() - input_smiles_dir.mkdir() - smiles_files = [] if smiles_path: smiles_files = ( @@ -1318,6 +1316,8 @@ def _run( 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)