diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index 20a5c67e..68922524 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -211,6 +211,12 @@ def _split_node(node: "_BFNode") -> tuple["_BFSubcluster", "_BFSubcluster"]: return new_subcluster1, new_subcluster2 +class _Assignment(tp.NamedTuple): + cluster_label: int + similarity: float + is_mergeable: bool + + class _BFNode: """Each node in a BitBirch tree is a _BFNode. @@ -356,6 +362,76 @@ def insert_bf_subcluster( ].packed_centroid return False + def find_highest_similarity_subcluster( + self, + fp: NDArray[np.uint8], + merge_accept_fn: MergeAcceptFunction, + threshold: float, + packed_fp: tp.Optional[NDArray[np.uint8]] = None, + use_medoids: bool = False, + unpacked_fitted_fps: tp.Optional[NDArray[np.uint8]] = None, + k_search_idx: int = 0, + k_search: int = 1, + ) -> _Assignment: + """Find the highest similarity subcluster to a given fingerprint. + Return a 3-tuple [subcluster index, similarity, is_mergeable]""" + # In this case the node *must have* subclusters always + # Within this node, find the closest subcluster to the one to-be-inserted + if packed_fp is None: + packed_fp = pack_fingerprints(fp) + + is_leaf_node = next(iter(self._subclusters)).child + + if use_medoids and is_leaf_node: + # Only use medoids for the leafs + if unpacked_fitted_fps is None: + raise ValueError("Unpacked fitted fps required if using medoids") + members = [s.mol_indices for s in self._subclusters] + unpacked_medoids = _unpacked_medoids_from_cluster_members( + unpacked_fitted_fps, members + ) + packed_centrals = pack_fingerprints(unpacked_medoids) + else: + packed_centrals = self.packed_centroids + + sim_matrix = _jt_sim_arr_vec_packed(packed_centrals, packed_fp) + + # Previous method (recall ~30%) + if k_search_idx == 0: + closest_idx = np.argmax(sim_matrix) + else: + # Get the kth-closer + closest_idx = np.argsort(sim_matrix)[::-1][k_search_idx] + + closest_subclust = self._subclusters[closest_idx] + closest_node = closest_subclust.child + # Return a 3-tuple , , + if closest_node is None: + return _Assignment( + closest_subclust._index, + sim_matrix[closest_idx], + closest_subclust.fp_is_mergeable(fp, threshold, merge_accept_fn), + ) + # Recurse the naive version is bounded by an exponential increase, with tree + # depth k^D. If using k_search - 1 it is bounded by a factorial, which is a bit + # better, since the max number of searches decreases with tree depth + assignment = _Assignment(-1, 0.0, False) + k_search = max(1, k_search - 1) + _range = min(len(closest_node._subclusters), k_search) + for k_search_idx in range(_range): + _assignment = closest_node.find_highest_similarity_subcluster( + fp, + merge_accept_fn, + threshold, + use_medoids=use_medoids, + unpacked_fitted_fps=unpacked_fitted_fps, + k_search_idx=k_search_idx, + k_search=k_search, + ) + if _assignment.similarity >= assignment.similarity: + assignment = _assignment + return assignment + class _BFSubcluster: r"""Each subcluster in a BFNode is called a BFSubcluster. @@ -390,7 +466,7 @@ class _BFSubcluster: """ # NOTE: Slots deactivates __dict__, and thus reduces memory usage of python objects - __slots__ = ("_buffer", "packed_centroid", "child", "mol_indices") + __slots__ = ("_buffer", "packed_centroid", "child", "mol_indices", "_index") def __init__( self, @@ -446,6 +522,7 @@ def __init__( ) # Will be overwritten self.mol_indices = list(mol_indices) self.child: tp.Optional["_BFNode"] = None + self._index: int = -1 @property def unpacked_centroid(self) -> NDArray[np.uint8]: @@ -525,6 +602,21 @@ def merge_subcluster( return True return False + def fp_is_mergeable( + self, + fp: NDArray[np.uint8], + threshold: float, + merge_accept_fn: MergeAcceptFunction, + ) -> bool: + """Check if a cluster is worthy enough to be merged""" + old_n = self.n_samples + new_n = old_n + 1 + old_ls = self.linear_sum + # np.add with explicit dtype is safe from overflows, e.g. : + # np.add(np.uint8(255), np.uint8(255), dtype=np.uint16) = np.uint16(510) + new_ls = np.add(old_ls, fp, dtype=min_safe_uint(new_n)) + return merge_accept_fn(threshold, new_ls, new_n, old_ls, fp, old_n, 1) + class _CentroidsMolIds(tp.TypedDict): centroids: list[NDArray[np.uint8]] @@ -702,6 +794,121 @@ def set_merge( if branching_factor is not None: self.branching_factor = branching_factor + def assign( + self, + X: _Input | Path | str, + /, + input_is_packed: bool = True, + n_features: int | None = None, + max_fps: int | None = None, + sorted_subclusters_order: bool = True, + kind: str = "tree", + unpacked_fitted_fps: tp.Optional[NDArray[np.uint8]] = None, + use_medoids: bool = False, + k_search: int = 1, + ) -> tp.Any: + r""":meta private:""" + # Returns a pandas dataframe, but pandas import is triggered by this function + # to avoid memory usage if this function is unused, so the return value is + # untyped + import pandas as pd + + if kind not in ("tree", "flat"): + raise ValueError("Assignment must be one of tree|flat") + + if isinstance(X, (Path, str)): + X = _mmap_file_and_madvise_sequential(Path(X), max_fps=max_fps) + mmanager = _ArrayMemPagesManager.from_bb_input(X) + else: + X = X[:max_fps] + mmanager = _ArrayMemPagesManager.from_bb_input(X, can_release=False) + + n_features = _validate_n_features(X, input_is_packed, n_features) + # Start a new tree the first time this function is called + if not self.is_init: + raise ValueError("Create a tree before attempting assignments") + self._root = cast("_BFNode", self._root) # After init, this is not None + + # The array iterator either copies, un-sparsifies, or does nothing + # with the array rows, depending on the kind of X passed + arr_iterable = _get_array_iterable(X, input_is_packed, n_features) + arr_iterable = cast(tp.Iterable[NDArray[np.uint8]], arr_iterable) + + threshold = self.threshold + merge_accept_fn = self._merge_accept_fn + + assignments = [] + arr_idx = 0 + + bfs = self._get_leaf_bfs(sorted_subclusters_order) + + if kind == "flat": + if use_medoids: + if unpacked_fitted_fps is None: + raise ValueError("Unpacked fitted fps required if using medoids") + + all_packed_centrals = self.get_medoids( + fps=unpacked_fitted_fps, + input_is_packed=False, + sort=sorted_subclusters_order, + pack=True, + ) + else: + all_packed_centrals = np.stack([bf.packed_centroid for bf in bfs]) + for fp in arr_iterable: + packed_fp = pack_fingerprints(fp) + sim_matrix = _jt_sim_arr_vec_packed(all_packed_centrals, packed_fp) + closest_idx = np.argmax(sim_matrix) + sim = sim_matrix[closest_idx] + is_mergeable = bfs[closest_idx].fp_is_mergeable( + fp, threshold, merge_accept_fn + ) + assignments.append( + _Assignment(closest_idx.item() + 1, sim, is_mergeable) + ) + arr_idx += 1 + if mmanager.can_release and mmanager.should_release_curr_page(arr_idx): + mmanager.release_curr_page_and_update_addr() + return pd.DataFrame(assignments) + + if self._only_has_leaves: + raise ValueError( + "Internal nodes were released, assignments can't use 'tree' method" + ) + + # 'tree' branch + # Assign the leaf bf indices + for i, bf in enumerate(bfs, 1): + bf._index = i + + for fp in arr_iterable: + # NOTE: In this case we never need to split the root + assignment = _Assignment(-1, 0.0, False) + # TODO: Parallelizing this is pretty hard since it requires pickling the + # whole tree + _range = min(len(self._root._subclusters), k_search) + for k_search_idx in range(_range): + _assignment = self._root.find_highest_similarity_subcluster( + fp, + merge_accept_fn, + threshold, + use_medoids=use_medoids, + unpacked_fitted_fps=unpacked_fitted_fps, + k_search_idx=k_search_idx, + k_search=k_search, + ) + if _assignment.similarity >= assignment.similarity: + assignment = _assignment + assignments.append(assignment) + arr_idx += 1 + if mmanager.can_release and mmanager.should_release_curr_page(arr_idx): + mmanager.release_curr_page_and_update_addr() + + # Reset the leaf bf indices + for bf in bfs: + bf._index = -1 + return pd.DataFrame(assignments) + def fit( self, X: _Input | Path | str, @@ -933,26 +1140,11 @@ def get_medoids_mol_ids( if input_is_packed: fps = _unpack_fingerprints(fps, n_features=n_features) - cluster_medoids = self._unpacked_medoids_from_members(fps, cluster_members) + cluster_medoids = _unpacked_medoids_from_cluster_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], @@ -1425,6 +1617,21 @@ def _centrals_global_clustering( return predictor.fit_predict(centrals) + 1 +def _unpacked_medoids_from_cluster_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 + + # There are 4 cases here: # (1) The input is a scipy.sparse array # (2) The input is a list of dense arrays (nothing required) diff --git a/examples/assignment_script.py b/examples/assignment_script.py new file mode 100644 index 00000000..bae098c0 --- /dev/null +++ b/examples/assignment_script.py @@ -0,0 +1,91 @@ +import time +from pathlib import Path +import numpy as np +import matplotlib.pyplot as plt +import pickle +from bblean.fingerprints import unpack_fingerprints + +num_fps_to_assign = 1000 +directory = "test-assignment-rdkit5" +use_medoids = False +# very high branching factor is needed for good recall (~3500-5000) +# Probably a branching factor of ~sqrt(N_fps) * 3.5-5 is good +k = 4 # More than 4 is too slow, number of searches increases as factorial(k) +# k = 1 seems to always give ~30% recall, mostly independent of branching factor +# For example, k=4 starts with 4 searches, then each search splits into 3, then +# each splits into 2, for a total of 24 searches. k=1 or k=2 is ideal +# The actual growth is slower since some nodes have les branches, but it is bounded by +# the factorial +# +# Mergeable recall is similar to the normal recall (a bit higher in general) +# k = 4 is around 75% recal with branching factor of 3500-5000 +# k = 2 seems to be ~ 10x faster than flat search (but it is non-parallelizable) +# k = 4 is only ~2x / ~1.5x faster +# k = 10 is ~ 4x SLOWER and gets around 96% recall (spawns around 90 searches) + +fps = np.load("../10M/packed-fps-rdkit-uint8-7f343532.006.npy")[:num_fps_to_assign] +with open(f"./{directory}/bitbirch.pkl", mode="rb") as f: + tree = pickle.load(f) + # NOTE: Using tolerance-diamenter doesn't change the number of mergeable fps very + # much + # tree.set_merge("tolerance-diameter", tolerance=0.05) + # + # Similar time for medoids or centroids + if use_medoids: + unpacked_fitted_fps = unpack_fingerprints( + np.load(list(Path(f"./{directory}/input-fps/").glob("*.npy"))[0]) + ) + flat_assignments = tree.assign( + fps, kind="flat", use_medoids=True, unpacked_fitted_fps=unpacked_fitted_fps + ) + tree_assignments = tree.assign( + fps, + kind="tree", + use_medoids=True, + unpacked_fitted_fps=unpacked_fitted_fps, + k_search=k, + ) + else: + _start = time.perf_counter() + tree_assignments = tree.assign(fps, kind="tree", k_search=k) + print(f"Time elapsed tree: {time.perf_counter() - _start} s", flush=True) + _start = time.perf_counter() + flat_assignments = tree.assign(fps, kind="flat") + print(f"Time elapsed flat: {time.perf_counter() - _start} s", flush=True) + + correct_assignments = flat_assignments["is_mergeable"] + + # Clearly what is happening in the rdkit fps is that the first centroid has a *ton* + # of 1s and the rest have much less 1s + idxs = flat_assignments["cluster_label"] + values, counts = np.unique(idxs, return_counts=True) + fig, ax = plt.subplots() + ax.bar(values, counts, width=200, alpha=0.25) + ax.set_ylabel(r"Counts") + ax.set_xlabel(r"Label") + ax.set_title(f"Flat ({'medoid' if use_medoids else 'centroid'}) assignment") + plt.show(block=False) + + idxs = tree_assignments["cluster_label"] + values, counts = np.unique(idxs, return_counts=True) + fig, ax = plt.subplots() + ax.bar(values, counts, width=200, alpha=0.25) + ax.set_ylabel(r"Counts") + ax.set_xlabel(r"Label") + ax.set_title(f"Tree ({'medoid' if use_medoids else 'centroid'}) assignment") + plt.show() + + num_matches = ( + flat_assignments["cluster_label"] == tree_assignments["cluster_label"] + ).sum() + num_mergeable_matches = ( + flat_assignments["cluster_label"][correct_assignments] + == tree_assignments["cluster_label"][correct_assignments] + ).sum() + + print(f"Total: {len(flat_assignments)}") + print(f"Total mergeable: {correct_assignments.sum()}") + print(f"Recall: {num_matches * 100 / len(flat_assignments)}") + print( + f"Mergeable recall: {num_mergeable_matches * 100 / len(flat_assignments[correct_assignments])}" # noqa + )