From 87646b8ea1a8d3007379625cf70d765ad2bb1f41 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 14 Oct 2026 13:06:36 -0400 Subject: [PATCH 1/3] Add Python BitBIRCH clustering and benchmark nvmolkit.clustering.bitbirch clusters packed int32 or uint32 fingerprints from any array input and returns labels, plus centroids on request. Tests compare labels and centroids with an independent CPU model of the batched schedule (nvmolkit/tests/_bitbirch_reference.py), whose trees are audited, across thresholds, branching factors up to 2048, multi-tile cascading splits, grouped merges past the ordered warm-up, and 4096-bit rows. benchmarks/bitbirch_clustering_bench.py times BitBIRCH on real molecules, optionally against bblean. --- benchmarks/bitbirch_clustering_bench.py | 383 ++++++++++++++++++ .../tests/test_bitbirch_clustering_bench.py | 110 +++++ docs/api/nvmolkit.rst | 1 + docs/index.rst | 5 + nvmolkit/CMakeLists.txt | 1 + nvmolkit/clustering.cpp | 33 ++ nvmolkit/clustering.py | 59 +++ nvmolkit/tests/_bitbirch_reference.py | 323 +++++++++++++++ nvmolkit/tests/test_bitbirch.py | 160 ++++++++ nvmolkit/tests/test_bitbirch_integration.py | 72 ++++ nvmolkit/tests/test_clustering.py | 105 ++++- 11 files changed, 1251 insertions(+), 1 deletion(-) create mode 100644 benchmarks/bitbirch_clustering_bench.py create mode 100644 benchmarks/tests/test_bitbirch_clustering_bench.py create mode 100644 nvmolkit/tests/_bitbirch_reference.py create mode 100644 nvmolkit/tests/test_bitbirch.py create mode 100644 nvmolkit/tests/test_bitbirch_integration.py diff --git a/benchmarks/bitbirch_clustering_bench.py b/benchmarks/bitbirch_clustering_bench.py new file mode 100644 index 00000000..3f9df073 --- /dev/null +++ b/benchmarks/bitbirch_clustering_bench.py @@ -0,0 +1,383 @@ +# 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 + + +def _cluster_count(labels: torch.Tensor) -> int: + return int(labels.max().item()) + 1 if labels.numel() else 0 + + +@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, 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) + return timing, _cluster_count(holder["labels"]) + + +@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"] + return timing, _cluster_distribution(tree.get_cluster_mol_ids(), len(fingerprints)) + + +def _cluster_distribution(clusters, num_fingerprints: int) -> dict[str, float | int]: + sizes = np.asarray([len(cluster) for cluster in clusters], dtype=np.int64) + 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, num_clusters = _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"({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), + "num_clusters": num_clusters, + } + ) + + 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..93777223 --- /dev/null +++ b/benchmarks/tests/test_bitbirch_clustering_bench.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import bitbirch_clustering_bench as bench +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([[0, 1, 2], [3], [4]], 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..28ff3f02 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -184,6 +184,11 @@ 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**: Native ordered-tree clustering for packed binary fingerprints + * Linear device storage without an all-pairs similarity matrix + * Batched snapshot routing with ordered leaf owners avoids concurrent topology mutation + * Supports bounded host-input and summary caches, host-resident labels, and optional packed majority 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..b8e670d9 100644 --- a/nvmolkit/clustering.py +++ b/nvmolkit/clustering.py @@ -17,8 +17,11 @@ 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. +``bitbirch()`` builds one BitBIRCH tree incrementally from packed fingerprints +and needs no distance matrix. """ +import math from dataclasses import dataclass from enum import Enum from typing import Literal, overload @@ -388,3 +391,59 @@ 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 with one BitBIRCH tree. + + Each insertion batch routes against a stable snapshot of the tree. One ordered + owner then updates each leaf, and summaries and splits are repaired at + barriers, so the tree is never mutated concurrently and no independent trees + are merged. + + 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 at least 3. + batch_size: Fingerprints routed per insertion batch. Larger batches + expose more parallelism but route against an older snapshot. + 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 branching_factor < 3: + raise ValueError(f"branching_factor must be at least 3, 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..8dfdb7d3 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,106 @@ 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}, "branching_factor must be at least 3"), + ({"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) From 9a1a81aa37eb8ca0eeccbddae8e8d2f0da3bf802 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 23 Oct 2026 09:36:22 -0400 Subject: [PATCH 2/3] Clarify clustering documentation and cite algorithm publications --- docs/index.rst | 5 +---- nvmolkit/clustering.py | 34 ++++++++++++++++++---------------- 2 files changed, 19 insertions(+), 20 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 28ff3f02..5fd4f6b0 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -184,10 +184,7 @@ 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**: Native ordered-tree clustering for packed binary fingerprints - * Linear device storage without an all-pairs similarity matrix - * Batched snapshot routing with ordered leaf owners avoids concurrent topology mutation - * Supports bounded host-input and summary caches, host-resident labels, and optional packed majority centroids +* **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 diff --git a/nvmolkit/clustering.py b/nvmolkit/clustering.py index b8e670d9..0e065b99 100644 --- a/nvmolkit/clustering.py +++ b/nvmolkit/clustering.py @@ -13,12 +13,11 @@ # 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. -``bitbirch()`` builds one BitBIRCH tree incrementally from packed fingerprints -and needs no distance matrix. +Butina and DISE accept precomputed distance matrices or compute distances as +needed through their ``fused_`` functions. BitBIRCH clusters packed binary +fingerprints directly. """ import math @@ -143,6 +142,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 @@ -194,6 +196,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. @@ -210,10 +215,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) @@ -275,6 +276,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 @@ -359,6 +362,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. @@ -402,20 +407,17 @@ def bitbirch( return_centroids: bool = False, stream: torch.cuda.Stream | None = None, ) -> AsyncGpuResult | tuple[AsyncGpuResult, AsyncGpuResult]: - """Cluster packed binary fingerprints with one BitBIRCH tree. + """Cluster packed binary fingerprints on the GPU using BitBIRCH. - Each insertion batch routes against a stable snapshot of the tree. One ordered - owner then updates each leaf, and summaries and splits are repaired at - barriers, so the tree is never mutated concurrently and no independent trees - are merged. + `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 at least 3. - batch_size: Fingerprints routed per insertion batch. Larger batches - expose more parallelism but route against an older snapshot. + 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. From 010b10410d8d43a776b4d77244e43b494ef06738 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Mon, 26 Oct 2026 12:55:55 -0400 Subject: [PATCH 3/3] Address BitBIRCH review: report nvMolKit cluster distribution, bound branching_factor The benchmark now reports the same cluster-size distribution for nvMolKit as for bblean, so throughput comparisons show whether the clusterings differ. bitbirch() validates and documents the 65534 branching_factor ceiling enforced by the C++ implementation instead of surfacing its generic error. --- benchmarks/bitbirch_clustering_bench.py | 21 ++++++++----------- .../tests/test_bitbirch_clustering_bench.py | 3 ++- nvmolkit/clustering.py | 7 ++++--- nvmolkit/tests/test_clustering.py | 3 ++- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/benchmarks/bitbirch_clustering_bench.py b/benchmarks/bitbirch_clustering_bench.py index 3f9df073..a1f1d043 100644 --- a/benchmarks/bitbirch_clustering_bench.py +++ b/benchmarks/bitbirch_clustering_bench.py @@ -47,10 +47,6 @@ def _load_bblean(enabled: bool) -> tuple[ModuleType | None, str, str]: return module, "available", version -def _cluster_count(labels: torch.Tensor) -> int: - return int(labels.max().item()) + 1 if labels.numel() else 0 - - @nvtx.annotate("bench_nvmolkit_bitbirch", color="red") def _bench_nvmolkit( fingerprints: torch.Tensor, @@ -60,7 +56,7 @@ def _bench_nvmolkit( batch_size: int, runs: int, warmups: int, -) -> tuple[object, int]: +) -> tuple[object, dict[str, float | int]]: holder: dict[str, torch.Tensor] = {} @nvtx.annotate("bitbirch_nvmolkit_run", color="orange") @@ -73,7 +69,8 @@ def run() -> None: ).torch() timing = time_it(run, runs=runs, warmups=warmups, gpu_sync=True) - return timing, _cluster_count(holder["labels"]) + labels = holder["labels"] + return timing, _cluster_distribution(torch.bincount(labels).cpu().numpy(), labels.numel()) @nvtx.annotate("bench_bblean_bitbirch", color="green") @@ -110,11 +107,11 @@ def run(_deadline=None) -> None: progress_target=1, ) tree = holder["tree"] - return timing, _cluster_distribution(tree.get_cluster_mol_ids(), len(fingerprints)) + 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(clusters, num_fingerprints: int) -> dict[str, float | int]: - sizes = np.asarray([len(cluster) for cluster in clusters], dtype=np.int64) +def _cluster_distribution(sizes: np.ndarray, num_fingerprints: int) -> dict[str, float | int]: if not sizes.size: return { "num_clusters": 0, @@ -292,7 +289,7 @@ def main() -> None: ) torch.cuda.cudart().cudaProfilerStart() try: - timing, num_clusters = _bench_nvmolkit( + timing, distribution = _bench_nvmolkit( device_fps, threshold=threshold, branching_factor=args.branching_factor, @@ -304,7 +301,7 @@ def main() -> None: torch.cuda.cudart().cudaProfilerStop() print( f" Finished nvMolKit GPU: {timing.mean_ms:.3f} +/- {timing.std_ms:.3f} ms " - f"({num_clusters:,} clusters)", + f"({distribution['num_clusters']:,} clusters)", flush=True, ) configuration_rows.append( @@ -314,7 +311,7 @@ def main() -> None: "batch_size": args.batch_size, "device": torch.cuda.get_device_name(device_fps.device), **_timing_fields(timing), - "num_clusters": num_clusters, + **distribution, } ) diff --git a/benchmarks/tests/test_bitbirch_clustering_bench.py b/benchmarks/tests/test_bitbirch_clustering_bench.py index 93777223..5ec90d41 100644 --- a/benchmarks/tests/test_bitbirch_clustering_bench.py +++ b/benchmarks/tests/test_bitbirch_clustering_bench.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import bitbirch_clustering_bench as bench +import numpy as np import pytest @@ -97,7 +98,7 @@ def missing_import(_name): def test_cluster_distribution_reports_sizes_and_singletons(): - distribution = bench._cluster_distribution([[0, 1, 2], [3], [4]], 5) + distribution = bench._cluster_distribution(np.array([3, 1, 1]), 5) assert distribution["num_clusters"] == 3 assert distribution["mean_cluster_size"] == pytest.approx(5 / 3) diff --git a/nvmolkit/clustering.py b/nvmolkit/clustering.py index 0e065b99..d5a7e6a6 100644 --- a/nvmolkit/clustering.py +++ b/nvmolkit/clustering.py @@ -35,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, ...], ...] @@ -416,7 +417,7 @@ def bitbirch( 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 at least 3. + 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)``. @@ -432,8 +433,8 @@ def bitbirch( """ if not math.isfinite(threshold) or not 0 <= threshold <= 1: raise ValueError("threshold must be finite and in [0, 1]") - if branching_factor < 3: - raise ValueError(f"branching_factor must be at least 3, got {branching_factor}") + 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) diff --git a/nvmolkit/tests/test_clustering.py b/nvmolkit/tests/test_clustering.py index 8dfdb7d3..b0cf5478 100644 --- a/nvmolkit/tests/test_clustering.py +++ b/nvmolkit/tests/test_clustering.py @@ -634,7 +634,8 @@ def test_bitbirch_accepts_noncontiguous_packed_input(): ({"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}, "branching_factor must be at least 3"), + ({"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"), ], )