Coverage for cosmolayer/store/_chalcedon/butina_cluster.py: 87%

1"""Butina clustering algorithm (Butina, JCICS 39 747-750, 1999). 

2 

3Vendored from https://github.com/rowansci/chalcedon at commit 

492da3cc5bd6ffb0d397cb49ea556f168d1d38b7e (MIT license, see ``NOTICE`` in 

5this directory). Modified from upstream to add the ``progress`` parameter 

6gating the tqdm bars, so a library call doesn't print to stderr by 

7default; to use ``tqdm.auto``; and to leave bars on screen when enabled. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import TYPE_CHECKING 

13 

14import numpy as np 

15from tqdm.auto import tqdm 

16 

17from cosmolayer.store._chalcedon.tanimoto_similarity import ( 

18 Precision, 

19 TanimotoSimilarity, 

20) 

21 

22if TYPE_CHECKING: 

23 from numpy.typing import NDArray 

24 

25_DEFAULT_COUNT_BLOCK_SIZE = 2500 

26_DEFAULT_ASSIGN_BATCH_SIZE = 500 

27_DIAGONAL_SPLIT_MINIMUM = 512 # diagonal blocks below this size aren't split 

28 

29 

30def butina_cluster( 

31 fingerprints: NDArray[np.integer | np.floating], 

32 cutoff: float = 0.65, 

33 count_block_size: int = _DEFAULT_COUNT_BLOCK_SIZE, 

34 assign_batch_size: int = _DEFAULT_ASSIGN_BATCH_SIZE, 

35 dtype: Precision = "float32", 

36 progress: bool = False, 

37) -> NDArray[np.intp]: 

38 """Cluster fingerprints using the Butina algorithm. 

39 

40 Count-sort-assign strategy: count neighbors via upper-triangle BLAS 

41 blocks, sort by count descending, greedily assign in batched sgemm 

42 passes. Produces clusters matching RDKit's reference implementation at 

43 typical cheminformatics cutoffs; at uncommon cutoffs a small fraction 

44 of boundary-pair decisions may differ due to float rounding. 

45 

46 Args: 

47 fingerprints: non-negative fingerprint matrix of shape `(n, d)`. 

48 Binary, count, and positive float vectors are all supported. 

49 cutoff: Tanimoto distance cutoff; pairs with distance ≤ `cutoff` 

50 are neighbors. 

51 count_block_size: side length of the count-phase BLAS blocks. 

52 Workspace cost scales as O(count_block_size**2). 

53 assign_batch_size: number of centers per assign-phase sgemm call. 

54 Workspace cost scales as O(assign_batch_size * n). 

55 dtype: working precision. Pass `"float64"` for higher precision at 

56 ≈2x runtime and ≈2x memory. 

57 progress: whether to print tqdm progress bars to stderr. Not part 

58 of upstream chalcedon; added for library use. Default False. 

59 

60 Returns: 

61 Cluster ID per point, shape `(n,)`. Cluster 0 is the largest. 

62 

63 Examples: 

64 >>> import numpy as np 

65 >>> fingerprints = np.array([ 

66 ... [1, 1, 0, 0], 

67 ... [1, 1, 1, 0], 

68 ... [0, 0, 1, 1], 

69 ... [0, 0, 0, 1], 

70 ... ], dtype=np.uint8) 

71 >>> butina_cluster(fingerprints, cutoff=0.5).tolist() 

72 [1, 1, 0, 0] 

