Skills › Scientific Agent Skills › medchem › filter_molecules.py
filter_molecules.py
330 lines · skills/medchem/scripts/filter_molecules.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Batch molecular filtering using the medchem library.
5Usage:
6 uv run python filter_molecules.py input.csv --rules rule_of_five,rule_of_cns --pains --output filtered.csv
7 uv run python filter_molecules.py input.sdf --rules rule_of_veber --nibr --complexity 99 --output results.csv
8 uv run python filter_molecules.py smiles.txt --query 'MATCHRULE("rule_of_five") AND NOT HASALERT("pains")' --output clean.csv
9"""
11import argparse
12import sys
13from pathlib import Path
14from typing import List, Tuple
16try:
17 import pandas as pd
18 import datamol as dm
19 import medchem as mc
20 from rdkit import Chem
21 from tqdm import tqdm
22except ImportError as e:
23 print(f"Error: Missing required package: {e}")
24 print("Install dependencies: uv pip install medchem datamol pandas tqdm")
25 sys.exit(1)
28def load_molecules(
29 input_file: Path, smiles_column: str = "smiles"
30) -> Tuple[pd.DataFrame, List[Chem.Mol]]:
31 """Load molecules from CSV/TSV, SDF, or plain SMILES text files."""
32 suffix = input_file.suffix.lower()
34 if suffix == ".sdf":
35 print(f"Loading SDF file: {input_file}")
36 supplier = Chem.SDMolSupplier(str(input_file))
37 mols = [mol for mol in supplier if mol is not None]
38 data = []
39 for mol in mols:
40 props = mol.GetPropsAsDict()
41 props["smiles"] = Chem.MolToSmiles(mol)
42 data.append(props)
43 df = pd.DataFrame(data)
45 elif suffix in [".csv", ".tsv"]:
46 print(f"Loading CSV/TSV file: {input_file}")
47 sep = "\t" if suffix == ".tsv" else ","
48 df = pd.read_csv(input_file, sep=sep)
49 if smiles_column not in df.columns:
50 print(f"Error: Column '{smiles_column}' not found")
51 print(f"Available columns: {', '.join(df.columns)}")
52 sys.exit(1)
53 print("Converting SMILES to molecules...")
54 mols = [
55 dm.to_mol(smi) if isinstance(smi, str) and smi.strip() else None
56 for smi in tqdm(df[smiles_column], desc="Parsing")
57 ]
59 elif suffix == ".txt":
60 print(f"Loading text file: {input_file}")
61 with open(input_file) as f:
62 smiles_list = [line.strip() for line in f if line.strip()]
63 df = pd.DataFrame({"smiles": smiles_list})
64 print("Converting SMILES to molecules...")
65 mols = [dm.to_mol(smi) for smi in tqdm(smiles_list, desc="Parsing")]
67 else:
68 print(f"Error: Unsupported file format: {suffix}")
69 print("Supported formats: .csv, .tsv, .sdf, .txt")
70 sys.exit(1)
72 valid_indices = [i for i, mol in enumerate(mols) if mol is not None]
73 if len(valid_indices) < len(mols):
74 n_invalid = len(mols) - len(valid_indices)
75 print(f"Warning: {n_invalid} invalid molecules removed")
76 df = df.iloc[valid_indices].reset_index(drop=True)
77 mols = [mols[i] for i in valid_indices]
79 print(f"Loaded {len(mols)} valid molecules")
80 return df, mols
83def apply_rule_filters(
84 mols: List[Chem.Mol], rules: List[str], n_jobs: int
85) -> pd.DataFrame:
86 """Apply medicinal chemistry rule filters. Returns rule result columns."""
87 print(f"\nApplying rule filters: {', '.join(rules)}")
88 rfilter = mc.rules.RuleFilters(rule_list=rules)
89 results = rfilter(mols=mols, n_jobs=n_jobs, progress=True)
90 return results.drop(columns=["mol"], errors="ignore")
93def apply_common_alerts(mols: List[Chem.Mol], n_jobs: int) -> pd.DataFrame:
94 """Apply ChEMBL-derived common structural alerts."""
95 print("\nApplying common structural alerts...")
96 alert_filter = mc.structural.CommonAlertsFilters(alerts_set=["BMS"])
97 results = alert_filter(mols=mols, n_jobs=n_jobs, progress=True)
98 return results.drop(columns=["mol"], errors="ignore").rename(
99 columns={"pass_filter": "passes_common_alerts", "status": "common_alert_status",
100 "reasons": "common_alert_reasons"}
101 )
104def apply_nibr(mols: List[Chem.Mol], n_jobs: int) -> pd.DataFrame:
105 """Apply NIBR screening-deck filters."""
106 print("\nApplying NIBR filters...")
107 nibr_filter = mc.structural.NIBRFilters()
108 results = nibr_filter(mols=mols, n_jobs=n_jobs, progress=True)
109 # The class excludes explicit exclusion alerts; enforce the accumulated
110 # severity cutoff as well, matching functional.nibr_filter's default.
111 results["pass_filter"] &= results["severity"] < 10
112 return results.drop(columns=["mol"], errors="ignore").rename(
113 columns={"pass_filter": "passes_nibr", "status": "nibr_status",
114 "reasons": "nibr_reasons"}
115 )
118def apply_alert_catalog(
119 mols: List[Chem.Mol], alerts: List[str], n_jobs: int
120):
121 """Apply named alert catalogs (pains, brenk, etc.)."""
122 print(f"\nApplying alert catalogs: {', '.join(alerts)}")
123 unknown = set(alerts) - set(mc.catalogs.list_named_catalogs())
124 if unknown:
125 raise ValueError(f"Unknown named catalogs: {', '.join(sorted(unknown))}")
126 for alert in alerts:
127 passes = mc.functional.catalog_filter(
128 mols=mols, catalogs=[alert], n_jobs=n_jobs, progress=True
129 )
130 yield alert, passes
133def apply_lilly(mols: List[Chem.Mol], max_demerits: int, n_jobs: int) -> pd.DataFrame:
134 """Apply Lilly demerits; a requested but unavailable filter is an error."""
135 print(f"\nApplying Lilly demerits filter (max={max_demerits})...")
136 try:
137 passes = mc.functional.lilly_demerit_filter(
138 mols=mols, max_demerits=max_demerits, n_jobs=n_jobs, progress=True
139 )
140 except ImportError as e:
141 raise RuntimeError(
142 "Lilly filter unavailable. Install native tools with: medchem install-lilly"
143 ) from e
144 return pd.DataFrame({"passes_lilly": passes})
147def apply_complexity(
148 mols: List[Chem.Mol], limit: str, method: str, n_jobs: int
149) -> pd.DataFrame:
150 """Filter by complexity percentile threshold."""
151 print(f"\nApplying complexity filter (metric={method}, limit={limit} percentile)...")
152 passes = mc.functional.complexity_filter(
153 mols=mols,
154 complexity_metric=method,
155 limit=limit,
156 n_jobs=n_jobs,
157 progress=True,
158 )
159 return pd.DataFrame({f"passes_complexity_{method}": passes})
162def apply_query(mols: List[Chem.Mol], query: str, n_jobs: int) -> pd.DataFrame:
163 """Apply medchem query language filter."""
164 print(f"\nApplying query: {query}")
165 qf = mc.query.QueryFilter(query)
166 passes = qf(mols=mols, n_jobs=n_jobs, progress=True)
167 return pd.DataFrame({"passes_query": passes})
170def apply_groups(mols: List[Chem.Mol], groups: List[str]) -> pd.DataFrame:
171 """Detect chemical group matches."""
172 print(f"\nDetecting chemical groups: {', '.join(groups)}")
173 unknown = set(groups) - set(mc.groups.list_default_chemical_groups())
174 if unknown:
175 raise ValueError(f"Unknown chemical groups: {', '.join(sorted(unknown))}")
176 results = {}
177 for group in groups:
178 detector = mc.groups.ChemicalGroup(groups=[group])
179 results[f"has_{group}"] = [detector.has_match(mol) for mol in mols]
180 return pd.DataFrame(results)
183def generate_summary(
184 df: pd.DataFrame, output_file: Path, pass_columns: list[str] | None = None
185) -> None:
186 """Write a text summary of filtering results."""
187 summary_file = output_file.parent / f"{output_file.stem}_summary.txt"
188 pass_cols = pass_columns if pass_columns is not None else [
189 c for c in df.columns if c.startswith("passes_") or c == "pass_all"
190 ]
192 with open(summary_file, "w") as f:
193 f.write("=" * 80 + "\nMEDCHEM FILTERING SUMMARY\n" + "=" * 80 + "\n\n")
194 f.write(f"Total molecules processed: {len(df)}\n\n")
196 for col in pass_cols:
197 if col in df.columns and df[col].dtype == bool:
198 n_pass = df[col].sum()
199 pct = 100 * n_pass / len(df) if len(df) else 0
200 f.write(f" {col}: {n_pass} passed ({pct:.1f}%)\n")
202 if pass_cols:
203 bool_cols = [c for c in pass_cols if c in df.columns and df[c].dtype == bool]
204 if bool_cols:
205 n_all = df[bool_cols].all(axis=1).sum()
206 pct = 100 * n_all / len(df) if len(df) else 0
207 f.write(f"\n All filters passed: {n_all} ({pct:.1f}%)\n")
209 f.write("\n" + "=" * 80 + "\n")
211 print(f"\nSummary report saved to: {summary_file}")
214def main() -> None:
215 parser = argparse.ArgumentParser(
216 description="Batch molecular filtering using medchem",
217 formatter_class=argparse.RawDescriptionHelpFormatter,
218 epilog=__doc__,
219 )
220 parser.add_argument("input", type=Path, help="Input file (CSV, TSV, SDF, or TXT)")
221 parser.add_argument("--output", "-o", type=Path, required=True, help="Output CSV file")
222 parser.add_argument("--smiles-column", default="smiles", help="SMILES column name (default: smiles)")
224 parser.add_argument("--rules", help="Comma-separated rules (e.g. rule_of_five,rule_of_cns)")
225 parser.add_argument("--query", help='Medchem query string (e.g. MATCHRULE("rule_of_five") AND NOT HASALERT("pains"))')
226 parser.add_argument("--common-alerts", action="store_true", help="Apply BMS common structural alerts")
227 parser.add_argument("--nibr", action="store_true", help="Apply NIBR filters")
228 parser.add_argument("--lilly", action="store_true", help="Apply Lilly demerits (requires medchem install-lilly)")
229 parser.add_argument("--lilly-max", type=int, default=160, help="Max Lilly demerits (default: 160)")
230 parser.add_argument(
231 "--alerts",
232 help="Comma-separated alert catalogs (e.g. pains,brenk). Shorthand for --pains when set to pains",
233 )
234 parser.add_argument("--pains", action="store_true", help="Apply PAINS filter (alias for --alerts pains)")
236 parser.add_argument(
237 "--complexity",
238 help="Complexity percentile limit (e.g. 99 keeps below 99th percentile on ZINC-15)",
239 )
240 parser.add_argument(
241 "--complexity-method",
242 default="bertz",
243 choices=["bertz", "sas", "qed", "clogp", "whitlock", "barone", "smcm", "twc"],
244 help="Complexity metric (default: bertz)",
245 )
246 parser.add_argument("--groups", help="Comma-separated chemical groups to detect")
247 parser.add_argument("--n-jobs", type=int, default=-1, help="Parallel jobs (-1 = all cores)")
248 parser.add_argument("--no-summary", action="store_true", help="Skip summary report")
249 parser.add_argument("--filter-output", action="store_true", help="Only output molecules passing all filters")
251 args = parser.parse_args()
253 if not any([args.rules, args.query, args.common_alerts, args.nibr, args.lilly,
254 args.pains, args.alerts, args.complexity, args.groups]):
255 print("Error: Specify at least one filter (--rules, --query, --pains, --nibr, etc.)")
256 sys.exit(1)
258 df, mols = load_molecules(args.input, args.smiles_column)
259 if not mols:
260 parser.error("No valid molecules remain after parsing; no filters were evaluated")
261 result_parts = [df.reset_index(drop=True)]
263 if args.rules:
264 rule_list = [r.strip() for r in args.rules.split(",")]
265 unknown = set(rule_list) - set(mc.rules.RuleFilters.list_available_rules_names())
266 if unknown:
267 print(f"Warning: Unknown rules (will error at runtime): {', '.join(sorted(unknown))}")
268 result_parts.append(apply_rule_filters(mols, rule_list, args.n_jobs))
270 if args.query:
271 result_parts.append(apply_query(mols, args.query, args.n_jobs))
273 if args.common_alerts:
274 result_parts.append(apply_common_alerts(mols, args.n_jobs))
276 if args.nibr:
277 result_parts.append(apply_nibr(mols, args.n_jobs))
279 if args.lilly:
280 result_parts.append(apply_lilly(mols, args.lilly_max, args.n_jobs))
282 alert_list = []
283 if args.pains:
284 alert_list.append("pains")
285 if args.alerts:
286 alert_list.extend(a.strip() for a in args.alerts.split(","))
287 for alert in dict.fromkeys(alert_list):
288 col_name = f"passes_{alert}"
289 _, passes = next(apply_alert_catalog(mols, [alert], args.n_jobs))
290 result_parts.append(pd.DataFrame({col_name: passes}))
292 if args.complexity:
293 result_parts.append(
294 apply_complexity(mols, args.complexity, args.complexity_method, args.n_jobs)
295 )
297 if args.groups:
298 group_list = [g.strip() for g in args.groups.split(",")]
299 result_parts.append(apply_groups(mols, group_list))
301 df_final = pd.concat(result_parts, axis=1)
302 # Input metadata may also use passes_* names. Only freshly evaluated
303 # filters may affect retention or the scientific summary.
304 pass_cols = [
305 column for part in result_parts[1:] for column in part.columns
306 if column.startswith("passes_") or column == "pass_all"
307 ]
308 duplicate_columns = df_final.columns[df_final.columns.duplicated()].tolist()
309 if duplicate_columns:
310 parser.error(f"Input and result columns collide: {duplicate_columns}; rename input columns")
312 args.output.parent.mkdir(parents=True, exist_ok=True)
313 if not args.no_summary:
314 generate_summary(df_final, args.output, pass_columns=pass_cols)
316 if args.filter_output:
317 bool_cols = [c for c in pass_cols if c in df_final.columns and df_final[c].dtype == bool]
318 if bool_cols:
319 mask = df_final[bool_cols].all(axis=1)
320 df_final = df_final[mask]
321 print(f"\nFiltered to {len(df_final)} molecules passing all filters")
323 df_final.to_csv(args.output, index=False)
324 print(f"\nResults saved to: {args.output}")
326 print("\nDone!")
329if __name__ == "__main__":
330 main()