bblean 0.6.0b1__cp312-cp312-macosx_10_13_universal2.whl

bblean/multiround.py

@@ -0,0 +1,489 @@
+# BitBIRCH-Lean Python Package: An open-source clustering module based on iSIM.
+#
+# If you find this software useful please cite the following articles:
+# - BitBIRCH: efficient clustering of large molecular libraries:
+#   https://doi.org/10.1039/D5DD00030K
+# - BitBIRCH Clustering Refinement Strategies:
+#   https://doi.org/10.1021/acs.jcim.5c00627
+# - BitBIRCh-Lean:
+#   (preprint) https://www.biorxiv.org/content/10.1101/2025.10.22.684015v1
+#
+# Copyright (C) 2025  The Miranda-Quintana Lab and other BitBirch developers, comprised
+# exclusively by:
+# - Ramon Alain Miranda Quintana <ramirandaq@gmail.com>, <quintana@chem.ufl.edu>
+# - Krisztina Zsigmond <kzsigmond@ufl.edu>
+# - Ignacio Pickering <ipickering@chem.ufl.edu>
+# - Kenneth Lopez Perez <klopezperez@chem.ufl.edu>
+# - Miroslav Lzicar <miroslav.lzicar@deepmedchem.com>
+#
+# Authors of this file are:
+# - Ramon Alain Miranda Quintana <ramirandaq@gmail.com>, <quintana@chem.ufl.edu>
+# - Ignacio Pickering <ipickering@chem.ufl.edu>
+#
+# This program is free software: you can redistribute it and/or modify it under the
+# terms of the GNU General Public License as published by the Free Software Foundation,
+# version 3 (SPDX-License-Identifier: GPL-3.0-only).
+#
+# Portions of ./bblean/bitbirch.py are licensed under the BSD 3-Clause License
+# Copyright (c) 2007-2024 The scikit-learn developers. All rights reserved.
+# (SPDX-License-Identifier: BSD-3-Clause). Copies or reproductions of code in the
+# ./bblean/bitbirch.py file must in addition adhere to the BSD-3-Clause license terms. A
+# copy of the BSD-3-Clause license can be located at the root of this repository, under
+# ./LICENSES/BSD-3-Clause.txt.
+#
+# Portions of ./bblean/bitbirch.py were previously licensed under the LGPL 3.0
+# license (SPDX-License-Identifier: LGPL-3.0-only), they are relicensed in this program
+# as GPL-3.0, with permission of all original copyright holders:
+# - Ramon Alain Miranda Quintana <ramirandaq@gmail.com>, <quintana@chem.ufl.edu>
+# - Vicky (Vic) Jung <jungvicky@ufl.edu>
+# - Kenneth Lopez Perez <klopezperez@chem.ufl.edu>
+# - Kate Huddleston <kdavis2@chem.ufl.edu>
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
+# PARTICULAR PURPOSE.  See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with this
+# program. This copy can be located at the root of this repository, under
+# ./LICENSES/GPL-3.0-only.txt.  If not, see <http://www.gnu.org/licenses/gpl-3.0.html>.
+r"""Multi-round BitBirch workflow for clustering huge datasets in parallel"""
+import sys
+import math
+import pickle
+import typing as tp
+import multiprocessing as mp
+from pathlib import Path
+
+from rich.console import Console
+import numpy as np
+from numpy.typing import NDArray
+
+
+from bblean._console import get_console
+from bblean._timer import Timer
+from bblean._config import DEFAULTS
+from bblean.utils import batched
+from bblean.bitbirch import BitBirch
+from bblean.fingerprints import _get_fps_file_num
+
+__all__ = ["run_multiround_bitbirch"]
+
+
+# Save a list of numpy arrays into a single array in a streaming fashion, avoiding
+# stacking them in memory
+def _numpy_streaming_save(
+    fp_list: list[NDArray[np.integer]] | NDArray[np.integer], path: Path | str
+) -> None:
+    first_arr = np.ascontiguousarray(fp_list[0])
+    header = np.lib.format.header_data_from_array_1_0(first_arr)
+    header["shape"] = (len(fp_list), len(first_arr))
+    path = Path(path)
+    if not path.suffix:
+        path = path.with_suffix(".npy")
+    with open(path, "wb") as f:
+        np.lib.format.write_array_header_1_0(f, header)
+        for arr in fp_list:
+            np.ascontiguousarray(arr).tofile(f)
+
+
+# Glob and sort by uint bits and label, if a console is passed then the number of output
+# files is printed
+def _get_prev_round_buf_and_mol_idxs_files(
+    path: Path, round_idx: int, console: Console | None = None
+) -> list[tuple[Path, Path]]:
+    path = Path(path)
+    # TODO: Important: What should be the logic for batching? currently there doesn't
+    # seem to be much logic for grouping the files
+    buf_files = sorted(path.glob(f"round-{round_idx - 1}-bufs*.npy"))
+    idx_files = sorted(path.glob(f"round-{round_idx - 1}-idxs*.pkl"))
+    if console is not None:
+        console.print(f"    - Collected {len(buf_files)} buffer-index file pairs")
+    return list(zip(buf_files, idx_files))
+
+
+def _sort_batch(b: tp.Sequence[tuple[Path, Path]]) -> tuple[tuple[Path, Path], ...]:
+    return tuple(
+        sorted(
+            b,
+            key=lambda b: int(b[0].name.split("uint")[-1].split(".")[0]),
+            reverse=True,
+        )
+    )
+
+
+def _chunk_file_pairs_in_batches(
+    file_pairs: tp.Sequence[tuple[Path, Path]],
+    bin_size: int,
+    console: Console | None = None,
+) -> list[tuple[str, tuple[tuple[Path, Path], ...]]]:
+    z = len(str(math.ceil(len(file_pairs) / bin_size)))
+    # Within each batch, sort the files by starting with the uint16 files, followed by
+    # uint8 files, this helps that (approximately) the largest clusters are fitted first
+    # which may improve final cluster quality
+    batches = [
+        (str(i).zfill(z), _sort_batch(b))
+        for i, b in enumerate(batched(file_pairs, bin_size))
+    ]
+    if console is not None:
+        console.print(f"    - Chunked files into {len(batches)} batches")
+    return batches
+
+
+def _save_bufs_and_mol_idxs(
+    out_dir: Path,
+    fps_bfs: dict[str, tp.Any],
+    mols_bfs: dict[str, tp.Any],
+    label: str,
+    round_idx: int,
+) -> None:
+    for dtype, buf_list in fps_bfs.items():
+        suffix = f".label-{label}-{dtype.replace('8', '08')}"
+        _numpy_streaming_save(buf_list, out_dir / f"round-{round_idx}-bufs{suffix}.npy")
+        with open(out_dir / f"round-{round_idx}-idxs{suffix}.pkl", mode="wb") as f:
+            pickle.dump(mols_bfs[dtype], f)
+
+
+class _InitialRound:
+    def __init__(
+        self,
+        branching_factor: int,
+        threshold: float,
+        tolerance: float,
+        out_dir: Path | str,
+        refinement_before_midsection: str,
+        refine_threshold_change: float,
+        refine_merge_criterion: str,
+        n_features: int | None = None,
+        max_fps: int | None = None,
+        merge_criterion: str = DEFAULTS.merge_criterion,
+        input_is_packed: bool = True,
+    ) -> None:
+        self.n_features = n_features
+        self.refinement_before_midsection = refinement_before_midsection
+        if refinement_before_midsection not in ["full", "split", "none"]:
+            raise ValueError(f"Unknown refinement kind {refinement_before_midsection}")
+        self.branching_factor = branching_factor
+        self.threshold = threshold
+        self.tolerance = tolerance
+        self.out_dir = Path(out_dir)
+        self.max_fps = max_fps
+        self.merge_criterion = merge_criterion
+        self.refine_merge_criterion = refine_merge_criterion
+        self.input_is_packed = input_is_packed
+        self.refine_threshold_change = refine_threshold_change
+
+    def __call__(self, file_info: tuple[str, Path, int, int]) -> None:
+        file_label, fp_file, start_idx, end_idx = file_info
+
+        # First fit the fps in each process, in parallel.
+        # `reinsert_indices` required to keep track of mol idxs in different processes.
+        tree = BitBirch(
+            branching_factor=self.branching_factor,
+            threshold=self.threshold,
+            merge_criterion=self.merge_criterion,
+        )
+        range_ = range(start_idx, end_idx)
+        tree.fit(
+            fp_file,
+            reinsert_indices=range_,
+            n_features=self.n_features,
+            input_is_packed=self.input_is_packed,
+            max_fps=self.max_fps,
+        )
+        # Extract the BitFeatures of the leaves, breaking the largest cluster(s) apart,
+        # to prepare for refinement
+        tree.delete_internal_nodes()
+        if self.refinement_before_midsection == "none":
+            fps_bfs, mols_bfs = tree._bf_to_np()
+        elif self.refinement_before_midsection in ["split", "full"]:
+            fps_bfs, mols_bfs = tree._bf_to_np_refine(fp_file, initial_mol=start_idx)
+            if self.refinement_before_midsection == "full":
+                # Finish the first refinement step internally in this round
+                tree.reset()
+                tree.set_merge(
+                    self.refine_merge_criterion,
+                    tolerance=self.tolerance,
+                    threshold=self.threshold + self.refine_threshold_change,
+                )
+                for bufs, mol_idxs in zip(fps_bfs.values(), mols_bfs.values()):
+                    tree._fit_buffers(bufs, reinsert_index_seqs=mol_idxs)
+                    del mol_idxs
+                    del bufs
+
+                tree.delete_internal_nodes()
+                fps_bfs, mols_bfs = tree._bf_to_np()
+
+        _save_bufs_and_mol_idxs(self.out_dir, fps_bfs, mols_bfs, file_label, 1)
+
+
+class _TreeMergingRound:
+    def __init__(
+        self,
+        branching_factor: int,
+        threshold: float,
+        tolerance: float,
+        round_idx: int,
+        out_dir: Path | str,
+        split_largest_cluster: bool,
+        criterion: str,
+        all_fp_paths: tp.Sequence[Path] = (),
+    ) -> None:
+        self.all_fp_paths = list(all_fp_paths)
+        self.branching_factor = branching_factor
+        self.threshold = threshold
+        self.tolerance = tolerance
+        self.round_idx = round_idx
+        self.out_dir = Path(out_dir)
+        self.split_largest_cluster = split_largest_cluster
+        self.criterion = criterion
+
+    def __call__(self, batch_info: tuple[str, tp.Sequence[tuple[Path, Path]]]) -> None:
+        batch_label, batch_path_pairs = batch_info
+        tree = BitBirch(
+            branching_factor=self.branching_factor,
+            threshold=self.threshold,
+            merge_criterion=self.criterion,
+            tolerance=self.tolerance,
+        )
+        # Rebuild a tree, inserting all BitFeatures from the corresponding batch
+        for buf_path, idx_path in batch_path_pairs:
+            with open(idx_path, "rb") as f:
+                mol_idxs = pickle.load(f)
+            tree._fit_buffers(buf_path, reinsert_index_seqs=mol_idxs)
+            del mol_idxs
+
+        # Either do a refinement step, or fetch and save the bufs and idxs for the next
+        # round
+        tree.delete_internal_nodes()
+        if self.split_largest_cluster:
+            fps_bfs, mols_bfs = tree._bf_to_np_refine(self.all_fp_paths)
+        else:
+            fps_bfs, mols_bfs = tree._bf_to_np()
+        _save_bufs_and_mol_idxs(
+            self.out_dir, fps_bfs, mols_bfs, batch_label, self.round_idx
+        )
+
+
+class _FinalTreeMergingRound(_TreeMergingRound):
+    def __init__(
+        self,
+        branching_factor: int,
+        threshold: float,
+        tolerance: float,
+        criterion: str,
+        out_dir: Path | str,
+        save_tree: bool,
+        save_centroids: bool,
+    ) -> None:
+        super().__init__(
+            branching_factor, threshold, tolerance, -1, out_dir, False, criterion, ()
+        )
+        self.save_tree = save_tree
+        self.save_centroids = save_centroids
+
+    def __call__(self, batch_info: tuple[str, tp.Sequence[tuple[Path, Path]]]) -> None:
+        batch_path_pairs = batch_info[1]
+        tree = BitBirch(
+            branching_factor=self.branching_factor,
+            threshold=self.threshold,
+            merge_criterion=self.criterion,
+            tolerance=self.tolerance,
+        )
+        # Rebuild a tree, inserting all BitFeatures from the corresponding batch
+        for buf_path, idx_path in batch_path_pairs:
+            with open(idx_path, "rb") as f:
+                mol_idxs = pickle.load(f)
+            tree._fit_buffers(buf_path, reinsert_index_seqs=mol_idxs)
+            del mol_idxs
+
+        # Save clusters and exit
+        if self.save_tree:
+            # TODO: BitBIRCH is highly recursive. pickling may crash python,
+            # an alternative solution would be better
+            _old_limit = sys.getrecursionlimit()
+            sys.setrecursionlimit(100_000)
+            with open(self.out_dir / "bitbirch.pkl", mode="wb") as f:
+                pickle.dump(tree, f)
+            sys.setrecursionlimit(_old_limit)
+        tree.delete_internal_nodes()
+        if self.save_centroids:
+            output = tree.get_centroids_mol_ids()
+            with open(self.out_dir / "clusters.pkl", mode="wb") as f:
+                pickle.dump(output["mol_ids"], f)
+            with open(self.out_dir / "cluster-centroids-packed.pkl", mode="wb") as f:
+                pickle.dump(output["centroids"], f)
+        else:
+            with open(self.out_dir / "clusters.pkl", mode="wb") as f:
+                pickle.dump(tree.get_cluster_mol_ids(), f)
+
+
+# Create a list of tuples of labels, file paths and start-end idxs
+def _get_files_range_tuples(
+    files: tp.Sequence[Path],
+) -> list[tuple[str, Path, int, int]]:
+    running_idx = 0
+    files_info = []
+    z = len(str(len(files)))
+    for i, file in enumerate(files):
+        start_idx = running_idx
+        end_idx = running_idx + _get_fps_file_num(file)
+        files_info.append((str(i).zfill(z), file, start_idx, end_idx))
+        running_idx = end_idx
+    return files_info
+
+
+# NOTE: 'full_refinement_before_midsection' indicates if the refinement of the batches
+# is fully done after the tree-merging rounds, or if the data is only split before the
+# tree-merging rounds
+def run_multiround_bitbirch(
+    input_files: tp.Sequence[Path],
+    out_dir: Path,
+    n_features: int | None = None,
+    input_is_packed: bool = True,
+    num_initial_processes: int = 10,
+    num_midsection_processes: int | None = None,
+    initial_merge_criterion: str = DEFAULTS.merge_criterion,
+    branching_factor: int = DEFAULTS.branching_factor,
+    threshold: float = DEFAULTS.threshold,
+    midsection_threshold_change: float = DEFAULTS.refine_threshold_change,
+    tolerance: float = DEFAULTS.tolerance,
+    # Advanced
+    num_midsection_rounds: int = 1,
+    bin_size: int = 10,
+    max_tasks_per_process: int = 1,
+    refinement_before_midsection: str = "full",
+    split_largest_after_each_midsection_round: bool = False,
+    midsection_merge_criterion: str = DEFAULTS.refine_merge_criterion,
+    final_merge_criterion: str | None = None,
+    mp_context: tp.Any = None,
+    save_tree: bool = False,
+    save_centroids: bool = True,
+    # Debug
+    max_fps: int | None = None,
+    verbose: bool = False,
+    cleanup: bool = True,
+) -> Timer:
+    r"""Perform (possibly parallel) multi-round BitBirch clustering
+
+    .. warning::
+
+        The functionality provided by this function is stable, but its API
+        (the arguments it takes and its return values) may change in the future.
+    """
+    if final_merge_criterion is None:
+        final_merge_criterion = midsection_merge_criterion
+    if mp_context is None:
+        mp_context = mp.get_context("forkserver" if sys.platform == "linux" else None)
+    # Returns timing and for the different rounds
+    # TODO: Also return peak-rss
+    console = get_console(silent=not verbose)
+
+    if num_midsection_processes is None:
+        num_midsection_processes = num_initial_processes
+    else:
+        # Sanity check
+        if num_midsection_processes > num_initial_processes:
+            raise ValueError("Num. midsection procs. must be <= num. initial processes")
+
+    # Common params to all rounds BitBIRCH
+    common_kwargs: dict[str, tp.Any] = dict(
+        branching_factor=branching_factor,
+        tolerance=tolerance,
+        out_dir=out_dir,
+    )
+    timer = Timer()
+    timer.init_timing("total")
+
+    # Get starting and ending idxs for each file, and collect them into tuples
+    files_range_tuples = _get_files_range_tuples(input_files)  # correct
+    num_files = len(input_files)
+
+    # Initial round of clustering
+    round_idx = 1
+    timer.init_timing(f"round-{round_idx}")
+    console.print(f"(Initial) Round {round_idx}: Cluster initial batch of fingerprints")
+
+    initial_fn = _InitialRound(
+        n_features=n_features,
+        refinement_before_midsection=refinement_before_midsection,
+        max_fps=max_fps,
+        merge_criterion=initial_merge_criterion,
+        input_is_packed=input_is_packed,
+        threshold=threshold,
+        refine_merge_criterion=midsection_merge_criterion,
+        refine_threshold_change=midsection_threshold_change,
+        **common_kwargs,
+    )
+    num_ps = min(num_initial_processes, num_files)
+    console.print(f"    - Processing {num_files} inputs with {num_ps} processes")
+    with console.status("[italic]BitBirching...[/italic]", spinner="dots"):
+        if num_ps == 1:
+            for tup in files_range_tuples:
+                initial_fn(tup)
+        else:
+            with mp_context.Pool(
+                processes=num_ps, maxtasksperchild=max_tasks_per_process
+            ) as pool:
+                pool.map(initial_fn, files_range_tuples)
+
+    timer.end_timing(f"round-{round_idx}", console)
+    console.print_peak_mem(out_dir)
+
+    # Mid-section "Tree-Merging" rounds of clustering
+    for _ in range(num_midsection_rounds):
+        round_idx += 1
+        timer.init_timing(f"round-{round_idx}")
+        console.print(f"(Midsection) Round {round_idx}: Re-clustering in chunks")
+
+        file_pairs = _get_prev_round_buf_and_mol_idxs_files(out_dir, round_idx, console)
+        batches = _chunk_file_pairs_in_batches(file_pairs, bin_size, console)
+        merging_fn = _TreeMergingRound(
+            round_idx=round_idx,
+            all_fp_paths=input_files,
+            split_largest_cluster=split_largest_after_each_midsection_round,
+            criterion=midsection_merge_criterion,
+            threshold=threshold + midsection_threshold_change,
+            **common_kwargs,
+        )
+        num_ps = min(num_midsection_processes, len(batches))
+        console.print(f"    - Processing {len(batches)} inputs with {num_ps} processes")
+        with console.status("[italic]BitBirching...[/italic]", spinner="dots"):
+            if num_ps == 1:
+                for batch_info in batches:
+                    merging_fn(batch_info)
+            else:
+                with mp_context.Pool(
+                    processes=num_ps, maxtasksperchild=max_tasks_per_process
+                ) as pool:
+                    pool.map(merging_fn, batches)
+
+        timer.end_timing(f"round-{round_idx}", console)
+        console.print_peak_mem(out_dir)
+
+    # Final "Tree-Merging" round of clustering
+    round_idx += 1
+    timer.init_timing(f"round-{round_idx}")
+    console.print(f"(Final) Round {round_idx}: Final round of clustering")
+    file_pairs = _get_prev_round_buf_and_mol_idxs_files(out_dir, round_idx, console)
+
+    final_fn = _FinalTreeMergingRound(
+        save_tree=save_tree,
+        save_centroids=save_centroids,
+        criterion=final_merge_criterion,
+        threshold=threshold + midsection_threshold_change,
+        **common_kwargs,
+    )
+    with console.status("[italic]BitBirching...[/italic]", spinner="dots"):
+        final_fn(("", file_pairs))
+
+    timer.end_timing(f"round-{round_idx}", console)
+    console.print_peak_mem(out_dir)
+    # Remove intermediate files
+    if cleanup:
+        for f in out_dir.glob("round-*.npy"):
+            f.unlink()
+        for f in out_dir.glob("round-*.pkl"):
+            f.unlink()
+    console.print()
+    timer.end_timing("total", console, indent=False)
+    return timer