#!/usr/bin/env python3
"""
Batch molecular filtering using medchem library.
This script provides a production-ready workflow for filtering compound libraries
using medchem rules, structural alerts, and custom constraints.
Usage:
python filter_molecules.py input.csv --rules rule_of_five,rule_of_cns --alerts nibr --output filtered.csv
python filter_molecules.py input.sdf --rules rule_of_drug --lilly --complexity 400 --output results.csv
python filter_molecules.py smiles.txt --nibr --pains --n-jobs -1 --output clean.csv
"""
import argparse
import sys
from pathlib import Path
from typing import List, Dict, Optional, Tuple
import json
try:
import pandas as pd
import datamol as dm
import medchem as mc
from rdkit import Chem
from tqdm import tqdm
except ImportError as e:
print(f"Error: Missing required package: {e}")
print("Install dependencies: pip install medchem datamol pandas tqdm")
sys.exit(1)
def load_molecules(input_file: Path, smiles_column: str = "smiles") -> Tuple[pd.DataFrame, List[Chem.Mol]]:
"""
Load molecules from various file formats.
Supports:
- CSV/TSV with SMILES column
- SDF files
- Plain text files with one SMILES per line
Returns:
Tuple of (DataFrame with metadata, list of RDKit molecules)
"""
suffix = input_file.suffix.lower()
if suffix == ".sdf":
print(f"Loading SDF file: {input_file}")
supplier = Chem.SDMolSupplier(str(input_file))
mols = [mol for mol in supplier if mol is not None]
# Create DataFrame from SDF properties
data = []
for mol in mols:
props = mol.GetPropsAsDict()
props["smiles"] = Chem.MolToSmiles(mol)
data.append(props)
df = pd.DataFrame(data)
elif suffix in [".csv", ".tsv"]:
print(f"Loading CSV/TSV file: {input_file}")
sep = "\t" if suffix == ".tsv" else ","
df = pd.read_csv(input_file, sep=sep)
if smiles_column not in df.columns:
print(f"Error: Column '{smiles_column}' not found in file")
print(f"Available columns: {', '.join(df.columns)}")
sys.exit(1)
print(f"Converting SMILES to molecules...")
mols = [dm.to_mol(smi) for smi in tqdm(df[smiles_column], desc="Parsing")]
elif suffix == ".txt":
print(f"Loading text file: {input_file}")
with open(input_file) as f:
smiles_list = [line.strip() for line in f if line.strip()]
df = pd.DataFrame({"smiles": smiles_list})
print(f"Converting SMILES to molecules...")
mols = [dm.to_mol(smi) for smi in tqdm(smiles_list, desc="Parsing")]
else:
print(f"Error: Unsupported file format: {suffix}")
print("Supported formats: .csv, .tsv, .sdf, .txt")
sys.exit(1)
# Filter out invalid molecules
valid_indices = [i for i, mol in enumerate(mols) if mol is not None]
if len(valid_indices) < len(mols):
n_invalid = len(mols) - len(valid_indices)
print(f"Warning: {n_invalid} invalid molecules removed")
df = df.iloc[valid_indices].reset_index(drop=True)
mols = [mols[i] for i in valid_indices]
print(f"Loaded {len(mols)} valid molecules")
return df, mols
def apply_rule_filters(mols: List[Chem.Mol], rules: List[str], n_jobs: int) -> pd.DataFrame:
"""Apply medicinal chemistry rule filters."""
print(f"\nApplying rule filters: {', '.join(rules)}")
rfilter = mc.rules.RuleFilters(rule_list=rules)
results = rfilter(mols=mols, n_jobs=n_jobs, progress=True)
# Convert to DataFrame
df_results = pd.DataFrame(results)
# Add summary column
df_results["passes_all_rules"] = df_results.all(axis=1)
return df_results
def apply_structural_alerts(mols: List[Chem.Mol], alert_type: str, n_jobs: int) -> pd.DataFrame:
"""Apply structural alert filters."""
print(f"\nApplying {alert_type} structural alerts...")
if alert_type == "common":
alert_filter = mc.structural.CommonAlertsFilters()
results = alert_filter(mols=mols, n_jobs=n_jobs, progress=True)
df_results = pd.DataFrame({
"has_common_alerts": [r["has_alerts"] for r in results],
"num_common_alerts": [r["num_alerts"] for r in results],
"common_alert_details": [", ".join(r["alert_details"]) if r["alert_details"] else "" for r in results]
})
elif alert_type == "nibr":
nibr_filter = mc.structural.NIBRFilters()
results = nibr_filter(mols=mols, n_jobs=n_jobs, progress=True)
df_results = pd.DataFrame({
"passes_nibr": results
})
elif alert_type == "lilly":
lilly_filter = mc.structural.LillyDemeritsFilters()
results = lilly_filter(mols=mols, n_jobs=n_jobs, progress=True)
df_results = pd.DataFrame({
"lilly_demerits": [r["demerits"] for r in results],
"passes_lilly": [r["passes"] for r in results],
"lilly_patterns": [", ".join([p["pattern"] for p in r["matched_patterns"]]) for r in results]
})
elif alert_type == "pains":
results = [mc.rules.basic_rules.pains_filter(mol) for mol in tqdm(mols, desc="PAINS")]
df_results = pd.DataFrame({
"passes_pains": results
})
else:
raise ValueError(f"Unknown alert type: {alert_type}")
return df_results
def apply_complexity_filter(mols: List[Chem.Mol], max_complexity: float, method: str = "bertz") -> pd.DataFrame:
"""Calculate molecular complexity."""
print(f"\nCalculating molecular complexity (method={method}, max={max_complexity})...")
complexity_scores = [
mc.complexity.calculate_complexity(mol, method=method)
for mol in tqdm(mols, desc="Complexity")
]
df_results = pd.DataFrame({
"complexity_score": complexity_scores,
"passes_complexity": [score <= max_complexity for score in complexity_scores]
})
return df_results
def apply_constraints(mols: List[Chem.Mol], constraints: Dict, n_jobs: int) -> pd.DataFrame:
"""Apply custom property constraints."""
print(f"\nApplying constraints: {constraints}")
constraint_filter = mc.constraints.Constraints(**constraints)
results = constraint_filter(mols=mols, n_jobs=n_jobs, progress=True)
df_results = pd.DataFrame({
"passes_constraints": [r["passes"] for r in results],
"constraint_violations": [", ".join(r["violations"]) if r["violations"] else "" for r in results]
})
return df_results
def apply_chemical_groups(mols: List[Chem.Mol], groups: List[str]) -> pd.DataFrame:
"""Detect chemical groups."""
print(f"\nDetecting chemical groups: {', '.join(groups)}")
group_detector = mc.groups.ChemicalGroup(groups=groups)
results = group_detector.get_all_matches(mols)
df_results = pd.DataFrame()
for group in groups:
df_results[f"has_{group}"] = [bool(r.get(group)) for r in results]
return df_results
def generate_summary(df: pd.DataFrame, output_file: Path):
"""Generate filtering summary report."""
summary_file = output_file.parent / f"{output_file.stem}_summary.txt"
with open(summary_file, "w") as f:
f.write("=" * 80 + "\n")
f.write("MEDCHEM FILTERING SUMMARY\n")
f.write("=" * 80 + "\n\n")
f.write(f"Total molecules processed: {len(df)}\n\n")
# Rule results
rule_cols = [col for col in df.columns if col.startswith("rule_") or col == "passes_all_rules"]
if rule_cols:
f.write("RULE FILTERS:\n")
f.write("-" * 40 + "\n")
for col in rule_cols:
if col in df.columns and df[col].dtype == bool:
n_pass = df[col].sum()
pct = 100 * n_pass / len(df)
f.write(f" {col}: {n_pass} passed ({pct:.1f}%)\n")
f.write("\n")
# Structural alerts
alert_cols = [col for col in df.columns if "alert" in col.lower() or "nibr" in col.lower() or "lilly" in col.lower() or "pains" in col.lower()]
if alert_cols:
f.write("STRUCTURAL ALERTS:\n")
f.write("-" * 40 + "\n")
if "has_common_alerts" in df.columns:
n_clean = (~df["has_common_alerts"]).sum()
pct = 100 * n_clean / len(df)
f.write(f" No common alerts: {n_clean} ({pct:.1f}%)\n")
if "passes_nibr" in df.columns:
n_pass = df["passes_nibr"].sum()
pct = 100 * n_pass / len(df)
f.write(f" Passes NIBR: {n_pass} ({pct:.1f}%)\n")
if "passes_lilly" in df.columns:
n_pass = df["passes_lilly"].sum()
pct = 100 * n_pass / len(df)
f.write(f" Passes Lilly: {n_pass} ({pct:.1f}%)\n")
avg_demerits = df["lilly_demerits"].mean()
f.write(f" Average Lilly demerits: {avg_demerits:.1f}\n")
if "passes_pains" in df.columns:
n_pass = df["passes_pains"].sum()
pct = 100 * n_pass / len(df)
f.write(f" Passes PAINS: {n_pass} ({pct:.1f}%)\n")
f.write("\n")
# Complexity
if "complexity_score" in df.columns:
f.write("COMPLEXITY:\n")
f.write("-" * 40 + "\n")
avg_complexity = df["complexity_score"].mean()
f.write(f" Average complexity: {avg_complexity:.1f}\n")
if "passes_complexity" in df.columns:
n_pass = df["passes_complexity"].sum()
pct = 100 * n_pass / len(df)
f.write(f" Within threshold: {n_pass} ({pct:.1f}%)\n")
f.write("\n")
# Constraints
if "passes_constraints" in df.columns:
f.write("CONSTRAINTS:\n")
f.write("-" * 40 + "\n")
n_pass = df["passes_constraints"].sum()
pct = 100 * n_pass / len(df)
f.write(f" Passes all constraints: {n_pass} ({pct:.1f}%)\n")
f.write("\n")
# Overall pass rate
pass_cols = [col for col in df.columns if col.startswith("passes_")]
if pass_cols:
df["passes_all_filters"] = df[pass_cols].all(axis=1)
n_pass = df["passes_all_filters"].sum()
pct = 100 * n_pass / len(df)
f.write("OVERALL:\n")
f.write("-" * 40 + "\n")
f.write(f" Molecules passing all filters: {n_pass} ({pct:.1f}%)\n")
f.write("\n" + "=" * 80 + "\n")
print(f"\nSummary report saved to: {summary_file}")
def main():
parser = argparse.ArgumentParser(
description="Batch molecular filtering using medchem",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__
)
# Input/Output
parser.add_argument("input", type=Path, help="Input file (CSV, TSV, SDF, or TXT)")
parser.add_argument("--output", "-o", type=Path, required=True, help="Output CSV file")
parser.add_argument("--smiles-column", default="smiles", help="Name of SMILES column (default: smiles)")
# Rule filters
parser.add_argument("--rules", help="Comma-separated list of rules (e.g., rule_of_five,rule_of_cns)")
# Structural alerts
parser.add_argument("--common-alerts", action="store_true", help="Apply common structural alerts")
parser.add_argument("--nibr", action="store_true", help="Apply NIBR filters")
parser.add_argument("--lilly", action="store_true", help="Apply Lilly demerits filter")
parser.add_argument("--pains", action="store_true", help="Apply PAINS filter")
# Complexity
parser.add_argument("--complexity", type=float, help="Maximum complexity threshold")
parser.add_argument("--complexity-method", default="bertz", choices=["bertz", "whitlock", "barone"],
help="Complexity calculation method")
# Constraints
parser.add_argument("--mw-range", help="Molecular weight range (e.g., 200,500)")
parser.add_argument("--logp-range", help="LogP range (e.g., -2,5)")
parser.add_argument("--tpsa-max", type=float, help="Maximum TPSA")
parser.add_argument("--hbd-max", type=int, help="Maximum H-bond donors")
parser.add_argument("--hba-max", type=int, help="Maximum H-bond acceptors")
parser.add_argument("--rotatable-bonds-max", type=int, help="Maximum rotatable bonds")
# Chemical groups
parser.add_argument("--groups", help="Comma-separated chemical groups to detect")
# Processing options
parser.add_argument("--n-jobs", type=int, default=-1, help="Number of parallel jobs (-1 = all cores)")
parser.add_argument("--no-summary", action="store_true", help="Don't generate summary report")
parser.add_argument("--filter-output", action="store_true", help="Only output molecules passing all filters")
args = parser.parse_args()
# Load molecules
df, mols = load_molecules(args.input, args.smiles_column)
# Apply filters
result_dfs = [df]
# Rules
if args.rules:
rule_list = [r.strip() for r in args.rules.split(",")]
df_rules = apply_rule_filters(mols, rule_list, args.n_jobs)
result_dfs.append(df_rules)
# Structural alerts
if args.common_alerts:
df_alerts = apply_structural_alerts(mols, "common", args.n_jobs)
result_dfs.append(df_alerts)
if args.nibr:
df_nibr = apply_structural_alerts(mols, "nibr", args.n_jobs)
result_dfs.append(df_nibr)
if args.lilly:
df_lilly = apply_structural_alerts(mols, "lilly", args.n_jobs)
result_dfs.append(df_lilly)
if args.pains:
df_pains = apply_structural_alerts(mols, "pains", args.n_jobs)
result_dfs.append(df_pains)
# Complexity
if args.complexity:
df_complexity = apply_complexity_filter(mols, args.complexity, args.complexity_method)
result_dfs.append(df_complexity)
# Constraints
constraints = {}
if args.mw_range:
mw_min, mw_max = map(float, args.mw_range.split(","))
constraints["mw_range"] = (mw_min, mw_max)
if args.logp_range:
logp_min, logp_max = map(float, args.logp_range.split(","))
constraints["logp_range"] = (logp_min, logp_max)
if args.tpsa_max:
constraints["tpsa_max"] = args.tpsa_max
if args.hbd_max:
constraints["hbd_max"] = args.hbd_max
if args.hba_max:
constraints["hba_max"] = args.hba_max
if args.rotatable_bonds_max:
constraints["rotatable_bonds_max"] = args.rotatable_bonds_max
if constraints:
df_constraints = apply_constraints(mols, constraints, args.n_jobs)
result_dfs.append(df_constraints)
# Chemical groups
if args.groups:
group_list = [g.strip() for g in args.groups.split(",")]
df_groups = apply_chemical_groups(mols, group_list)
result_dfs.append(df_groups)
# Combine results
df_final = pd.concat(result_dfs, axis=1)
# Filter output if requested
if args.filter_output:
pass_cols = [col for col in df_final.columns if col.startswith("passes_")]
if pass_cols:
df_final["passes_all"] = df_final[pass_cols].all(axis=1)
df_final = df_final[df_final["passes_all"]]
print(f"\nFiltered to {len(df_final)} molecules passing all filters")
# Save results
args.output.parent.mkdir(parents=True, exist_ok=True)
df_final.to_csv(args.output, index=False)
print(f"\nResults saved to: {args.output}")
# Generate summary
if not args.no_summary:
generate_summary(df_final, args.output)
print("\nDone!")
if __name__ == "__main__":
main()