Coverage for cosmolayer/store/clustering.py: 100%
62 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""Butina clustering of molecules by Morgan-fingerprint Tanimoto
2similarity.
4Used by ``SegmentStore.from_cosmo_files`` to assign each stored molecule a
5``cluster_id`` (see the ``molecules.parquet`` column of that name), so
6callers can build train/test splits or diversity subsets that don't leak
7near-duplicate structures across a split.
8"""
10from dataclasses import dataclass
11from typing import cast
13import numpy as np
14from numpy.typing import NDArray
15from rdkit import Chem, rdBase
16from rdkit.Chem import rdFingerprintGenerator
17from tqdm.auto import tqdm
19from cosmolayer.store._chalcedon.butina_cluster import (
20 butina_cluster as _chalcedon_butina_cluster,
21)
22from cosmolayer.store._chalcedon.tanimoto_similarity import TanimotoSimilarity
25@dataclass(frozen=True)
26class ClusteringSpecs:
27 """Parameters for fingerprinting and Butina-clustering molecules.
29 Parameters
30 ----------
31 cutoff : float
32 Tanimoto distance threshold: molecules within ``cutoff`` of a
33 cluster centroid join that cluster. Default 0.65.
34 radius : int
35 Morgan fingerprint radius. Default 2.
36 fp_size : int
37 Morgan fingerprint bit-vector length. Default 2048.
38 include_chirality : bool
39 Whether the fingerprint distinguishes stereoisomers. Default True.
40 """
42 cutoff: float = 0.65
43 radius: int = 2
44 fp_size: int = 2048
45 include_chirality: bool = True
48class FingerprintGenerator:
49 """Generates Morgan fingerprints for molecules under a fixed
50 ``ClusteringSpecs``.
52 Parameters
53 ----------
54 specs : ClusteringSpecs
55 Fingerprint parameters (radius, size, chirality).
56 """
58 def __init__(self, specs: ClusteringSpecs) -> None:
59 self.specs = specs
60 self._generator = rdFingerprintGenerator.GetMorganGenerator(
61 radius=specs.radius,
62 fpSize=specs.fp_size,
63 includeChirality=specs.include_chirality,
64 )
66 def generate(self, mol: Chem.Mol) -> NDArray[np.int8]:
67 """Generate a molecule's fingerprint as a dense bit array.
69 Parameters
70 ----------
71 mol : Chem.Mol
72 Molecule to fingerprint.
74 Returns
75 -------
76 np.ndarray, shape (fp_size,)
77 Dense 0/1 bit array.
78 """
79 with rdBase.BlockLogs():
80 fp = self._generator.GetFingerprintAsNumPy(mol).astype(np.int8)
81 return cast(NDArray[np.int8], fp)
84def butina_cluster(
85 fingerprints: NDArray[np.int8], cutoff: float, *, progress: bool = False
86) -> NDArray[np.int64]:
87 """Butina-cluster molecules by fingerprint Tanimoto distance.
89 Delegates to chalcedon's count-sort-assign Butina implementation
90 (vendored in ``cosmolayer.store._chalcedon``), which produces
91 the same partition as RDKit's reference implementation at typical
92 cheminformatics cutoffs but scales substantially better with ``n``.
93 Pairwise work is still effectively ``O(n**2)``, so this is intended
94 for dataset-scale molecule counts rather than unbounded streaming.
96 Parameters
97 ----------
98 fingerprints : np.ndarray, shape (n, fp_size)
99 One fingerprint per molecule, as produced by
100 ``FingerprintGenerator.generate``.
101 cutoff : float
102 Tanimoto distance threshold: molecules within ``cutoff`` of a
103 cluster centroid join that cluster.
104 progress : bool, optional
105 If True, show chalcedon's tqdm bars (neighbor counting and
106 cluster assignment) on stderr. Default False, so a library call
107 stays quiet.
109 Returns
110 -------
111 np.ndarray, shape (n,)
112 Cluster id per molecule, in ``[0, num_clusters)``. Cluster ids are
113 ordered by cluster formation order (largest neighbor lists first,
114 per Butina's algorithm), not by input order.
115 """
116 n = fingerprints.shape[0]
117 if n == 0:
118 return np.empty(n, dtype=np.int64)
119 if n == 1:
120 return np.zeros(1, dtype=np.int64)
122 cluster_ids = _chalcedon_butina_cluster(
123 fingerprints, cutoff=cutoff, progress=progress
124 )
125 return cluster_ids.astype(np.int64)
128# Bound (batch, k) Tanimoto workspace so large clusters never allocate k×k.
129_MEDOID_SCORE_MAX_CELLS = 4_000_000
132def cluster_medoid_distances(
133 fingerprints: NDArray[np.int8],
134 cluster_ids: NDArray[np.int64],
135 *,
136 progress: bool = False,
137) -> NDArray[np.float64]:
138 """Tanimoto distance of each molecule to its cluster's Tanimoto medoid.
140 The medoid of a cluster is the member that maximizes the sum of
141 Tanimoto similarities to the other members (equivalently, minimizes
142 the sum of Tanimoto distances). Ties break to the lowest row index.
143 The medoid's own distance is exactly 0.
145 Parameters
146 ----------
147 fingerprints : np.ndarray, shape (n, fp_size)
148 One fingerprint per molecule, as produced by
149 ``FingerprintGenerator.generate``.
150 cluster_ids : np.ndarray, shape (n,)
151 Cluster id per molecule, as produced by ``butina_cluster``.
152 progress : bool, optional
153 If True, show a tqdm bar over clusters. Default False.
155 Returns
156 -------
157 np.ndarray, shape (n,)
158 Tanimoto distance to that molecule's cluster medoid, ``float64``.
160 Raises
161 ------
162 ValueError
163 If ``fingerprints`` and ``cluster_ids`` have different lengths.
164 """
165 n = fingerprints.shape[0]
166 if cluster_ids.shape[0] != n:
167 raise ValueError(
168 "fingerprints and cluster_ids must have the same length, "
169 f"got {fingerprints.shape[0]} and {cluster_ids.shape[0]}."
170 )
171 distances = np.zeros(n, dtype=np.float64)
172 if n == 0:
173 return distances
175 fps_all = np.asarray(fingerprints, dtype=np.float64)
176 for cluster_id in tqdm(
177 np.unique(cluster_ids),
178 desc="Computing medoids",
179 disable=not progress,
180 ):
181 members = np.flatnonzero(cluster_ids == cluster_id)
182 k = int(members.shape[0])
183 if k == 1:
184 continue
185 cluster_fps = fps_all[members]
186 similarity = TanimotoSimilarity(cluster_fps, dtype="float64")
187 batch = max(1, min(k, _MEDOID_SCORE_MAX_CELLS // k))
188 scores = np.empty(k, dtype=np.float64)
189 for start in range(0, k, batch):
190 end = min(start + batch, k)
191 scores[start:end] = similarity.chunk(start, end).sum(axis=1)
192 local_medoid = int(np.argmax(scores))
193 medoid_fp = cluster_fps[local_medoid]
194 norms = np.einsum("ij,ij->i", cluster_fps, cluster_fps)
195 dots = cluster_fps @ medoid_fp
196 unions = norms + norms[local_medoid] - dots
197 sims = np.divide(dots, unions, out=np.zeros_like(dots), where=unions > 0)
198 dists = 1.0 - sims
199 dists[local_medoid] = 0.0
200 distances[members] = dists
201 return distances