Medchem

scripts/filter_molecules.py

#!/usr/bin/env python3
"""
Batch molecular filtering using the medchem library.

Usage:
    uv run python filter_molecules.py input.csv --rules rule_of_five,rule_of_cns --pains --output filtered.csv
    uv run python filter_molecules.py input.sdf --rules rule_of_veber --nibr --complexity 99 --output results.csv
    uv run python filter_molecules.py smiles.txt --query 'MATCHRULE("rule_of_five") AND NOT HASALERT("pains")' --output clean.csv
"""

import argparse
import sys
from pathlib import Path
from typing import List, Tuple

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: uv 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 CSV/TSV, SDF, or plain SMILES text files."""
    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]
        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")
            print(f"Available columns: {', '.join(df.columns)}")
            sys.exit(1)
        print("Converting SMILES to molecules...")
        mols = [
            dm.to_mol(smi) if isinstance(smi, str) and smi.strip() else None
            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("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)

    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. Returns rule result columns."""
    print(f"\nApplying rule filters: {', '.join(rules)}")
    rfilter = mc.rules.RuleFilters(rule_list=rules)
    results = rfilter(mols=mols, n_jobs=n_jobs, progress=True)
    return results.drop(columns=["mol"], errors="ignore")


def apply_common_alerts(mols: List[Chem.Mol], n_jobs: int) -> pd.DataFrame:
    """Apply ChEMBL-derived common structural alerts."""
    print("\nApplying common structural alerts...")
    alert_filter = mc.structural.CommonAlertsFilters(alerts_set=["BMS"])
    results = alert_filter(mols=mols, n_jobs=n_jobs, progress=True)
    return results.drop(columns=["mol"], errors="ignore").rename(
        columns={"pass_filter": "passes_common_alerts", "status": "common_alert_status",
                 "reasons": "common_alert_reasons"}
    )


def apply_nibr(mols: List[Chem.Mol], n_jobs: int) -> pd.DataFrame:
    """Apply NIBR screening-deck filters."""
    print("\nApplying NIBR filters...")
    nibr_filter = mc.structural.NIBRFilters()
    results = nibr_filter(mols=mols, n_jobs=n_jobs, progress=True)
    # The class excludes explicit exclusion alerts; enforce the accumulated
    # severity cutoff as well, matching functional.nibr_filter's default.
    results["pass_filter"] &= results["severity"] < 10
    return results.drop(columns=["mol"], errors="ignore").rename(
        columns={"pass_filter": "passes_nibr", "status": "nibr_status",
                 "reasons": "nibr_reasons"}
    )


def apply_alert_catalog(
    mols: List[Chem.Mol], alerts: List[str], n_jobs: int
):
    """Apply named alert catalogs (pains, brenk, etc.)."""
    print(f"\nApplying alert catalogs: {', '.join(alerts)}")
    unknown = set(alerts) - set(mc.catalogs.list_named_catalogs())
    if unknown:
        raise ValueError(f"Unknown named catalogs: {', '.join(sorted(unknown))}")
    for alert in alerts:
        passes = mc.functional.catalog_filter(
            mols=mols, catalogs=[alert], n_jobs=n_jobs, progress=True
        )
        yield alert, passes


def apply_lilly(mols: List[Chem.Mol], max_demerits: int, n_jobs: int) -> pd.DataFrame:
    """Apply Lilly demerits; a requested but unavailable filter is an error."""
    print(f"\nApplying Lilly demerits filter (max={max_demerits})...")
    try:
        passes = mc.functional.lilly_demerit_filter(
            mols=mols, max_demerits=max_demerits, n_jobs=n_jobs, progress=True
        )
    except ImportError as e:
        raise RuntimeError(
            "Lilly filter unavailable. Install native tools with: medchem install-lilly"
        ) from e
    return pd.DataFrame({"passes_lilly": passes})


def apply_complexity(
    mols: List[Chem.Mol], limit: str, method: str, n_jobs: int
) -> pd.DataFrame:
    """Filter by complexity percentile threshold."""
    print(f"\nApplying complexity filter (metric={method}, limit={limit} percentile)...")
    passes = mc.functional.complexity_filter(
        mols=mols,
        complexity_metric=method,
        limit=limit,
        n_jobs=n_jobs,
        progress=True,
    )
    return pd.DataFrame({f"passes_complexity_{method}": passes})


def apply_query(mols: List[Chem.Mol], query: str, n_jobs: int) -> pd.DataFrame:
    """Apply medchem query language filter."""
    print(f"\nApplying query: {query}")
    qf = mc.query.QueryFilter(query)
    passes = qf(mols=mols, n_jobs=n_jobs, progress=True)
    return pd.DataFrame({"passes_query": passes})


def apply_groups(mols: List[Chem.Mol], groups: List[str]) -> pd.DataFrame:
    """Detect chemical group matches."""
    print(f"\nDetecting chemical groups: {', '.join(groups)}")
    unknown = set(groups) - set(mc.groups.list_default_chemical_groups())
    if unknown:
        raise ValueError(f"Unknown chemical groups: {', '.join(sorted(unknown))}")
    results = {}
    for group in groups:
        detector = mc.groups.ChemicalGroup(groups=[group])
        results[f"has_{group}"] = [detector.has_match(mol) for mol in mols]
    return pd.DataFrame(results)


