Skills › Scientific Agent Skills › rdkit › similarity_search.py
similarity_search.py
295 lines · skills/rdkit/scripts/similarity_search.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Molecular Similarity Search
5Perform fingerprint-based similarity screening against a database of molecules.
6Supports multiple fingerprint types and similarity metrics.
8Usage:
9 python similarity_search.py "CCO" database.smi --threshold 0.7
10 python similarity_search.py query.smi database.sdf --method morgan --output hits.csv
11"""
13import argparse
14import sys
15import math
16from pathlib import Path
18try:
19 from rdkit import Chem, rdBase
20 from rdkit.Chem import MACCSkeys, rdFingerprintGenerator
21 from rdkit import DataStructs
22except ImportError:
23 print("Error: RDKit not installed. Install with: uv pip install rdkit")
24 sys.exit(1)
27from _common import read_molecule_records, molecule_smiles
30FINGERPRINT_METHODS = {
31 'morgan': 'Morgan fingerprint (ECFP-like)',
32 'rdkit': 'RDKit topological fingerprint',
33 'maccs': 'MACCS structural keys',
34 'atompair': 'Atom pair fingerprint',
35 'torsion': 'Topological torsion fingerprint'
36}
39def generate_fingerprint(mol, method='morgan', radius=2, n_bits=2048,
40 include_chirality=False):
41 """Generate molecular fingerprint based on specified method."""
42 if mol is None or mol.GetNumAtoms() == 0:
43 return None
45 method = method.lower()
46 if radius < 0 or n_bits <= 0:
47 raise ValueError("radius must be nonnegative and n_bits must be positive")
48 if include_chirality and method in {'rdkit', 'maccs'}:
49 raise ValueError(f"{method} does not support include_chirality")
51 if method == 'morgan':
52 gen = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits, includeChirality=include_chirality)
53 return gen.GetFingerprint(mol)
54 elif method == 'rdkit':
55 gen = rdFingerprintGenerator.GetRDKitFPGenerator(maxPath=7, fpSize=n_bits)
56 return gen.GetFingerprint(mol)
57 elif method == 'maccs':
58 return MACCSkeys.GenMACCSKeys(mol)
59 elif method == 'atompair':
60 gen = rdFingerprintGenerator.GetAtomPairGenerator(fpSize=n_bits, includeChirality=include_chirality)
61 return gen.GetFingerprint(mol)
62 elif method == 'torsion':
63 gen = rdFingerprintGenerator.GetTopologicalTorsionGenerator(fpSize=n_bits, includeChirality=include_chirality)
64 return gen.GetFingerprint(mol)
65 else:
66 raise ValueError(f"Unknown fingerprint method: {method}")
69def load_molecules(file_path):
70 """Load molecules from file."""
71 return [
72 {'index': index,
73 'name': mol.GetProp('_Name') if mol.HasProp('_Name') else f'Mol_{index}',
74 'smiles': molecule_smiles(mol), 'mol': mol}
75 for index, mol in read_molecule_records(file_path)
76 ]
79def similarity_search(query_mol, database, method='morgan', threshold=0.7,
80 radius=2, n_bits=2048, metric='tanimoto', include_chirality=False):
81 """
82 Perform similarity search.
84 Args:
85 query_mol: Query molecule (RDKit Mol object)
86 database: List of database molecules
87 method: Fingerprint method
88 threshold: Similarity threshold (0-1)
89 radius: Morgan fingerprint radius
90 n_bits: Fingerprint size
91 metric: Similarity metric (tanimoto, dice, cosine)
93 Returns:
94 List of hits with similarity scores
95 """
96 if not math.isfinite(threshold) or not 0 <= threshold <= 1:
97 raise ValueError("threshold must be finite and between 0 and 1")
98 if query_mol is None or query_mol.GetNumAtoms() == 0:
99 print("Error: Invalid query molecule")
100 return []
102 # Generate query fingerprint
103 query_fp = generate_fingerprint(query_mol, method, radius, n_bits, include_chirality)
104 if query_fp is None:
105 print("Error: Failed to generate query fingerprint")
106 return []
108 # Choose similarity function
109 if metric.lower() == 'tanimoto':
110 sim_func = DataStructs.TanimotoSimilarity
111 elif metric.lower() == 'dice':
112 sim_func = DataStructs.DiceSimilarity
113 elif metric.lower() == 'cosine':
114 sim_func = DataStructs.CosineSimilarity
115 else:
116 raise ValueError(f"Unknown similarity metric: {metric}")
118 # Search database
119 hits = []
120 for db_entry in database:
121 db_fp = generate_fingerprint(db_entry['mol'], method, radius, n_bits, include_chirality)
122 if db_fp is None:
123 continue
125 similarity = sim_func(query_fp, db_fp)
127 if similarity >= threshold:
128 hits.append({
129 'index': db_entry['index'],
130 'name': db_entry['name'],
131 'smiles': db_entry['smiles'],
132 'similarity': similarity
133 })
135 # Sort by similarity (descending)
136 hits.sort(key=lambda x: x['similarity'], reverse=True)
138 return hits
141def write_results(hits, output_file):
142 """Write results to CSV file."""
143 import csv
145 with open(output_file, 'w', newline='') as f:
146 fieldnames = ['Rank', 'Index', 'Name', 'SMILES', 'Similarity']
147 writer = csv.DictWriter(f, fieldnames=fieldnames)
148 writer.writeheader()
150 for rank, hit in enumerate(hits, 1):
151 writer.writerow({
152 'Rank': rank,
153 'Index': hit['index'],
154 'Name': hit['name'],
155 'SMILES': hit['smiles'],
156 'Similarity': f"{hit['similarity']:.4f}"
157 })
160def print_results(hits, max_display=20):
161 """Print results to console."""
162 if not hits:
163 print("\nNo hits found above threshold")
164 return
166 print(f"\nFound {len(hits)} similar molecules:")
167 print("="*80)
168 print(f"{'Rank':<6} {'Index':<8} {'Similarity':<12} {'Name':<20} {'SMILES'}")
169 print("-"*80)
171 for rank, hit in enumerate(hits[:max_display], 1):
172 name = hit['name'][:18] + '..' if len(hit['name']) > 20 else hit['name']
173 smiles = hit['smiles'][:40] + '...' if len(hit['smiles']) > 43 else hit['smiles']
174 print(f"{rank:<6} {hit['index']:<8} {hit['similarity']:<12.4f} {name:<20} {smiles}")
176 if len(hits) > max_display:
177 print(f"\n... and {len(hits) - max_display} more")
179 print("="*80)
182def main():
183 parser = argparse.ArgumentParser(
184 description='Molecular similarity search using fingerprints',
185 formatter_class=argparse.RawDescriptionHelpFormatter,
186 epilog=f"""
187Available fingerprint methods:
188{chr(10).join(f' {k:12s} - {v}' for k, v in FINGERPRINT_METHODS.items())}
190Similarity metrics:
191 tanimoto - Tanimoto coefficient (default)
192 dice - Dice coefficient
193 cosine - Cosine similarity
195Examples:
196 # Search with SMILES query
197 python similarity_search.py "CCO" database.smi --threshold 0.7
199 # Use different fingerprint
200 python similarity_search.py query.smi database.sdf --method maccs
202 # Save results
203 python similarity_search.py "c1ccccc1" database.smi --output hits.csv
205 # Adjust Morgan radius
206 python similarity_search.py "CCO" database.smi --method morgan --radius 3
207 """
208 )
210 parser.add_argument('query', help='Query SMILES or file')
211 parser.add_argument('database', help='Database file (SDF or SMILES)')
212 parser.add_argument('--method', '-m', default='morgan',
213 choices=FINGERPRINT_METHODS.keys(),
214 help='Fingerprint method (default: morgan)')
215 parser.add_argument('--threshold', '-t', type=float, default=0.7,
216 help='Similarity threshold (default: 0.7)')
217 parser.add_argument('--radius', '-r', type=int, default=2,
218 help='Morgan fingerprint radius (default: 2)')
219 parser.add_argument('--bits', '-b', type=int, default=2048,
220 help='Fingerprint size (default: 2048)')
221 parser.add_argument('--chirality', action='store_true',
222 help='Include chirality (Morgan, atom pair, torsion only)')
223 parser.add_argument('--metric', default='tanimoto',
224 choices=['tanimoto', 'dice', 'cosine'],
225 help='Similarity metric (default: tanimoto)')
226 parser.add_argument('--output', '-o', help='Output CSV file')
227 parser.add_argument('--max-display', type=int, default=20,
228 help='Maximum hits to display (default: 20)')
230 args = parser.parse_args()
232 if not math.isfinite(args.threshold) or not 0 <= args.threshold <= 1:
233 parser.error("threshold must be finite and between 0 and 1")
234 if args.radius < 0 or args.bits <= 0 or args.max_display <= 0:
235 parser.error("radius must be nonnegative; bits and max-display must be positive")
236 if args.chirality and args.method in {'rdkit', 'maccs'}:
237 parser.error(f"{args.method} does not support --chirality")
239 # Load query
240 query_path = Path(args.query)
241 if query_path.exists():
242 # Query is a file
243 query_mols = load_molecules(args.query)
244 if not query_mols:
245 print("Error: No valid molecules in query file")
246 sys.exit(1)
247 if len(query_mols) != 1:
248 parser.error('Query file must contain exactly one valid molecule')
249 query_mol = query_mols[0]['mol']
250 query_smiles = query_mols[0]['smiles']
251 else:
252 # Query is SMILES string
253 query_mol = Chem.MolFromSmiles(args.query)
254 query_smiles = args.query
255 if query_mol is None or query_mol.GetNumAtoms() == 0:
256 print(f"Error: Failed to parse query SMILES: {args.query}")
257 sys.exit(1)
259 print(f"Query: {query_smiles}")
260 print(f"RDKit: {rdBase.rdkitVersion}; method={args.method}; radius={args.radius}; "
261 f"bits={167 if args.method == 'maccs' else args.bits}; "
262 f"chirality={args.chirality}; metric={args.metric}")
263 print(f"Threshold: {args.threshold}")
264 print(f"Loading database: {args.database}...")
266 # Load database
267 database = load_molecules(args.database)
268 if not database:
269 print("Error: No valid molecules in database")
270 sys.exit(1)
272 print(f"Loaded {len(database)} molecules")
273 print("Searching...")
275 # Perform search
276 hits = similarity_search(
277 query_mol, database,
278 method=args.method,
279 threshold=args.threshold,
280 radius=args.radius,
281 n_bits=args.bits,
282 metric=args.metric,
283 include_chirality=args.chirality
284 )
286 # Output results
287 if args.output:
288 write_results(hits, args.output)
289 print(f"\nResults written to: {args.output}")
291 print_results(hits, args.max_display)
294if __name__ == '__main__':
295 main()