diff --git a/benchmarks/bitbirch_clustering_bench.py b/benchmarks/bitbirch_clustering_bench.py new file mode 100644 index 00000000..a1f1d043 --- /dev/null +++ b/benchmarks/bitbirch_clustering_bench.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark nvMolKit BitBIRCH on real molecular fingerprints. + +The optional ``bblean`` backend receives the exact same packed fingerprints as +nvMolKit. Fingerprint generation and host/device conversion are intentionally +outside the measured clustering region. +""" + +import argparse +import importlib +import importlib.metadata +import math +import sys +import time +from types import ModuleType + +import numpy as np +import nvtx +import torch +from bench_utils import ( + load_smiles, + print_csv_rows, + throughput_per_s, + time_it, + write_csv_rows, +) + +from nvmolkit.clustering import bitbirch +from nvmolkit.fingerprints import MorganFingerprintGenerator + +DEFAULT_NUM_MOLS = [1_000, 10_000, 50_000, 100_000, 250_000, 500_000, 1_000_000] + + +def _load_bblean(enabled: bool) -> tuple[ModuleType | None, str, str]: + if not enabled: + return None, "disabled", "" + try: + module = importlib.import_module("bblean") + except ImportError as error: + return None, "not_installed", str(error) + try: + version = importlib.metadata.version("bblean") + except importlib.metadata.PackageNotFoundError: + version = "unknown" + return module, "available", version + + +@nvtx.annotate("bench_nvmolkit_bitbirch", color="red") +def _bench_nvmolkit( + fingerprints: torch.Tensor, + *, + threshold: float, + branching_factor: int, + batch_size: int, + runs: int, + warmups: int, +) -> tuple[object, dict[str, float | int]]: + holder: dict[str, torch.Tensor] = {} + + @nvtx.annotate("bitbirch_nvmolkit_run", color="orange") + def run() -> None: + holder["labels"] = bitbirch( + fingerprints, + threshold, + branching_factor=branching_factor, + batch_size=batch_size, + ).torch() + + timing = time_it(run, runs=runs, warmups=warmups, gpu_sync=True) + labels = holder["labels"] + return timing, _cluster_distribution(torch.bincount(labels).cpu().numpy(), labels.numel()) + + +@nvtx.annotate("bench_bblean_bitbirch", color="green") +def _bench_bblean( + bblean: ModuleType, + fingerprints: np.ndarray, + *, + threshold: float, + branching_factor: int, + runs: int, + warmups: int, + max_seconds: float, +) -> tuple[object, dict[str, float | int]]: + holder: dict[str, object] = {} + completed = [0] + + @nvtx.annotate("bitbirch_bblean_run", color="yellow") + def run(_deadline=None) -> None: + tree = bblean.BitBirch( + threshold=threshold, + branching_factor=branching_factor, + merge_criterion="diameter", + ) + tree.fit(fingerprints) + holder["tree"] = tree + completed[0] = 1 + + timing = time_it( + run, + runs=runs, + warmups=warmups, + max_seconds=max_seconds, + progress_getter=lambda: completed[0], + progress_target=1, + ) + tree = holder["tree"] + sizes = np.asarray([len(cluster) for cluster in tree.get_cluster_mol_ids()], dtype=np.int64) + return timing, _cluster_distribution(sizes, len(fingerprints)) + + +def _cluster_distribution(sizes: np.ndarray, num_fingerprints: int) -> dict[str, float | int]: + if not sizes.size: + return { + "num_clusters": 0, + "mean_cluster_size": math.nan, + "cluster_size_p50": math.nan, + "cluster_size_p90": math.nan, + "cluster_size_p99": math.nan, + "max_cluster_size": 0, + "singleton_cluster_fraction": math.nan, + "items_in_singletons_fraction": math.nan, + "largest_cluster_fraction": math.nan, + } + quantiles = np.quantile(sizes, [0.5, 0.9, 0.99]) + singleton_count = int(np.count_nonzero(sizes == 1)) + return { + "num_clusters": int(sizes.size), + "mean_cluster_size": float(sizes.mean()), + "cluster_size_p50": float(quantiles[0]), + "cluster_size_p90": float(quantiles[1]), + "cluster_size_p99": float(quantiles[2]), + "max_cluster_size": int(sizes.max()), + "singleton_cluster_fraction": singleton_count / int(sizes.size), + "items_in_singletons_fraction": singleton_count / num_fingerprints, + "largest_cluster_fraction": int(sizes.max()) / num_fingerprints, + } + + +def _timing_fields(timing) -> dict[str, float | int]: + return { + "time_ms": timing.mean_ms, + "std_ms": timing.std_ms, + "runs_completed": len(timing.times_ms), + } + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--smiles", "-s", required=True, help="Path to a real SMILES/CXSMILES input file") + parser.add_argument( + "--num_mols", + "--num-mols", + "-n", + type=int, + nargs="+", + default=DEFAULT_NUM_MOLS, + help="Molecule counts to benchmark", + ) + parser.add_argument("--threshold", type=float, default=0.55) + parser.add_argument( + "--thresholds", + type=float, + nargs="+", + default=None, + help="Sweep multiple thresholds; overrides --threshold", + ) + parser.add_argument("--branching_factor", "--branching-factor", type=int, default=254) + parser.add_argument( + "--batch_size", + "--batch-size", + type=int, + default=1024, + help="nvMolKit insertion batch size (default: 1024)", + ) + parser.add_argument("--radius", type=int, default=2) + parser.add_argument("--fp_size", "--fp-size", type=int, default=1024) + parser.add_argument("--runs", "-r", type=int, default=3) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--no-bblean", + "--no_bblean", + dest="no_bblean", + action="store_true", + help="Skip the bblean benchmark", + ) + parser.add_argument( + "--no-nvmolkit", + "--no_nvmolkit", + dest="no_nvmolkit", + action="store_true", + help="Skip the nvMolKit benchmark", + ) + parser.add_argument( + "--bblean-max-seconds", + "--bblean_max_seconds", + dest="bblean_max_seconds", + type=float, + default=0.0, + help=( + "Stop starting bblean timing repetitions after this many wall-clock seconds per configuration; " + "0 disables the limit" + ), + ) + parser.add_argument("--output", "-o", default=None, help="Optional path to write CSV results") + return parser + + +def main() -> None: + parser = _build_parser() + args = parser.parse_args() + thresholds = args.thresholds if args.thresholds is not None else [args.threshold] + + if not args.num_mols or any(num_mols <= 0 for num_mols in args.num_mols): + parser.error("--num_mols must contain positive integers") + if args.runs <= 0 or args.warmups < 0: + parser.error("--runs must be positive and --warmups must be nonnegative") + if args.batch_size <= 0: + parser.error("--batch-size must be positive") + if args.bblean_max_seconds < 0: + parser.error("--bblean-max-seconds must be nonnegative") + if any(not math.isfinite(threshold) or not 0 <= threshold <= 1 for threshold in thresholds): + parser.error("thresholds must be finite and in [0, 1]") + if args.no_nvmolkit and args.no_bblean: + parser.error("no benchmark backend selected") + + molecule_counts = sorted(set(args.num_mols)) + max_size = molecule_counts[-1] + bblean, bblean_status, bblean_version = _load_bblean(not args.no_bblean) + if not args.no_bblean and bblean is None: + print(f"bblean unavailable ({bblean_status}): {bblean_version}", file=sys.stderr) + if args.no_nvmolkit and bblean is None: + parser.error("bblean is unavailable and nvMolKit is disabled") + + print( + f"Scanning {args.smiles} and sampling molecules for a maximum benchmark size of {max_size:,}", + flush=True, + ) + molecule_load_start = time.perf_counter() + with nvtx.annotate("bitbirch_molecule_loading", color="cyan"): + molecules = load_smiles(args.smiles, max_count=max_size + 100, sanitize=True, seed=args.seed) + molecule_load_seconds = time.perf_counter() - molecule_load_start + print(f"Molecule loading finished in {molecule_load_seconds:.3f} s", flush=True) + if len(molecules) < max_size: + parser.error(f"requested {max_size} molecules, but only {len(molecules)} valid molecules were loaded") + + print(f"Generating {max_size:,} Morgan fingerprints with all available CPU threads", flush=True) + fingerprint_start = time.perf_counter() + with nvtx.annotate("bitbirch_fingerprint_generation", color="blue"): + fingerprints = ( + MorganFingerprintGenerator(args.radius, args.fp_size) + .GetFingerprints(molecules[:max_size], num_threads=0) + .torch() + ) + torch.cuda.synchronize() + fingerprint_seconds = time.perf_counter() - fingerprint_start + print(f"Fingerprint generation finished in {fingerprint_seconds:.3f} s", flush=True) + + results = [] + try: + for size in molecule_counts: + device_fps = fingerprints[:size].contiguous() + host_fps = device_fps.cpu().numpy().view(np.uint8).reshape(size, -1).copy() if bblean is not None else None + for threshold in thresholds: + print(f"\nConfiguration: {size:,} fingerprints, threshold={threshold}", flush=True) + common_fields = { + "num_mols": size, + "threshold": threshold, + "branching_factor": args.branching_factor, + "radius": args.radius, + "fp_size": args.fp_size, + "runs": args.runs, + "warmups": args.warmups, + "seed": args.seed, + "fingerprint_total_seconds": fingerprint_seconds, + "input_file": args.smiles, + "input_type": "smiles", + } + configuration_rows = [] + + if not args.no_nvmolkit: + print( + f" Starting nvMolKit GPU: {args.warmups} warmup(s), {args.runs} timed run(s), " + f"batch_size={args.batch_size}", + flush=True, + ) + torch.cuda.cudart().cudaProfilerStart() + try: + timing, distribution = _bench_nvmolkit( + device_fps, + threshold=threshold, + branching_factor=args.branching_factor, + batch_size=args.batch_size, + runs=args.runs, + warmups=args.warmups, + ) + finally: + torch.cuda.cudart().cudaProfilerStop() + print( + f" Finished nvMolKit GPU: {timing.mean_ms:.3f} +/- {timing.std_ms:.3f} ms " + f"({distribution['num_clusters']:,} clusters)", + flush=True, + ) + configuration_rows.append( + { + "method": "nvmolkit", + **common_fields, + "batch_size": args.batch_size, + "device": torch.cuda.get_device_name(device_fps.device), + **_timing_fields(timing), + **distribution, + } + ) + + if bblean is not None: + budget = f", {args.bblean_max_seconds:g} s timing budget" if args.bblean_max_seconds > 0 else "" + print( + f" Starting bblean CPU: {args.warmups} warmup(s), up to {args.runs} timed run(s){budget}", + flush=True, + ) + timing, distribution = _bench_bblean( + bblean, + host_fps, + threshold=threshold, + branching_factor=args.branching_factor, + runs=args.runs, + warmups=args.warmups, + max_seconds=args.bblean_max_seconds, + ) + print( + f" Finished bblean CPU: {timing.mean_ms:.3f} +/- {timing.std_ms:.3f} ms " + f"over {len(timing.times_ms)} run(s) ({distribution['num_clusters']:,} clusters)", + flush=True, + ) + configuration_rows.append( + { + "method": "bblean", + **common_fields, + "batch_size": "N/A", + "device": "CPU", + "version": bblean_version, + **_timing_fields(timing), + **distribution, + } + ) + + bblean_throughput = next( + ( + throughput_per_s(size, result_row["time_ms"]) + for result_row in configuration_rows + if result_row["method"] == "bblean" + ), + math.nan, + ) + for result_row in configuration_rows: + throughput = throughput_per_s(size, result_row["time_ms"]) + result_row["fingerprints_per_second"] = throughput + result_row["vs_bblean_throughput_ratio"] = ( + throughput / bblean_throughput + if result_row["method"] == "nvmolkit" and math.isfinite(bblean_throughput) + else "N/A" + ) + + results.extend(configuration_rows) + write_csv_rows(results, args.output) + except Exception: + write_csv_rows(results, args.output) + raise + + print("\nCSV Results:") + print_csv_rows(results) + if args.output: + print(f"\nWrote results to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/tests/test_bitbirch_clustering_bench.py b/benchmarks/tests/test_bitbirch_clustering_bench.py new file mode 100644 index 00000000..5ec90d41 --- /dev/null +++ b/benchmarks/tests/test_bitbirch_clustering_bench.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import bitbirch_clustering_bench as bench +import numpy as np +import pytest + + +@pytest.mark.parametrize("flag", ["--no-bblean", "--no_bblean"]) +def test_backend_selection_accepts_bblean_disable_spellings(flag): + args = bench._build_parser().parse_args(["--smiles", "input.smi", flag]) + + assert args.no_bblean + assert not args.no_nvmolkit + + +@pytest.mark.parametrize("flag", ["--no-nvmolkit", "--no_nvmolkit"]) +def test_backend_selection_accepts_nvmolkit_disable_spellings(flag): + args = bench._build_parser().parse_args(["--smiles", "input.smi", flag]) + + assert args.no_nvmolkit + assert not args.no_bblean + + +def test_timing_fields_use_standard_mean_schema(): + timing = bench.time_it(lambda: None, runs=2, warmups=0) + + fields = bench._timing_fields(timing) + + assert set(fields) == {"time_ms", "std_ms", "runs_completed"} + assert fields["time_ms"] == timing.mean_ms + assert fields["std_ms"] == timing.std_ms + assert fields["runs_completed"] == 2 + + +def test_bblean_deadline_accepts_standard_spellings(): + parser = bench._build_parser() + + hyphenated = parser.parse_args(["--smiles", "input.smi", "--bblean-max-seconds", "12.5"]) + underscored = parser.parse_args(["--smiles", "input.smi", "--bblean_max_seconds", "7.5"]) + + assert hyphenated.bblean_max_seconds == 12.5 + assert underscored.bblean_max_seconds == 7.5 + + +def test_parser_uses_standard_input_and_workload_options(): + args = bench._build_parser().parse_args( + [ + "--smiles", + "input.smi", + "--num_mols", + "10", + "20", + "--branching_factor", + "64", + "--batch_size", + "512", + "--fp_size", + "512", + "-r", + "2", + ] + ) + + assert args.smiles == "input.smi" + assert args.num_mols == [10, 20] + assert args.branching_factor == 64 + assert args.batch_size == 512 + assert args.fp_size == 512 + assert args.runs == 2 + assert args.output is None + + +def test_bblean_is_not_imported_when_disabled(monkeypatch): + def unexpected_import(_name): + raise AssertionError("bblean import should not be attempted") + + monkeypatch.setattr(bench.importlib, "import_module", unexpected_import) + + module, status, detail = bench._load_bblean(False) + + assert module is None + assert status == "disabled" + assert detail == "" + + +def test_missing_bblean_is_reported_as_optional(monkeypatch): + def missing_import(_name): + raise ImportError("no bblean") + + monkeypatch.setattr(bench.importlib, "import_module", missing_import) + + module, status, detail = bench._load_bblean(True) + + assert module is None + assert status == "not_installed" + assert detail == "no bblean" + + +def test_cluster_distribution_reports_sizes_and_singletons(): + distribution = bench._cluster_distribution(np.array([3, 1, 1]), 5) + + assert distribution["num_clusters"] == 3 + assert distribution["mean_cluster_size"] == pytest.approx(5 / 3) + assert distribution["cluster_size_p50"] == 1 + assert distribution["cluster_size_p90"] == pytest.approx(2.6) + assert distribution["cluster_size_p99"] == pytest.approx(2.96) + assert distribution["max_cluster_size"] == 3 + assert distribution["singleton_cluster_fraction"] == pytest.approx(2 / 3) + assert distribution["items_in_singletons_fraction"] == pytest.approx(2 / 5) + assert distribution["largest_cluster_fraction"] == pytest.approx(3 / 5) diff --git a/docs/api/nvmolkit.rst b/docs/api/nvmolkit.rst index 107a7c8c..d122a1aa 100644 --- a/docs/api/nvmolkit.rst +++ b/docs/api/nvmolkit.rst @@ -92,6 +92,7 @@ Clustering clustering.fused_butina clustering.dise clustering.fused_dise + clustering.bitbirch .. autosummary:: :toctree: generated/ diff --git a/docs/index.rst b/docs/index.rst index 34aa525c..5fd4f6b0 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -184,6 +184,8 @@ nvMolKit currently supports the following features: * Supports generic dense similarity matrix inputs from Tanimoto or other methods * Supports fused in-loop Tanimoto/Cosine similarity for large datasets +* **BitBIRCH clustering**: GPU-accelerated clustering of binary fingerprints, with optional cluster centroids + * **Substructure Search**: GPU-accelerated substructure matching against batches of molecules * Supports SMILES and recursive SMARTS-based query molecules via RDKit * ``SubstructLibrary`` keeps target molecules on one or more GPUs for repeated queries, as a GPU counterpart to RDKit's ``SubstructLibrary`` diff --git a/nvmolkit/CMakeLists.txt b/nvmolkit/CMakeLists.txt index 79c437c6..0ac60942 100644 --- a/nvmolkit/CMakeLists.txt +++ b/nvmolkit/CMakeLists.txt @@ -102,6 +102,7 @@ target_link_libraries( PUBLIC ${Boost_LIBRARIES} ${PYTHON_LIBRARIES} device_vector + bitbirch butina aap diversity_pickers diff --git a/nvmolkit/clustering.cpp b/nvmolkit/clustering.cpp index eae17255..81ac13de 100644 --- a/nvmolkit/clustering.cpp +++ b/nvmolkit/clustering.cpp @@ -27,6 +27,7 @@ #include "nvmolkit/array_helpers.h" #include "nvmolkit/boost_python_utils.h" #include "src/aap.h" +#include "src/bitbirch.h" #include "src/butina.h" #include "src/diversity_pickers.h" #include "src/utils/device.h" @@ -164,9 +165,41 @@ nvMolKit::FingerprintSimilarityMetric parseFingerprintMetric(const std::string& throw std::invalid_argument("metric must be one of ['tanimoto', 'cosine']"); } +//! Returns (labels, centroids); centroids is None unless requested. +boost::python::object bitBirch(const boost::python::dict& fingerprints, + const double threshold, + const int branchingFactor, + const int batchSize, + const bool returnCentroids, + const std::uintptr_t streamPtr) { + const auto stream = requireStream(streamPtr); + const auto input = parseFingerprints(fingerprints); + nvMolKit::BitBirchOptions options; + options.branchingFactor = branchingFactor; + options.batchSize = batchSize; + options.returnCentroids = returnCentroids; + auto result = nvMolKit::bitBirchGpu(input.fingerprints, input.numItems, input.numWords, threshold, options, stream); + auto labels = toOwnedPyArray(nvMolKit::makePyArray(result.clusterIds, boost::python::make_tuple(input.numItems))); + boost::python::object centroids; + if (returnCentroids) { + centroids = toOwnedPyArray( + nvMolKit::makePyArray(result.centroids, boost::python::make_tuple(result.numClusters, result.numWords))); + } + return boost::python::make_tuple(labels, centroids); +} + } // namespace BOOST_PYTHON_MODULE(_clustering) { + boost::python::def("bitbirch", + bitBirch, + (boost::python::arg("fingerprints"), + boost::python::arg("threshold"), + boost::python::arg("branching_factor") = 1024, + boost::python::arg("batch_size") = 1024, + boost::python::arg("return_centroids") = false, + boost::python::arg("stream") = 0)); + boost::python::register_exception_translator(translateAapInvalidMolecules); boost::python::def( diff --git a/nvmolkit/clustering.py b/nvmolkit/clustering.py index 71b035e0..d5a7e6a6 100644 --- a/nvmolkit/clustering.py +++ b/nvmolkit/clustering.py @@ -13,12 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GPU-accelerated clustering. +"""GPU-accelerated Butina, directed sphere exclusion (DISE), and BitBIRCH clustering. -Each algorithm has a matrix form that takes a precomputed distance matrix and a -``fused_`` form that computes distances from fingerprints or molecules as needed. +Butina and DISE accept precomputed distance matrices or compute distances as +needed through their ``fused_`` functions. BitBIRCH clusters packed binary +fingerprints directly. """ +import math from dataclasses import dataclass from enum import Enum from typing import Literal, overload @@ -33,6 +35,7 @@ from nvmolkit.types import ArrayInput, AsyncGpuResult, _as_cuda_tensor, _resolve_cuda_stream _VALID_NEIGHBORLIST_SIZES = (8, 16, 24, 32, 64, 128) +_BITBIRCH_MAX_BRANCHING_FACTOR = 65534 # Matches kMaxBranchingFactor in src/bitbirch.cu. _RDKitClusters = tuple[tuple[int, ...], ...] @@ -140,6 +143,9 @@ def dise( ) -> ClusterDeviceResult | _RDKitClusters: """Cluster a distance matrix by directed sphere exclusion (DISE). + `Gobbi et al., J. Cheminform. 7, 11 (2015) + `_. + Centroids are the leaders selected by :func:`~nvmolkit.pickers.leader`. With ``assignment="first"``, each item joins the first centroid that excluded it; with ``"nearest"``, each non-centroid joins its nearest centroid, with @@ -191,6 +197,9 @@ def fused_dise( ) -> ClusterDeviceResult | _RDKitClusters: """Cluster by directed sphere exclusion (DISE), computing distances as needed. + `Gobbi et al., J. Cheminform. 7, 11 (2015) + `_. + Equivalent to :func:`dise` on the matrix of ``1 - similarity`` values, with memory that scales as ``O(N)``. Distances are computed in float32, so results can differ slightly from RDKit's double-precision results. @@ -207,10 +216,6 @@ def fused_dise( Returns: A :class:`ClusterDeviceResult` for ``OutputMode.DEVICE``, or centroid-first tuples of input indices for ``OutputMode.RDKIT``. - - Note: - For the method, see `Gobbi et al. (2015) - `_. """ _validate_output(output) _validate_assignment(assignment) @@ -272,6 +277,8 @@ def butina( ) -> _RDKitClusters | ClusterDeviceResult: """Perform Butina clustering on a distance matrix. + `Butina, JCICS 39, 747-750 (1999) `_. + The Butina algorithm is a deterministic clustering method that groups items based on distance thresholds. It iteratively: 1. Finds the item with the most neighbors within the cutoff distance @@ -356,6 +363,8 @@ def fused_butina( ) -> _RDKitClusters | ClusterDeviceResult: """Perform Butina clustering on fingerprints, computing distances as needed. + `Butina, JCICS 39, 747-750 (1999) `_. + Equivalent to :func:`butina` on the matrix of ``1 - similarity`` values, without forming the ``N x N`` matrix. @@ -388,3 +397,56 @@ def fused_butina( x.__cuda_array_interface__, cutoff, True, _packed_metric_name(resolved), active_stream.cuda_stream ) return _resolve_butina_output(result, output) + + +def bitbirch( + x: ArrayInput, + threshold: float, + *, + branching_factor: int = 1024, + batch_size: int = 1024, + return_centroids: bool = False, + stream: torch.cuda.Stream | None = None, +) -> AsyncGpuResult | tuple[AsyncGpuResult, AsyncGpuResult]: + """Cluster packed binary fingerprints on the GPU using BitBIRCH. + + `Pickering et al., BitBIRCH-Lean, bioRxiv (2025) + `_. + + Args: + x: Shape ``(N, W)`` packed int32 or uint32 fingerprints, ``W <= 128`` + (at most 4096 bits). Host arrays are copied to the GPU. + threshold: Minimum combined-cluster iSIM Jaccard--Tanimoto similarity. + branching_factor: Maximum entries per tree node. Must be in [3, 65534]. + batch_size: Number of fingerprints processed together. Must be positive. + return_centroids: Also return packed majority centroids with shape + ``(num_clusters, W)``. + stream: CUDA stream to use. If None, uses the current stream. + + Returns: + One cluster ID per input fingerprint, numbered by each cluster's earliest + member, and the centroids in cluster-ID order if requested. + + Notes: + A row represents exactly ``32 * W`` bits. Results are deterministic for + fixed input and options; changing ``batch_size`` can change the clustering. + """ + if not math.isfinite(threshold) or not 0 <= threshold <= 1: + raise ValueError("threshold must be finite and in [0, 1]") + if not 3 <= branching_factor <= _BITBIRCH_MAX_BRANCHING_FACTOR: + raise ValueError(f"branching_factor must be in [3, {_BITBIRCH_MAX_BRANCHING_FACTOR}], got {branching_factor}") + if batch_size < 1: + raise ValueError(f"batch_size must be positive, got {batch_size}") + (x,), active_stream = _prepare_packed_fingerprints(("x", x), stream=stream) + if x.shape[1] > 128: + raise ValueError(f"bitbirch supports at most 128 fingerprint words (4096 bits), got {x.shape[1]}") + with torch.cuda.stream(active_stream): + labels, centroids = _clustering.bitbirch( + x.__cuda_array_interface__, + threshold, + branching_factor, + batch_size, + return_centroids, + active_stream.cuda_stream, + ) + return (AsyncGpuResult(labels), AsyncGpuResult(centroids)) if return_centroids else AsyncGpuResult(labels) diff --git a/nvmolkit/tests/_bitbirch_reference.py b/nvmolkit/tests/_bitbirch_reference.py new file mode 100644 index 00000000..1f5387af --- /dev/null +++ b/nvmolkit/tests/_bitbirch_reference.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Independent semantic model of nvMolKit's concurrent BitBIRCH schedule. + +This CPU semantic model is intentionally not a scalable storage implementation. +It uses unpacked input and uint64 sums to make membership auditing straightforward. +Logical batches and structural subrounds define the algorithm. No bblean +implementation is used here. +""" + +from collections import Counter, defaultdict +from dataclasses import dataclass, field + +import numpy as np + +POPCOUNT = np.asarray([value.bit_count() for value in range(256)], dtype=np.uint8) +ORDERED_WARMUP_SIZE = 16384 + + +def isim(sums: np.ndarray, count: int) -> float: + values = sums.astype(np.float64) + common = np.sum(values * (values - 1) * 0.5) + mismatch = np.sum(values * (count - values)) + return float(common / (common + mismatch)) if common + mismatch else 1.0 + + +def centroid(sums: np.ndarray, count: int) -> np.ndarray: + return np.packbits(sums >= (count + 1) // 2) + + +def similarities(queries: np.ndarray, candidates: np.ndarray) -> np.ndarray: + intersection = POPCOUNT[queries[:, None, :] & candidates[None, :, :]].sum(axis=2) + union = POPCOUNT[queries[:, None, :] | candidates[None, :, :]].sum(axis=2) + return np.divide(intersection, union, out=np.ones_like(intersection, dtype=float), where=union != 0) + + +@dataclass(eq=False) +class Entry: + uid: int + count: int + sums: np.ndarray + packed: np.ndarray + members: list[int] = field(default_factory=list) + child: "Node | None" = None + + +@dataclass(eq=False) +class Node: + uid: int + leaf: bool + entries: list[Entry] = field(default_factory=list) + parent: "Node | None" = None + cached_centroids: np.ndarray | None = None + + def matrix(self) -> np.ndarray: + if self.cached_centroids is None: + self.cached_centroids = np.stack([entry.packed for entry in self.entries]) + return self.cached_centroids + + +class BatchedBitBirch: + def __init__( + self, + threshold=0.25, + branching_factor=254, + batch_size=256, + ): + if not np.isfinite(threshold) or not 0 <= threshold <= 1: + raise ValueError("threshold must be finite and in [0, 1]") + if branching_factor < 3 or batch_size < 1: + raise ValueError("branching factor >= 3 and positive batch size required") + self.threshold = threshold + self.branching_factor = branching_factor + self.batch_size = batch_size + self.root = Node(0, True) + self.next_node = 1 + self.stats = Counter() + self.bits = np.empty((0, 0), dtype=np.uint8) + self.packed = np.empty((0, 0), dtype=np.uint8) + + def _node(self, leaf: bool, entries: list[Entry], parent=None) -> Node: + result = Node(self.next_node, leaf, entries, parent) + self.next_node += 1 + for entry in entries: + if entry.child is not None: + entry.child.parent = result + return result + + def _summary(self, node: Node) -> Entry: + sums = np.sum([entry.sums for entry in node.entries], axis=0, dtype=np.uint64) + count = sum(entry.count for entry in node.entries) + return Entry(-node.uid - 1, count, sums, centroid(sums, count), child=node) + + def _refresh(self, node: Node) -> None: + if not node.leaf: + for index, entry in enumerate(node.entries): + self._refresh(entry.child) + node.entries[index] = self._summary(entry.child) + node.cached_centroids = None + + def _route(self, molecule_ids: np.ndarray): + pending = [(self.root, molecule_ids)] + groups = defaultdict(list) + leaves = {} + entries = {} + while pending: + node, ids = pending.pop() + leaves[node.uid] = node + if not node.entries: + groups[(node.uid, None)].extend(ids.tolist()) + continue + choices = np.argmax(similarities(self.packed[ids], node.matrix()), axis=1) + for choice in np.unique(choices): + selected = ids[choices == choice] + entry = node.entries[int(choice)] + if node.leaf: + groups[(node.uid, entry.uid)].extend(selected.tolist()) + entries[entry.uid] = entry + else: + pending.append((entry.child, selected)) + self.stats["route_attempts"] += len(molecule_ids) + self.stats["max_cluster_queue"] = max(self.stats["max_cluster_queue"], *map(len, groups.values())) + return groups, leaves, entries + + def _commit(self, entry: Entry, ids: list[int], sums: np.ndarray) -> None: + entry.sums = sums + entry.count += len(ids) + entry.packed = centroid(sums, entry.count) + entry.members.extend(ids) + + def _try_group(self, entry: Entry, ids: list[int]) -> list[int]: + residual = [] + eligible = [] + for molecule in ids: + if isim(entry.sums + self.bits[molecule], entry.count + 1) >= self.threshold: + eligible.append(molecule) + else: + residual.append(molecule) + if not eligible: + return residual + combined = entry.sums + self.bits[eligible].sum(axis=0, dtype=np.uint64) + if isim(combined, entry.count + len(eligible)) >= self.threshold: + self._commit(entry, eligible, combined) + return residual + for molecule in eligible: + combined = entry.sums + self.bits[molecule] + if isim(combined, entry.count + 1) >= self.threshold: + self._commit(entry, [molecule], combined) + else: + residual.append(molecule) + return residual + + def _insert_residual(self, node: Node, molecule: int) -> None: + if node.entries: + selected = int(np.argmax(similarities(self.packed[molecule : molecule + 1], node.matrix())[0])) + entry = node.entries[selected] + combined = entry.sums + self.bits[molecule] + self.stats["residual_checks"] += 1 + if isim(combined, entry.count + 1) >= self.threshold: + self._commit(entry, [molecule], combined) + node.cached_centroids = None + self.stats["residual_committed"] += 1 + return + sums = self.bits[molecule].astype(np.uint64) + node.entries.append(Entry(molecule, 1, sums, self.packed[molecule].copy(), [molecule])) + node.cached_centroids = None + self.stats["created_clusters"] += 1 + + def _split(self, node: Node) -> None: + count = len(node.entries) + assert count == self.branching_factor + 1 + centroids = node.matrix() + distances = similarities(centroids, centroids) + distances[np.tril_indices(count)] = np.inf + lhs, rhs = np.unravel_index(np.argmin(distances), distances.shape) + seeds = similarities(centroids, centroids[[lhs, rhs]]) + left = [node.entries[lhs]] + right = [node.entries[rhs]] + capacity = (count + 1) // 2 + for index, entry in enumerate(node.entries): + if index in (lhs, rhs): + continue + take_left = seeds[index, 0] > seeds[index, 1] or ( + seeds[index, 0] == seeds[index, 1] and len(left) <= len(right) + ) + if len(left) >= capacity: + take_left = False + elif len(right) >= capacity: + take_left = True + if take_left: + left.append(entry) + else: + right.append(entry) + node.entries = left + node.cached_centroids = None + for entry in left: + if entry.child is not None: + entry.child.parent = node + sibling = self._node(node.leaf, right, node.parent) + if node.parent is None: + self.root = self._node(False, [self._summary(node), self._summary(sibling)]) + else: + parent = node.parent + for index, entry in enumerate(parent.entries): + if entry.child is node: + parent.entries[index] = self._summary(node) + break + parent.entries.append(self._summary(sibling)) + parent.cached_centroids = None + if len(parent.entries) > self.branching_factor: + self._split(parent) + self.stats["splits"] += 1 + + def _epoch(self, ids: np.ndarray) -> None: + ordered_leaf = ids[0] < ORDERED_WARMUP_SIZE + pending = ids + while len(pending): + self.stats["subrounds"] += 1 + groups, nodes, entries = self._route(pending) + residual = defaultdict(list) + # All routing has finished. Leaf owners touch disjoint data. + keys = sorted(groups) + for node_id, entry_id in keys: + ids = groups[(node_id, entry_id)] + if entry_id is None or ordered_leaf: + residual[node_id].extend(ids) + else: + residual[node_id].extend(self._try_group(entries[entry_id], ids)) + nodes[node_id].cached_centroids = None + overflow = [] + parked = [] + for node_id in sorted(residual): + node = nodes[node_id] + ids = sorted(residual[node_id]) + self.stats["max_residual_leaf_queue"] = max(self.stats["max_residual_leaf_queue"], len(ids)) + for offset, molecule in enumerate(ids): + self._insert_residual(node, molecule) + if len(node.entries) > self.branching_factor: + overflow.append(node) + parked.extend(ids[offset + 1 :]) + break + # All payload updates finish before topology changes. Refresh first so + # cascading splits see current summaries of every sibling. + self._refresh(self.root) + for node in sorted(overflow, key=node_uid): + self._split(node) + # A previous split may have created/moved an ancestor of the next. + self._refresh(self.root) + pending = np.asarray(sorted(parked), dtype=np.int64) + self.stats["parked_for_split"] += len(pending) + + def fit(self, packed: np.ndarray): + data = np.asarray(packed) + if data.ndim != 2 or data.dtype != np.uint8 or data.shape[1] == 0: + raise ValueError("expected packed uint8 matrix with nonzero width") + if self.root.entries: + raise ValueError("reference fit requires a fresh tree") + self.packed = data + self.bits = np.unpackbits(data, axis=1) + begin = 0 + while begin < len(data): + boundary = min(ORDERED_WARMUP_SIZE, len(data)) if begin < ORDERED_WARMUP_SIZE else len(data) + end = min(begin + self.batch_size, boundary) + self.stats["epochs"] += 1 + self._epoch(np.arange(begin, end, dtype=np.int64)) + begin = end + return self + + def clusters(self) -> list[Entry]: + pending = [self.root] + result = [] + while pending: + node = pending.pop() + if node.leaf: + result.extend(node.entries) + else: + pending.extend(entry.child for entry in node.entries) + return sorted(result, key=first_member) + + def labels(self) -> np.ndarray: + labels = np.full(len(self.packed), -1, dtype=np.int64) + for label, entry in enumerate(self.clusters()): + labels[entry.members] = label + return labels + + def audit(self) -> dict: + """Recompute from actual members, independently of incremental updates.""" + members = [] + scores = [] + pending = [self.root] + while pending: + node = pending.pop() + assert len(node.entries) <= self.branching_factor + if node is not self.root: + assert len(node.entries) >= 2 + for entry in node.entries: + if node.leaf: + assert entry.child is None + assert entry.count == len(entry.members) + expected = self.bits[entry.members].sum(axis=0, dtype=np.uint64) + members.extend(entry.members) + value = isim(expected, entry.count) + assert value + 1e-12 >= self.threshold + scores.append(value) + else: + assert entry.child.parent is node + assert not entry.members + assert entry.count == sum(child.count for child in entry.child.entries) + expected = np.sum([child.sums for child in entry.child.entries], axis=0, dtype=np.uint64) + pending.append(entry.child) + np.testing.assert_array_equal(entry.sums, expected) + np.testing.assert_array_equal(entry.packed, centroid(expected, entry.count)) + assert sorted(members) == list(range(len(self.packed))) + return {"audited_molecules": len(members), "minimum_isim": min(scores, default=1.0)} + + +def node_uid(node: Node) -> int: + return node.uid + + +def first_member(entry: Entry) -> int: + return min(entry.members) diff --git a/nvmolkit/tests/test_bitbirch.py b/nvmolkit/tests/test_bitbirch.py new file mode 100644 index 00000000..49ef5b14 --- /dev/null +++ b/nvmolkit/tests/test_bitbirch.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Correctness tests for batched BitBIRCH against an independent CPU model.""" + +import numpy as np +import pytest +import torch +from _bitbirch_reference import ORDERED_WARMUP_SIZE, BatchedBitBirch + +from nvmolkit.clustering import bitbirch + + +def reference(packed, threshold, branching_factor, batch_size): + bytes_ = packed.view(np.uint8).reshape(len(packed), packed.shape[1] * 4) + tree = BatchedBitBirch(threshold, branching_factor, batch_size).fit(bytes_) + tree.audit() + centroids = np.stack([entry.packed for entry in tree.clusters()]).view(np.uint32).reshape(-1, packed.shape[1]) + return tree.labels(), centroids + + +def assert_matches_reference(packed, threshold, branching_factor, batch_size): + expected, expected_centroids = reference(packed, threshold, branching_factor, batch_size) + labels, centroids = bitbirch( + packed, threshold, branching_factor=branching_factor, batch_size=batch_size, return_centroids=True + ) + np.testing.assert_array_equal(labels.numpy(), expected) + np.testing.assert_array_equal(centroids.numpy(), expected_centroids) + + +@pytest.mark.parametrize("branching_factor,batch_size,words", [(3, 7, 1), (7, 32, 3), (254, 32, 32)]) +@pytest.mark.parametrize("threshold", [0.0, 0.3, 0.7, 1.0]) +def test_matches_independent_reference(branching_factor, batch_size, words, threshold): + packed = np.random.default_rng(740 + words).integers(0, 2**32, size=(97, words), dtype=np.uint32) + packed[30:40] = packed[:10] + assert_matches_reference(packed, threshold, branching_factor, batch_size) + + +@pytest.mark.parametrize("words", [127, 128]) +def test_widest_fingerprints_match_reference(words): + packed = np.random.default_rng(202 + words).integers(0, 2**32, size=(43, words), dtype=np.uint32) + packed[30:40] = packed[:10] + assert_matches_reference(packed, 0.4, 3, 16) + + +@pytest.mark.parametrize("branching_factor,words", [(63, 2), (64, 2), (257, 1), (256, 64), (300, 32), (600, 64)]) +def test_leaves_spanning_split_tiles_match_reference(branching_factor, words): + # Overflowing leaves of branching_factor + 1 entries cover one or more 64-entry split tiles. + rows = 4 * branching_factor + 77 + packed = np.random.default_rng(1234 + branching_factor).integers(0, 2**32, size=(rows, words), dtype=np.uint32) + packed[rows // 2 : rows // 2 + 60] = packed[:60] + assert_matches_reference(packed, 0.55, branching_factor, 512) + + +def test_default_branching_factor_matches_reference(): + packed = np.random.default_rng(5).integers(0, 2**32, size=(4 * 1025 + 77, 32), dtype=np.uint32) + packed[2000:2100] = packed[:100] + assert_matches_reference(packed, 0.55, 1024, 512) + + +def test_wide_nodes_with_tied_similarities_match_reference(): + # One-word fingerprints give few distinct similarities, so ties decide most choices. + packed = np.random.default_rng(17).integers(0, 2**32, size=(5000, 1), dtype=np.uint32) + assert_matches_reference(packed, 1.0, 2048, 1024) + + +@pytest.mark.parametrize("threshold", [0.6, 1.0]) +def test_cascading_splits_of_multi_tile_inner_nodes_match_reference(threshold): + # Over 70 leaves force inner nodes of 71 entries (two split tiles) to split as well. + packed = np.random.default_rng(41).integers(0, 2**32, size=(6000, 4), dtype=np.uint32) + assert_matches_reference(packed, threshold, 70, 512) + + +def test_grouped_merges_after_warmup_match_reference(): + # Clustered rows past the ordered warm-up, in batches that do not divide it. + rng = np.random.default_rng(5) + base = rng.integers(0, 2**32, size=(400, 2), dtype=np.uint32) + noise = np.bitwise_and.reduce([rng.integers(0, 2**32, size=(20000, 2), dtype=np.uint32) for _ in range(3)]) + packed = base[rng.integers(0, len(base), len(noise))] ^ noise + assert len(packed) > ORDERED_WARMUP_SIZE and ORDERED_WARMUP_SIZE % 1000 != 0 + assert_matches_reference(packed, 0.5, 50, 1000) + + +def test_split_wider_than_one_block_matches_reference(): + packed = np.random.default_rng(82).integers(0, 2**32, size=(529, 1), dtype=np.uint32) + assert_matches_reference(packed, 1.0, 257, 512) + + +@pytest.mark.parametrize("branching_factor", [254, 255]) +def test_full_1024bit_leaves_match_reference(branching_factor): + packed = np.random.default_rng(991).integers(0, 2**32, size=(529, 32), dtype=np.uint32) + packed[510:520] = packed[:10] + assert_matches_reference(packed, 0.55, branching_factor, 512) + + +def test_multiple_owners_and_cascading_splits(): + packed = np.random.default_rng(331).integers(0, 2**32, size=(43, 2), dtype=np.uint32) + packed[29:36] = packed[:7] + assert_matches_reference(packed, 0.4, 3, 16) + + +def test_concurrent_leaf_allocation_is_schedule_independent(): + packed = np.random.default_rng(73).integers(0, 2**32, size=(503, 3), dtype=np.uint32) + expected, _ = reference(packed, 0.55, 3, 128) + for _ in range(5): + labels = bitbirch(packed, 0.55, branching_factor=3, batch_size=128).numpy() + np.testing.assert_array_equal(labels, expected) + + +def test_empty_zeros_and_duplicates(): + labels, centroids = bitbirch(np.zeros((0, 2), dtype=np.uint32), 0.5, return_centroids=True) + assert labels.numpy().shape == (0,) + assert centroids.numpy().shape == (0, 2) + packed = np.zeros((65537, 1), dtype=np.uint32) + labels, centroids = bitbirch(packed, 1.0, return_centroids=True) + np.testing.assert_array_equal(labels.numpy(), np.zeros(len(packed), dtype=np.int32)) + np.testing.assert_array_equal(centroids.numpy(), np.zeros((1, 1), dtype=np.uint32)) + + +def test_explicit_stream_orders_input_clustering_and_consumption(): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + packed = torch.tensor([[1], [2], [4], [8], [1], [2], [4], [8]], dtype=torch.int32, device="cuda") + labels = bitbirch(packed, 1.0, branching_factor=3, batch_size=4, stream=stream) + copied = labels.torch().clone() + stream.synchronize() + expected, _ = reference(packed.cpu().numpy().view(np.uint32), 1.0, 3, 4) + np.testing.assert_array_equal(copied.cpu().numpy(), expected) + + +def test_memory_mapped_input_matches_device_input(tmp_path): + packed = np.random.default_rng(871).integers(0, 2**32, (529, 3), dtype=np.uint32) + packed[250:300] = packed[:50] + path = tmp_path / "fingerprints.npy" + np.save(path, packed) + mapped = np.load(path, mmap_mode="r") + options = {"branching_factor": 3, "batch_size": 32, "return_centroids": True} + expected, expected_centroids = bitbirch(torch.from_numpy(packed.view(np.int32)).cuda(), 0.4, **options) + labels, centroids = bitbirch(mapped, 0.4, **options) + np.testing.assert_array_equal(labels.numpy(), expected.numpy()) + np.testing.assert_array_equal(centroids.numpy(), expected_centroids.numpy()) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"threshold": float("nan")}, + {"threshold": -1}, + {"threshold": 0.5, "branching_factor": 2}, + {"threshold": 0.5, "batch_size": 0}, + ], +) +def test_invalid_options_rejected(kwargs): + with pytest.raises(ValueError): + bitbirch(np.zeros((1, 1), dtype=np.uint32), **kwargs) + + +def test_fingerprints_wider_than_4096_bits_rejected(): + with pytest.raises(ValueError, match="at most 128 fingerprint words"): + bitbirch(np.zeros((4, 129), dtype=np.uint32), 0.5) diff --git a/nvmolkit/tests/test_bitbirch_integration.py b/nvmolkit/tests/test_bitbirch_integration.py new file mode 100644 index 00000000..9df25457 --- /dev/null +++ b/nvmolkit/tests/test_bitbirch_integration.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch +from _bitbirch_reference import BatchedBitBirch +from rdkit.Chem import rdFingerprintGenerator + +from nvmolkit.clustering import bitbirch +from nvmolkit.fingerprints import MorganFingerprintGenerator + + +def _reference(packed, threshold, branching_factor, batch_size): + byte_rows = packed.view(np.uint8).reshape(len(packed), packed.shape[1] * 4) + tree = BatchedBitBirch(threshold, branching_factor, batch_size).fit(byte_rows) + tree.audit() + centroids = np.stack([entry.packed for entry in tree.clusters()]).view(np.uint32) + return tree.labels(), centroids + + +def test_rdkit_fingerprints_to_bitbirch_matches_independent_schedule(size_limited_mols): + mols = size_limited_mols[:80] + generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=1024) + bits = np.asarray([generator.GetFingerprint(mol).ToList() for mol in mols], dtype=np.uint8) + packed = np.packbits(bits, axis=1, bitorder="little").view(np.uint32) + expected_labels, expected_centroids = _reference(packed, 0.25, 7, 32) + + labels, centroids = bitbirch( + packed, + 0.25, + branching_factor=7, + batch_size=32, + return_centroids=True, + ) + + np.testing.assert_array_equal(labels.numpy(), expected_labels) + np.testing.assert_array_equal(centroids.numpy(), expected_centroids) + + +def test_gpu_fingerprint_generation_and_clustering_chain_on_explicit_stream(size_limited_mols): + mols = size_limited_mols[:32] * 2 + stream = torch.cuda.Stream() + fingerprints = MorganFingerprintGenerator(radius=3, fpSize=512).GetFingerprints( + mols, + num_threads=1, + stream=stream, + ) + labels = bitbirch( + fingerprints, + threshold=1.0, + branching_factor=4, + batch_size=32, + stream=stream, + ).torch() + stream.synchronize() + + packed = fingerprints.torch().cpu().numpy().view(np.uint32) + expected, _ = _reference(packed, 1.0, 4, 32) + np.testing.assert_array_equal(labels.cpu().numpy(), expected) + + +def test_host_rdkit_and_device_nvmolkit_fingerprints_cluster_identically(size_limited_mols): + mols = size_limited_mols[:40] + rdkit_generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=256) + rdkit_bits = np.asarray([rdkit_generator.GetFingerprint(mol).ToList() for mol in mols], dtype=np.uint8) + host_fingerprints = np.packbits(rdkit_bits, axis=1, bitorder="little").view(np.uint32) + device_fingerprints = MorganFingerprintGenerator(radius=2, fpSize=256).GetFingerprints(mols, num_threads=1) + + host_labels = bitbirch(host_fingerprints, threshold=0.5, branching_factor=7, batch_size=16).numpy() + device_labels = bitbirch(device_fingerprints, threshold=0.5, branching_factor=7, batch_size=16).numpy() + + np.testing.assert_array_equal(device_labels, host_labels) diff --git a/nvmolkit/tests/test_clustering.py b/nvmolkit/tests/test_clustering.py index 58bd1a7d..b0cf5478 100644 --- a/nvmolkit/tests/test_clustering.py +++ b/nvmolkit/tests/test_clustering.py @@ -19,7 +19,7 @@ from rdkit.ML.Cluster.Butina import ClusterData import nvmolkit.clustering as clustering -from nvmolkit.clustering import ClusterDeviceResult, OutputMode, butina, fused_butina +from nvmolkit.clustering import ClusterDeviceResult, OutputMode, bitbirch, butina, fused_butina from nvmolkit.similarity import AAPMetric, CosineMetric, TanimotoMetric from nvmolkit.types import AsyncGpuResult @@ -545,3 +545,107 @@ def test_fused_butina_rdkit_and_device_outputs_agree(): assert rdkit_result == fused_butina_clusters(x, cutoff=0.4) assert [cluster[0] for cluster in rdkit_result] == device_result.centroids.numpy().tolist() assert device_result.cluster_sizes.torch().tolist() == [len(cluster) for cluster in rdkit_result] + + +# --------------------------------------------------------------------------- +# BitBIRCH tests +# --------------------------------------------------------------------------- + + +def test_bitbirch_merges_at_exact_diameter_threshold_and_returns_centroids(): + x = torch.tensor([[0b0011], [0b0001], [0b1100], [0b0100]], dtype=torch.int32, device="cuda") + labels, centroids = bitbirch(x, threshold=0.5, branching_factor=3, return_centroids=True) + torch.testing.assert_close(labels.torch().cpu(), torch.tensor([0, 0, 1, 1], dtype=torch.int32)) + torch.testing.assert_close( + centroids.torch().cpu().view(torch.int32), + torch.tensor([[0b0011], [0b1100]], dtype=torch.int32), + ) + + +@pytest.mark.parametrize("branching_factor", [3, 4, 7]) +def test_bitbirch_cascading_splits_are_deterministic(branching_factor): + x = ( + torch.ones(24, dtype=torch.int32, device="cuda") << torch.arange(24, dtype=torch.int32, device="cuda") + ).reshape(-1, 1) + labels = bitbirch(x, threshold=0.9, branching_factor=branching_factor).torch() + torch.testing.assert_close(labels, torch.arange(24, dtype=torch.int32, device="cuda")) + + +@pytest.mark.parametrize("input_kind", ["async", "cpu_tensor", "numpy"]) +def test_bitbirch_accepts_array_input_types(input_kind): + x = torch.tensor([[3], [1], [12], [4]], dtype=torch.int32, device="cuda") + if input_kind == "async": + inp = AsyncGpuResult(x) + elif input_kind == "cpu_tensor": + inp = x.cpu() + else: + inp = x.cpu().numpy() + labels = bitbirch(inp, threshold=0.5, branching_factor=3).torch().cpu() + torch.testing.assert_close(labels, torch.tensor([0, 0, 1, 1], dtype=torch.int32)) + + +def test_bitbirch_explicit_stream(): + x = torch.tensor([[3], [1], [12], [4]], dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + labels = bitbirch(x, threshold=0.5, branching_factor=3, stream=stream).torch() + stream.synchronize() + torch.testing.assert_close(labels, torch.tensor([0, 0, 1, 1], dtype=torch.int32, device="cuda")) + + +def test_bitbirch_empty_and_all_zero_fingerprints(): + empty = torch.empty((0, 3), dtype=torch.uint32, device="cuda") + empty_labels, empty_centroids = bitbirch(empty, threshold=0.5, return_centroids=True) + assert empty_labels.torch().shape == (0,) + assert empty_centroids.torch().shape == (0, 3) + + zeros = torch.zeros((12, 3), dtype=torch.uint32, device="cuda") + labels, centroids = bitbirch(zeros, threshold=1.0, branching_factor=3, return_centroids=True) + torch.testing.assert_close(labels.torch(), torch.zeros(12, dtype=torch.int32, device="cuda")) + torch.testing.assert_close(centroids.torch(), torch.zeros((1, 3), dtype=torch.uint32, device="cuda")) + + +def test_bitbirch_threshold_extremes_and_all_one_majority_centroid(): + x = torch.tensor([[0], [1], [3], [0xFFFFFFFF]], dtype=torch.uint32, device="cuda") + labels, centroids = bitbirch(x, threshold=0.0, branching_factor=3, return_centroids=True) + torch.testing.assert_close(labels.torch(), torch.zeros(4, dtype=torch.int32, device="cuda")) + torch.testing.assert_close(centroids.torch(), torch.tensor([[3]], dtype=torch.uint32, device="cuda")) + + duplicates = torch.full((17, 2), 0xFFFFFFFF, dtype=torch.uint32, device="cuda") + labels, centroids = bitbirch(duplicates, threshold=1.0, return_centroids=True) + torch.testing.assert_close(labels.torch(), torch.zeros(17, dtype=torch.int32, device="cuda")) + torch.testing.assert_close(centroids.torch(), torch.full((1, 2), 0xFFFFFFFF, dtype=torch.uint32, device="cuda")) + + +def test_bitbirch_accepts_noncontiguous_packed_input(): + storage = torch.tensor( + [[3, 99, 1, 99], [1, 99, 1, 99], [12, 99, 4, 99], [4, 99, 4, 99]], + dtype=torch.uint32, + device="cuda", + ) + x = storage[:, ::2] + assert not x.is_contiguous() + labels = bitbirch(x, threshold=0.5, branching_factor=3).torch() + torch.testing.assert_close(labels, torch.tensor([0, 0, 1, 1], dtype=torch.int32, device="cuda")) + + +@pytest.mark.parametrize( + "kwargs,error", + [ + ({"threshold": -0.1}, r"threshold must be finite and in \[0, 1\]"), + ({"threshold": np.nan}, r"threshold must be finite and in \[0, 1\]"), + ({"threshold": np.inf}, r"threshold must be finite and in \[0, 1\]"), + ({"threshold": 0.5, "branching_factor": 2}, r"branching_factor must be in \[3, 65534\]"), + ({"threshold": 0.5, "branching_factor": 65535}, r"branching_factor must be in \[3, 65534\]"), + ({"threshold": 0.5, "batch_size": 0}, "batch_size must be positive"), + ], +) +def test_bitbirch_validation(kwargs, error): + x = torch.zeros((2, 1), dtype=torch.int32) + with pytest.raises(ValueError, match=error): + bitbirch(x, **kwargs) + + +def test_bitbirch_rejects_zero_word_fingerprints(): + x = torch.empty((4, 0), dtype=torch.uint32, device="cuda") + with pytest.raises(ValueError, match="at least one fingerprint word"): + bitbirch(x, threshold=0.5)