73 """ 

74 similarity = TanimotoSimilarity(fingerprints, dtype=dtype) 

75 fingerprint_count = similarity.fingerprint_count 

76 block_size = min(count_block_size, fingerprint_count) 

77 batch_size = min(assign_batch_size, fingerprint_count) 

78 

79 # Shared workspace, reshaped per block. 

80 max_workspace = max(block_size * block_size, batch_size * fingerprint_count) 

81 dot_products_flat = np.empty(max_workspace, dtype=similarity.precision) 

82 unions_flat = np.empty(max_workspace, dtype=similarity.precision) 

83 boolean_flat = np.empty(max_workspace, dtype=bool) 

84 

85 # Phase 1: count neighbors via upper-triangle blocks. Diagonal blocks split 

86 # recursively for ssyrk-equivalent FLOPs savings in pure numpy. 

87 neighbor_counts = np.zeros(fingerprint_count, dtype=np.intp) 

88 

89 def count_block(row_start: int, row_end: int, column_start: int, column_end: int) -> None: 

90 """Add neighbor counts contributed by one upper-triangle block.""" 

91 row_count = row_end - row_start 

92 column_count = column_end - column_start 

93 size = row_count * column_count 

94 neighbors = similarity._block_neighbors( 

95 row_start, 

96 row_end, 

97 column_start, 

98 column_end, 

99 cutoff, 

100 dot_products_flat[:size].reshape(row_count, column_count), 

101 unions_flat[:size].reshape(row_count, column_count), 

102 boolean_flat[:size].reshape(row_count, column_count), 

103 ) 

104 neighbor_counts[row_start:row_end] += np.count_nonzero(neighbors, axis=1) 

105 if row_start != column_start: 

106 neighbor_counts[column_start:column_end] += np.count_nonzero(neighbors, axis=0) 

107 

108 def count_diagonal(start: int, end: int) -> None: 

109 """Count neighbors in a diagonal block, recursively splitting if large enough.""" 

110 if end - start >= _DIAGONAL_SPLIT_MINIMUM: 

111 middle = start + (end - start) // 2 

112 count_diagonal(start, middle) 

113 count_diagonal(middle, end) 

114 count_block(start, middle, middle, end) 

115 else: 

116 count_block(start, end, start, end) 

117 

118 for row_start in tqdm( 

119 range(0, fingerprint_count, block_size), 

120 desc="Counting neighbors", 

121 leave=True, 

122 disable=not progress, 

123 ): 

124 row_end = min(row_start + block_size, fingerprint_count) 

125 count_diagonal(row_start, row_end) 

126 for column_start in range(row_start + block_size, fingerprint_count, block_size): 

127 column_end = min(column_start + block_size, fingerprint_count) 

128 count_block(row_start, row_end, column_start, column_end) 

129 

130 # Phase 2: sort by count desc, ties by higher index first. 

131 order = np.lexsort((-np.arange(fingerprint_count, dtype=np.intp), -neighbor_counts)) 

132 

133 # Phase 3: batched greedy assign against a shrinking compact unassigned list. 

134 cluster_id = np.full(fingerprint_count, -1, dtype=np.intp) 

135 unassigned = np.arange(fingerprint_count, dtype=np.intp) 

136 next_cluster_id = 0 

137 cursor = 0 

138 with tqdm( 

139 total=fingerprint_count, 

140 desc="Assigning clusters", 

141 leave=True, 

142 disable=not progress, 

143 ) as progress_bar: 

144 while cursor < fingerprint_count: 

145 pending = order[cursor:] 

146 pending_positions = np.flatnonzero(cluster_id[pending] == -1) 

147 if len(pending_positions) == 0: 

148 progress_bar.update(fingerprint_count - cursor) 

149 break 

150 center_count = min(batch_size, len(pending_positions)) 

151 step = int(pending_positions[center_count - 1]) + 1 

152 centers = pending[pending_positions[:center_count]] 

153 cursor += step 

154 progress_bar.update(step) 

155 

156 unassigned_count = len(unassigned) 

157 size = center_count * unassigned_count 

158 is_neighbor = similarity._rows_neighbors_against( 

159 centers, 

160 unassigned, 

161 cutoff, 

162 dot_products_flat[:size].reshape(center_count, unassigned_count), 

163 unions_flat[:size].reshape(center_count, unassigned_count), 

164 boolean_flat[:size].reshape(center_count, unassigned_count), 

165 ) 

166 # Avoids recomputing `cluster_id[unassigned] == -1` per center. 

167 still_unassigned = np.ones(unassigned_count, dtype=bool) 

168 for batch_index, center in enumerate(centers.tolist()): 

169 if cluster_id[center] != -1: 

170 continue 

171 # Drop points already claimed by earlier centers in this batch. 

172 member_mask = is_neighbor[batch_index] & still_unassigned 

173 cluster_id[unassigned[member_mask]] = next_cluster_id 

174 still_unassigned[member_mask] = False 

175 next_cluster_id += 1 

176 

177 unassigned = np.flatnonzero(cluster_id == -1) 

178 

179 return cluster_id