Coverage for cosmolayer/store/_chalcedon/tanimoto_similarity.py: 100%
64 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""Tanimoto similarity computation for fingerprint arrays.
3Vendored unmodified from https://github.com/rowansci/chalcedon at commit
492da3cc5bd6ffb0d397cb49ea556f168d1d38b7e (MIT license, see ``NOTICE`` in
5this directory).
6"""
8from __future__ import annotations
10from dataclasses import InitVar, dataclass, field
11from typing import TYPE_CHECKING, Any, Literal
13import numpy as np
15if TYPE_CHECKING:
16 from numpy.typing import NDArray
18Precision = Literal["float32", "float64"] | type[np.float32] | type[np.float64]
19"""Working-precision spec: either the dtype name or the numpy scalar type."""
22def pairwise_tanimoto(
23 fingerprints: NDArray[Any],
24 dtype: Precision = "float32",
25) -> NDArray[np.floating]:
26 """Compute pairwise Tanimoto similarity matrix.
28 Uses BLAS-accelerated matrix multiplication for vectorized computation.
29 Memory usage is O(n^2) for the output matrix.
31 Args:
32 fingerprints: non-negative fingerprint matrix of shape `(n, d)`.
33 Binary, count, and positive float vectors are all supported.
34 dtype: working precision. Pass `"float64"` for higher precision at
35 ≈2x runtime and ≈2x memory.
37 Returns:
38 Symmetric similarity matrix of shape `(n, n)` with values in [0, 1].
40 Examples:
41 >>> import numpy as np
42 >>> fingerprints = np.array([[1, 1, 0], [1, 0, 1], [0, 0, 0]], dtype=np.uint8)
43 >>> similarity = pairwise_tanimoto(fingerprints)
44 >>> float(similarity[0, 1])
45 0.3333333432674408
46 >>> float(similarity[2, 0])
47 0.0
48 """
49 fingerprints_float = np.asarray(fingerprints, dtype=dtype)
50 fingerprint_count = fingerprints_float.shape[0]
51 dot_products = np.empty((fingerprint_count, fingerprint_count), dtype=fingerprints_float.dtype)
52 chunk_size = max(1, (2**31 - 1) // fingerprint_count)
53 for start in range(0, fingerprint_count, chunk_size):
54 dot_products[start : start + chunk_size] = (
55 fingerprints_float[start : start + chunk_size] @ fingerprints_float.T
56 )
57 norms = dot_products.diagonal()
58 unions = norms[:, None] + norms[None, :] - dot_products
59 # In-place divide into `dot_products`; `where=unions > 0` leaves empty-vs-empty
60 # cells untouched, and those cells already hold 0 (dot == 0 when norm == 0).
61 np.divide(dot_products, unions, out=dot_products, where=unions > 0)
62 return dot_products
65@dataclass(slots=True, frozen=True)
66class TanimotoSimilarity:
67 """Row-at-a-time Tanimoto similarity for streaming algorithms.
69 Precomputes per-fingerprint norms so each `row(i)` call needs only one
70 matrix-vector product. Memory is O(n * d); no pairwise storage.
72 Args:
73 fingerprints: non-negative fingerprint matrix of shape `(n, d)`.
74 Binary, count, and positive float vectors are all supported.
75 dtype: working precision. Pass `"float64"` for higher precision at
76 ≈2x runtime and ≈2x memory.
78 Examples:
79 >>> import numpy as np
80 >>> fingerprints = np.array([[1, 1, 0], [1, 0, 1], [0, 0, 1]], dtype=np.uint8)
81 >>> similarity = TanimotoSimilarity(fingerprints)
82 >>> similarity.fingerprint_count
83 3
84 >>> float(similarity.row(0)[1])
85 0.3333333432674408
86 """
88 fingerprints: InitVar[NDArray[Any]]
89 dtype: InitVar[Precision] = "float32"
90 _fingerprints: NDArray[np.floating] = field(init=False)
91 _norms: NDArray[np.floating] = field(init=False)
93 def __post_init__(self, fingerprints: NDArray[Any], dtype: Precision) -> None:
94 """Cast fingerprints to the requested dtype and precompute per-row norms."""
95 cast_fingerprints = np.asarray(fingerprints, dtype=dtype)
96 object.__setattr__(self, "_fingerprints", cast_fingerprints)
97 object.__setattr__(
98 self, "_norms", np.einsum("ij,ij->i", cast_fingerprints, cast_fingerprints)
99 )
101 @property
102 def fingerprint_count(self) -> int:
103 """Number of fingerprints."""
104 return self._fingerprints.shape[0]
106 @property
107 def precision(self) -> np.dtype:
108 """Working precision of the cached fingerprint matrix."""
109 return self._fingerprints.dtype
111 def row(self, index: int) -> NDArray[np.floating]:
112 """Tanimoto similarity of point `index` to all points.
114 Returns:
115 Similarity array of shape `(n,)`. Self-similarity is 0.
116 """
117 dot_products = self._fingerprints @ self._fingerprints[index]
118 unions = self._norms + self._norms[index] - dot_products
119 zero = np.asarray(0.0, dtype=self._fingerprints.dtype)
120 similarities = np.where(unions > 0, dot_products / unions, zero)
121 similarities[index] = 0.0
122 return similarities
124 def chunk(self, start: int, end: int) -> NDArray[np.floating]:
125 """Tanimoto similarity for rows `[start, end)` to all points.
127 Returns:
128 Similarity matrix of shape `(end - start, n)`. Self-similarities
129 are 0.
130 """
131 dot_products = self._fingerprints[start:end] @ self._fingerprints.T
132 unions = self._norms[start:end, None] + self._norms[None, :] - dot_products
133 np.divide(dot_products, unions, out=dot_products, where=unions > 0)
134 np.fill_diagonal(dot_products[:, start:end], 0.0)
135 return dot_products
137 def _fill_dot_products_and_unions(
138 self,
139 row_fingerprints: NDArray[np.floating],
140 column_fingerprints: NDArray[np.floating],
141 row_norms: NDArray[np.floating],
142 column_norms: NDArray[np.floating],
143 dot_products_buffer: NDArray[np.floating],
144 unions_buffer: NDArray[np.floating],
145 ) -> None:
146 """Compute `row_fingerprints @ column_fingerprints.T` and the unions in-place."""
147 np.matmul(row_fingerprints, column_fingerprints.T, out=dot_products_buffer)
148 np.add(row_norms[:, None], column_norms[None, :], out=unions_buffer)
149 unions_buffer -= dot_products_buffer
151 def _rows_neighbors_against(
152 self,
153 indices: NDArray[np.intp],
154 against: NDArray[np.intp],
155 cutoff: float,
156 dot_products_buffer: NDArray[np.floating],
157 unions_buffer: NDArray[np.floating],
158 boolean_buffer: NDArray[np.bool_],
159 ) -> NDArray[np.bool_]:
160 """Boolean matrix: distance <= cutoff for `rows[indices]` vs `rows[against]`.
162 Internal kernel. All three buffers are written in-place; callers own
163 them and reuse them across iterations to keep peak RSS bounded.
164 Self-pairs are NOT masked: if `i == against[j]` the cell will be True.
165 """
166 self._fill_dot_products_and_unions(
167 self._fingerprints[indices],
168 self._fingerprints[against],
169 self._norms[indices],
170 self._norms[against],
171 dot_products_buffer,
172 unions_buffer,
173 )
174 # Fuse `(1 - dot/union) <= cutoff` as `dot >= (1-cutoff)*union`; scaling
175 # `unions_buffer` in place avoids an extra workspace.
176 unions_buffer *= np.asarray(1.0 - cutoff, dtype=unions_buffer.dtype)
177 np.greater_equal(dot_products_buffer, unions_buffer, out=boolean_buffer)
178 # Guard empty-vs-empty pairs (union == 0 ⇒ dot == 0, which would spuriously
179 # satisfy the >=). Truthy-float AND avoids a `unions > 0` temp.
180 np.logical_and(boolean_buffer, unions_buffer, out=boolean_buffer)
181 return boolean_buffer
183 def _block_neighbors(
184 self,
185 row_start: int,
186 row_end: int,
187 column_start: int,
188 column_end: int,
189 cutoff: float,
190 dot_products_buffer: NDArray[np.floating],
191 unions_buffer: NDArray[np.floating],
192 boolean_buffer: NDArray[np.bool_],
193 ) -> NDArray[np.bool_]:
194 """Boolean matrix: Tanimoto distance <= cutoff for sub-block.
196 Internal kernel. All three buffers are written in-place; callers own
197 them and reuse them across iterations to keep peak RSS bounded.
198 Row and column ranges must be either fully aligned (diagonal block,
199 self-pairs are zeroed) or fully disjoint.
200 """
201 self._fill_dot_products_and_unions(
202 self._fingerprints[row_start:row_end],
203 self._fingerprints[column_start:column_end],
204 self._norms[row_start:row_end],
205 self._norms[column_start:column_end],
206 dot_products_buffer,
207 unions_buffer,
208 )
209 # Fuse `(1 - dot/union) <= cutoff` as `dot >= (1-cutoff)*union`; scaling
210 # `unions_buffer` in place avoids an extra workspace.
211 unions_buffer *= np.asarray(1.0 - cutoff, dtype=unions_buffer.dtype)
212 np.greater_equal(dot_products_buffer, unions_buffer, out=boolean_buffer)
213 # Guard empty-vs-empty pairs (union == 0 ⇒ dot == 0, which would spuriously
214 # satisfy the >=). Truthy-float AND avoids a `unions > 0` temp.
215 np.logical_and(boolean_buffer, unions_buffer, out=boolean_buffer)
216 if row_start == column_start:
217 np.fill_diagonal(boolean_buffer, False)
218 return boolean_buffer