def generate_summary(
    df: pd.DataFrame, output_file: Path, pass_columns: list[str] | None = None
) -> None:
    """Write a text summary of filtering results."""
    summary_file = output_file.parent / f"{output_file.stem}_summary.txt"
    pass_cols = pass_columns if pass_columns is not None else [
        c for c in df.columns if c.startswith("passes_") or c == "pass_all"
    ]

    with open(summary_file, "w") as f:
        f.write("=" * 80 + "\nMEDCHEM FILTERING SUMMARY\n" + "=" * 80 + "\n\n")
        f.write(f"Total molecules processed: {len(df)}\n\n")

        for col in pass_cols:
            if col in df.columns and df[col].dtype == bool:
                n_pass = df[col].sum()
                pct = 100 * n_pass / len(df) if len(df) else 0
                f.write(f"  {col}: {n_pass} passed ({pct:.1f}%)\n")

        if pass_cols:
            bool_cols = [c for c in pass_cols if c in df.columns and df[c].dtype == bool]
            if bool_cols:
                n_all = df[bool_cols].all(axis=1).sum()
                pct = 100 * n_all / len(df) if len(df) else 0
                f.write(f"\n  All filters passed: {n_all} ({pct:.1f}%)\n")

        f.write("\n" + "=" * 80 + "\n")

    print(f"\nSummary report saved to: {summary_file}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Batch molecular filtering using medchem",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    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="SMILES column name (default: smiles)")

    parser.add_argument("--rules", help="Comma-separated rules (e.g. rule_of_five,rule_of_cns)")
    parser.add_argument("--query", help='Medchem query string (e.g. MATCHRULE("rule_of_five") AND NOT HASALERT("pains"))')
    parser.add_argument("--common-alerts", action="store_true", help="Apply BMS 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 (requires medchem install-lilly)")
    parser.add_argument("--lilly-max", type=int, default=160, help="Max Lilly demerits (default: 160)")
    parser.add_argument(
        "--alerts",
        help="Comma-separated alert catalogs (e.g. pains,brenk). Shorthand for --pains when set to pains",
    )
    parser.add_argument("--pains", action="store_true", help="Apply PAINS filter (alias for --alerts pains)")

    parser.add_argument(
        "--complexity",
        help="Complexity percentile limit (e.g. 99 keeps below 99th percentile on ZINC-15)",
    )
    parser.add_argument(
        "--complexity-method",
        default="bertz",
        choices=["bertz", "sas", "qed", "clogp", "whitlock", "barone", "smcm", "twc"],
        help="Complexity metric (default: bertz)",
    )
    parser.add_argument("--groups", help="Comma-separated chemical groups to detect")
    parser.add_argument("--n-jobs", type=int, default=-1, help="Parallel jobs (-1 = all cores)")
    parser.add_argument("--no-summary", action="store_true", help="Skip summary report")
    parser.add_argument("--filter-output", action="store_true", help="Only output molecules passing all filters")

    args = parser.parse_args()

    if not any([args.rules, args.query, args.common_alerts, args.nibr, args.lilly,
                args.pains, args.alerts, args.complexity, args.groups]):
        print("Error: Specify at least one filter (--rules, --query, --pains, --nibr, etc.)")
        sys.exit(1)

    df, mols = load_molecules(args.input, args.smiles_column)
    if not mols:
        parser.error("No valid molecules remain after parsing; no filters were evaluated")
    result_parts = [df.reset_index(drop=True)]

    if args.rules:
        rule_list = [r.strip() for r in args.rules.split(",")]
        unknown = set(rule_list) - set(mc.rules.RuleFilters.list_available_rules_names())
        if unknown:
            print(f"Warning: Unknown rules (will error at runtime): {', '.join(sorted(unknown))}")
        result_parts.append(apply_rule_filters(mols, rule_list, args.n_jobs))

    if args.query:
        result_parts.append(apply_query(mols, args.query, args.n_jobs))

    if args.common_alerts:
        result_parts.append(apply_common_alerts(mols, args.n_jobs))

    if args.nibr:
        result_parts.append(apply_nibr(mols, args.n_jobs))

    if args.lilly:
        result_parts.append(apply_lilly(mols, args.lilly_max, args.n_jobs))

    alert_list = []
    if args.pains:
        alert_list.append("pains")
    if args.alerts:
        alert_list.extend(a.strip() for a in args.alerts.split(","))
    for alert in dict.fromkeys(alert_list):
        col_name = f"passes_{alert}"
        _, passes = next(apply_alert_catalog(mols, [alert], args.n_jobs))
        result_parts.append(pd.DataFrame({col_name: passes}))

    if args.complexity:
        result_parts.append(
            apply_complexity(mols, args.complexity, args.complexity_method, args.n_jobs)
        )

    if args.groups:
        group_list = [g.strip() for g in args.groups.split(",")]
        result_parts.append(apply_groups(mols, group_list))

    df_final = pd.concat(result_parts, axis=1)
    # Input metadata may also use passes_* names. Only freshly evaluated
    # filters may affect retention or the scientific summary.
    pass_cols = [
        column for part in result_parts[1:] for column in part.columns
        if column.startswith("passes_") or column == "pass_all"
    ]
    duplicate_columns = df_final.columns[df_final.columns.duplicated()].tolist()
    if duplicate_columns:
        parser.error(f"Input and result columns collide: {duplicate_columns}; rename input columns")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    if not args.no_summary:
        generate_summary(df_final, args.output, pass_columns=pass_cols)

    if args.filter_output:
        bool_cols = [c for c in pass_cols if c in df_final.columns and df_final[c].dtype == bool]
        if bool_cols:
            mask = df_final[bool_cols].all(axis=1)
            df_final = df_final[mask]
            print(f"\nFiltered to {len(df_final)} molecules passing all filters")

    df_final.to_csv(args.output, index=False)
    print(f"\nResults saved to: {args.output}")

    print("\nDone!")


if __name__ == "__main__":
    main()