From 438c2879b8e2af6f37d9a8ed82b5a49b1269b83f Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 11 Dec 2025 14:44:40 -0500 Subject: [PATCH 1/5] Proof of concept implementation of assignment --- bblean/bitbirch.py | 133 ++++++++++++++++++++++++++++++++++++++++++- bblean/cli.py | 2 +- bblean/multiround.py | 2 +- 3 files changed, 134 insertions(+), 3 deletions(-) diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index d9d2a9b4..5c0cbd05 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_idx: int + similarity: float + is_mergeable: bool + + class _BFNode: """Each node in a BitBirch tree is a _BFNode. @@ -356,6 +362,36 @@ 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, + ) -> _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) + + sim_matrix = _jt_sim_arr_vec_packed(self.packed_centroids, packed_fp) + closest_idx = np.argmax(sim_matrix) + 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 + return closest_node.find_highest_similarity_subcluster( + fp, merge_accept_fn, threshold, packed_fp + ) + class _BFSubcluster: r"""Each subcluster in a BFNode is called a BFSubcluster. @@ -390,7 +426,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 +482,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 +562,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 +754,85 @@ 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, + ) -> tp.Any: + r"""Assign a set of fingerprints according to the current BFTree + + Parameters + ---------- + + X : {array-like, sparse matrix} of shape (n_samples, n_features) + Input data. + + input_is_packed: bool + Whether the input fingerprints are packed + + n_features: int + Number of featurs of input fingerprints. Only required for packed inputs if + it is not a multiple of 8, otherwise it is redundant. + + Returns + ------- + Pandas DataFrame: with cluster_idx, similarity, is_mergeable columns. The + cluster index corresponds to sorted subclusters by default. + """ + # 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 + r""":meta private:""" + 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 self._only_has_leaves: + raise ValueError("Internal nodes were released, call reset() before fit()") + 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 + # Assign the leaf bf indices + bfs = self._get_leaf_bfs(sorted_subclusters_order) + for i, bf in enumerate(bfs): + bf._index = i + + for fp in arr_iterable: + # In this case we never need to split the root + assignment = self._root.find_highest_similarity_subcluster( + fp, merge_accept_fn, threshold + ) + 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, diff --git a/bblean/cli.py b/bblean/cli.py index 58aa0545..67cedc75 100644 --- a/bblean/cli.py +++ b/bblean/cli.py @@ -1101,7 +1101,7 @@ def _run( console.print("Can't save tree for non-lean variants", style="red") else: # TODO: Find alternative solution - tree.save_pickle(out_dir / "bitbirch.pkl") + tree.save(out_dir / "bitbirch.pkl") if variant == "lean": tree.delete_internal_nodes() # Dump outputs (peak memory, timings, config, cluster ids) diff --git a/bblean/multiround.py b/bblean/multiround.py index dfa08abf..3210e9f4 100644 --- a/bblean/multiround.py +++ b/bblean/multiround.py @@ -299,7 +299,7 @@ def __call__(self, batch_info: tuple[str, tp.Sequence[tuple[Path, Path]]]) -> No # Save clusters and exit if self.save_tree: # TODO: Find alternative solution - tree.save_pickle(self.out_dir / "bitbirch.pkl") + tree.save(self.out_dir / "bitbirch.pkl") tree.delete_internal_nodes() if self.save_centroids: output = tree.get_centroids_mol_ids() From 789e97868779cd08c11f5004b5832d848f9da2d0 Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 11 Dec 2025 20:04:40 -0500 Subject: [PATCH 2/5] Add k_search option --- bblean/_py_similarity.py | 10 +-- bblean/bitbirch.py | 162 +++++++++++++++++++++++++++++-------- bblean/csrc/similarity.cpp | 103 +++++++++++++++++------ bblean/similarity.py | 65 ++++++++++++--- tests/test_similarity.py | 11 ++- 5 files changed, 270 insertions(+), 81 deletions(-) diff --git a/bblean/_py_similarity.py b/bblean/_py_similarity.py index 10355cfe..33d450db 100644 --- a/bblean/_py_similarity.py +++ b/bblean/_py_similarity.py @@ -76,18 +76,10 @@ def jt_compl_isim( warnings.warn(msg, RuntimeWarning, stacklevel=2) return np.full(len(fps), fill_value=np.nan, dtype=np.float64) linear_sum = np.sum(fps, axis=0) - n_objects = len(fps) - 1 comp_sims = [jt_isim_from_sum(linear_sum - fp, n_objects) for fp in fps] - return np.array(comp_sims, dtype=np.float64) -def _jt_isim_medoid_index( - fps: NDArray[np.uint8], input_is_packed: bool = True, n_features: int | None = None -) -> int: - return np.argmin(jt_compl_isim(fps, input_is_packed, n_features)).item() - - def jt_isim_medoid( fps: NDArray[np.uint8], input_is_packed: bool = True, @@ -110,7 +102,7 @@ def jt_isim_medoid( if len(fps) < 3: idx = 0 # Medoid undefined for sets of 3 or more fingerprints else: - idx = _jt_isim_medoid_index(fps, input_is_packed=False) + idx = np.argmin(jt_compl_isim(fps, input_is_packed, n_features)).item() m = fps[idx] if pack: return idx, pack_fingerprints(m) diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index 733cee9f..1a148065 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -212,7 +212,7 @@ def _split_node(node: "_BFNode") -> tuple["_BFSubcluster", "_BFSubcluster"]: class _Assignment(tp.NamedTuple): - cluster_idx: int + cluster_label: int similarity: float is_mergeable: bool @@ -368,6 +368,10 @@ def find_highest_similarity_subcluster( 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]""" @@ -376,8 +380,29 @@ def find_highest_similarity_subcluster( if packed_fp is None: packed_fp = pack_fingerprints(fp) - sim_matrix = _jt_sim_arr_vec_packed(self.packed_centroids, packed_fp) - closest_idx = np.argmax(sim_matrix) + 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 , , @@ -387,10 +412,25 @@ def find_highest_similarity_subcluster( sim_matrix[closest_idx], closest_subclust.fp_is_mergeable(fp, threshold, merge_accept_fn), ) - # Recurse - return closest_node.find_highest_similarity_subcluster( - fp, merge_accept_fn, threshold, packed_fp - ) + # 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: @@ -762,7 +802,13 @@ def assign( 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: + # TODO: Do this with medoids instead of centroids + r"""Assign a set of fingerprints according to the current BFTree Parameters @@ -780,7 +826,7 @@ def assign( Returns ------- - Pandas DataFrame: with cluster_idx, similarity, is_mergeable columns. The + Pandas DataFrame: with cluster_label, similarity, is_mergeable columns. The cluster index corresponds to sorted subclusters by default. """ # Returns a pandas dataframe, but pandas import is triggered by this function @@ -788,7 +834,9 @@ def assign( # untyped import pandas as pd - r""":meta private:""" + 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) @@ -798,8 +846,6 @@ def assign( n_features = _validate_n_features(X, input_is_packed, n_features) # Start a new tree the first time this function is called - if self._only_has_leaves: - raise ValueError("Internal nodes were released, call reset() before fit()") 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 @@ -814,16 +860,66 @@ def assign( assignments = [] arr_idx = 0 - # Assign the leaf bf indices + bfs = self._get_leaf_bfs(sorted_subclusters_order) - for i, bf in enumerate(bfs): + + 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: - # In this case we never need to split the root - assignment = self._root.find_highest_similarity_subcluster( - fp, merge_accept_fn, threshold - ) + # 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): @@ -1065,26 +1161,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], @@ -1557,6 +1638,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/bblean/csrc/similarity.cpp b/bblean/csrc/similarity.cpp index 1fd6f3f5..69cccaa8 100644 --- a/bblean/csrc/similarity.cpp +++ b/bblean/csrc/similarity.cpp @@ -300,6 +300,75 @@ double jt_isim_from_sum(const CArrayForcecast& linear_sum, return a / ((a + (n_objects * sum_kq)) - sum_kqsq); } +// NOTE: This is only *slightly* faster for C++ than numpy, **only if the +// array is uint8_t** if the array is uint64 already, it is slower +template +py::array_t add_rows(const CArrayForcecast& arr) { + if (arr.ndim() != 2) { + throw std::runtime_error("Input array must be 2-dimensional"); + } + auto arr_ptr = arr.data(); + auto out = py::array_t(arr.shape(1)); + auto out_ptr = out.mutable_data(); + std::memset(out_ptr, 0, out.nbytes()); + py::ssize_t n_samples = arr.shape(0); + py::ssize_t n_features = arr.shape(1); + // Check GCC / CLang vectorize this + for (py::ssize_t i = 0; i < n_samples; ++i) { + const uint8_t* arr_row_ptr = arr_ptr + i * n_features; + for (py::ssize_t j = 0; j < n_features; ++j) { + out_ptr[j] += static_cast(arr_row_ptr[j]); + } + } + return out; +} +py::array_t _nochecks_jt_compl_isim_unpacked_u8( + const py::array_t& fps) { + py::ssize_t n_objects = fps.shape(0); + py::ssize_t n_features = fps.shape(1); + auto out = py::array_t(n_objects); + auto out_ptr = out.mutable_data(); + + if (n_objects < 3) { + PyErr_WarnEx(PyExc_RuntimeWarning, + "Invalid num fps in compl_isim. Expected n_objects >= 3", + 1); + for (py::ssize_t i{0}; i != n_objects; ++i) { + out_ptr[i] = std::numeric_limits::quiet_NaN(); + } + return out; + } + + auto linear_sum = add_rows(fps); + auto ls_cptr = linear_sum.data(); + + py::array_t shifted_linear_sum(n_features); + auto shifted_ls_ptr = shifted_linear_sum.mutable_data(); + + auto in_cptr = fps.data(); + for (py::ssize_t i{0}; i != n_objects; ++i) { + for (py::ssize_t j{0}; j != n_features; ++j) { + shifted_ls_ptr[j] = ls_cptr[j] - in_cptr[i * n_features + j]; + } + // For all compl isim N is n_objects - 1 + out_ptr[i] = jt_isim_from_sum(shifted_linear_sum, n_objects - 1); + } + return out; +} + +py::array_t jt_compl_isim( + const CArrayForcecast& fps, bool input_is_packed = true, + std::optional n_features_opt = std::nullopt) { + if (fps.ndim() != 2) { + throw std::runtime_error("fps arr must be 2D"); + } + if (input_is_packed) { + return _nochecks_jt_compl_isim_unpacked_u8( + _nochecks_unpack_fingerprints_2d(fps, n_features_opt)); + } + return _nochecks_jt_compl_isim_unpacked_u8(fps); +} + // Contraint: T must be uint64_t or uint8_t template void _calc_arr_vec_jt(const py::array_t& arr, @@ -372,33 +441,10 @@ py::array_t jt_sim_packed_precalc_cardinalities( } py::array_t _jt_sim_arr_vec_packed(const py::array_t& arr, - const py::array_t& vec) { + const py::array_t& vec) { return jt_sim_packed_precalc_cardinalities(arr, vec, _popcount_2d(arr)); } -// NOTE: This is only *slightly* faster for C++ than numpy, **only if the -// array is uint8_t** if the array is uint64 already, it is slower -template -py::array_t add_rows(const CArrayForcecast& arr) { - if (arr.ndim() != 2) { - throw std::runtime_error("Input array must be 2-dimensional"); - } - auto arr_ptr = arr.data(); - auto out = py::array_t(arr.shape(1)); - auto out_ptr = out.mutable_data(); - std::memset(out_ptr, 0, out.nbytes()); - py::ssize_t n_samples = arr.shape(0); - py::ssize_t n_features = arr.shape(1); - // Check GCC / CLang vectorize this - for (py::ssize_t i = 0; i < n_samples; ++i) { - const uint8_t* arr_row_ptr = arr_ptr + i * n_features; - for (py::ssize_t j = 0; j < n_features; ++j) { - out_ptr[j] += static_cast(arr_row_ptr[j]); - } - } - return out; -} - double jt_isim_unpacked_u8(const CArrayForcecast& arr) { return jt_isim_from_sum(add_rows(arr), arr.shape(0)); } @@ -406,8 +452,9 @@ double jt_isim_unpacked_u8(const CArrayForcecast& arr) { double jt_isim_packed_u8( const CArrayForcecast& arr, std::optional n_features_opt = std::nullopt) { - return jt_isim_from_sum(add_rows(unpack_fingerprints(arr, n_features_opt)), - arr.shape(0)); + return jt_isim_from_sum( + add_rows(unpack_fingerprints(arr, n_features_opt)), + arr.shape(0)); } py::tuple jt_most_dissimilar_packed( @@ -510,6 +557,10 @@ PYBIND11_MODULE(_cpp_similarity, m) { m.def("jt_isim_unpacked_u8", &jt_isim_unpacked_u8, "iSIM Tanimoto calculation", py::arg("arr")); + m.def("jt_compl_isim", &jt_compl_isim, "Complementary iSIM tanimoto", + py::arg("fps"), py::arg("input_is_packed") = true, + py::arg("n_features") = std::nullopt); + m.def("_jt_sim_arr_vec_packed", &_jt_sim_arr_vec_packed, "Tanimoto similarity between a matrix of packed fps and a single " "packed fp", diff --git a/bblean/similarity.py b/bblean/similarity.py index c5660def..50cb1a96 100644 --- a/bblean/similarity.py +++ b/bblean/similarity.py @@ -34,12 +34,8 @@ "jt_sim_matrix_packed", ] -from bblean._py_similarity import ( - centroid_from_sum, - centroid, - jt_compl_isim, - jt_isim_medoid, -) +from bblean._py_similarity import centroid_from_sum, centroid +from bblean.fingerprints import pack_fingerprints, unpack_fingerprints # jt_isim_packed and jt_isim_unpacked are not exposed, only used within functions for # speed @@ -49,6 +45,7 @@ jt_isim_from_sum, jt_isim_unpacked, jt_isim_packed, + jt_compl_isim, _jt_sim_arr_vec_packed, jt_most_dissimilar_packed, ) @@ -56,11 +53,13 @@ try: from bblean._cpp_similarity import ( # type: ignore jt_isim_from_sum, - _jt_sim_arr_vec_packed, jt_isim_unpacked_u8, jt_isim_packed_u8, + jt_compl_isim, # TODO: Does it need wrappers for non-uint8? + _jt_sim_arr_vec_packed, jt_most_dissimilar_packed, - unpack_fingerprints, + # Needed for wrappers + unpack_fingerprints as _unpack_fingerprints, ) # Wrap these two since doing @@ -80,7 +79,7 @@ def jt_isim_packed( # type: ignore if arr.dtype == np.uint64: return jt_isim_from_sum( np.sum( - unpack_fingerprints(arr, n_features), # type: ignore + _unpack_fingerprints(arr, n_features), # type: ignore axis=0, dtype=np.uint64, ), @@ -93,6 +92,7 @@ def jt_isim_packed( # type: ignore jt_isim_from_sum, jt_isim_unpacked, jt_isim_packed, + jt_compl_isim, _jt_sim_arr_vec_packed, jt_most_dissimilar_packed, ) @@ -103,6 +103,35 @@ def jt_isim_packed( # type: ignore ) +def jt_isim_medoid( + fps: NDArray[np.uint8], + input_is_packed: bool = True, + n_features: int | None = None, + pack: bool = True, +) -> tuple[int, NDArray[np.uint8]]: + r"""Calculate the (Tanimoto) medoid of a set of fingerprints, using iSIM + + Returns both the index of the medoid in the input array and the medoid itself + + .. note:: + Returns the first (or only) fingerprint for array of size 2 and 1 respectively. + Raises ValueError for arrays of size 0 + + """ + if not fps.size: + raise ValueError("Size of fingerprints set must be > 0") + if input_is_packed: + fps = unpack_fingerprints(fps, n_features) + if len(fps) < 3: + idx = 0 # Medoid undefined for sets of 3 or more fingerprints + else: + idx = np.argmin(jt_compl_isim(fps, input_is_packed, n_features)).item() + m = fps[idx] + if pack: + return idx, pack_fingerprints(m) + return idx, m + + def jt_isim( fps: NDArray[np.integer], input_is_packed: bool = True, @@ -149,7 +178,11 @@ def jt_isim_diameter( r"""Calculate the Tanimoto diameter of a set of fingerprints""" return jt_isim_diameter_from_sum( np.sum( - unpack_fingerprints(arr, n_features) if input_is_packed else arr, + ( + unpack_fingerprints(arr.astype(np.uint8, copy=False), n_features) + if input_is_packed + else arr + ), axis=0, dtype=np.uint64, ), # type: ignore @@ -165,7 +198,11 @@ def jt_isim_radius( r"""Calculate the Tanimoto radius of a set of fingerprints""" return jt_isim_radius_from_sum( np.sum( - unpack_fingerprints(arr, n_features) if input_is_packed else arr, + ( + unpack_fingerprints(arr.astype(np.uint8, copy=False), n_features) + if input_is_packed + else arr + ), axis=0, dtype=np.uint64, ), # type: ignore @@ -181,7 +218,11 @@ def jt_isim_radius_compl( r"""Calculate the complement of the Tanimoto radius of a set of fingerprints""" return jt_isim_radius_compl_from_sum( np.sum( - unpack_fingerprints(arr, n_features) if input_is_packed else arr, + ( + unpack_fingerprints(arr.astype(np.uint8, copy=False), n_features) + if input_is_packed + else arr + ), axis=0, dtype=np.uint64, ), # type: ignore diff --git a/tests/test_similarity.py b/tests/test_similarity.py index b39193b8..b1879486 100644 --- a/tests/test_similarity.py +++ b/tests/test_similarity.py @@ -254,8 +254,12 @@ def test_jt_compl_isim() -> None: with pytest.warns(RuntimeWarning): _ = pysim.jt_compl_isim(fps) + with pytest.warns(RuntimeWarning): + _ = csim.jt_compl_isim(fps) + fps = make_fake_fingerprints(10, seed=17408390758220920002, pack=False) - assert pysim.jt_compl_isim(fps).tolist() == snapshot( + output = pysim.jt_compl_isim(fps).tolist() + assert output == snapshot( [ 0.20256457907452147, 0.24748926949201983, @@ -269,10 +273,15 @@ def test_jt_compl_isim() -> None: 0.2225069540267648, ] ) + assert csim.jt_compl_isim(fps).tolist() == output assert ( pysim.jt_compl_isim(np.zeros((10, 512), dtype=np.uint8)) == np.ones(10, dtype=np.float64) ).all() + assert ( + csim.jt_compl_isim(np.zeros((10, 512), dtype=np.uint8)) + == np.ones(10, dtype=np.float64) + ).all() def test_jt_isim_medoid() -> None: From 41226ce47453402ca8a8fc5ce655ec9c136dfbfa Mon Sep 17 00:00:00 2001 From: ipickering Date: Thu, 11 Dec 2025 20:06:16 -0500 Subject: [PATCH 3/5] Add assignment script to examples --- examples/assignment_script.py | 91 +++++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 examples/assignment_script.py 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 + ) From 942c332f05eb1a98919d5d11c8cfa5cd17f0f5f0 Mon Sep 17 00:00:00 2001 From: ipickering Date: Sat, 13 Dec 2025 18:59:47 -0500 Subject: [PATCH 4/5] Run csim compl tests conditionally --- tests/test_similarity.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/test_similarity.py b/tests/test_similarity.py index b1879486..6668a2ab 100644 --- a/tests/test_similarity.py +++ b/tests/test_similarity.py @@ -253,9 +253,9 @@ def test_jt_compl_isim() -> None: fps = make_fake_fingerprints(2, seed=17408390758220920002, pack=False) with pytest.warns(RuntimeWarning): _ = pysim.jt_compl_isim(fps) - - with pytest.warns(RuntimeWarning): - _ = csim.jt_compl_isim(fps) + if CSIM_AVAIL: + with pytest.warns(RuntimeWarning): + _ = csim.jt_compl_isim(fps) fps = make_fake_fingerprints(10, seed=17408390758220920002, pack=False) output = pysim.jt_compl_isim(fps).tolist() @@ -273,15 +273,17 @@ def test_jt_compl_isim() -> None: 0.2225069540267648, ] ) - assert csim.jt_compl_isim(fps).tolist() == output + if CSIM_AVAIL: + assert csim.jt_compl_isim(fps).tolist() == output assert ( pysim.jt_compl_isim(np.zeros((10, 512), dtype=np.uint8)) == np.ones(10, dtype=np.float64) ).all() - assert ( - csim.jt_compl_isim(np.zeros((10, 512), dtype=np.uint8)) - == np.ones(10, dtype=np.float64) - ).all() + if CSIM_AVAIL: + assert ( + csim.jt_compl_isim(np.zeros((10, 512), dtype=np.uint8)) + == np.ones(10, dtype=np.float64) + ).all() def test_jt_isim_medoid() -> None: From 011004e40b76cea66331f3fa179b10ea31e73341 Mon Sep 17 00:00:00 2001 From: ipickering Date: Tue, 16 Dec 2025 16:11:01 -0500 Subject: [PATCH 5/5] Remove docstring --- bblean/bitbirch.py | 23 +---------------------- 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/bblean/bitbirch.py b/bblean/bitbirch.py index 1a148065..701389a7 100644 --- a/bblean/bitbirch.py +++ b/bblean/bitbirch.py @@ -807,28 +807,7 @@ def assign( use_medoids: bool = False, k_search: int = 1, ) -> tp.Any: - # TODO: Do this with medoids instead of centroids - - r"""Assign a set of fingerprints according to the current BFTree - - Parameters - ---------- - - X : {array-like, sparse matrix} of shape (n_samples, n_features) - Input data. - - input_is_packed: bool - Whether the input fingerprints are packed - - n_features: int - Number of featurs of input fingerprints. Only required for packed inputs if - it is not a multiple of 8, otherwise it is redundant. - - Returns - ------- - Pandas DataFrame: with cluster_label, similarity, is_mergeable columns. The - cluster index corresponds to sorted subclusters by default. - """ + 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