Spaces:
Paused
Paused
| """ | |
| High-Throughput Virtual Screening & Batch Profiling Engine. | |
| Screens libraries of candidate SMILES across ESOL, BBBP, Tox21, and ADMET filters, | |
| and ranks compounds using a Multi-Objective Drug Desirability Score. | |
| """ | |
| from typing import List, Dict, Any, Optional | |
| import pandas as pd | |
| import numpy as np | |
| from rdkit import Chem, RDLogger | |
| from rdkit.Chem import Descriptors, QED | |
| from src.admet import compute_admet_profile | |
| from src.ensemble import predict_all_endpoints | |
| RDLogger.DisableLog('rdApp.*') | |
| def calculate_desirability_score( | |
| qed: float, | |
| lipinski_pass: bool, | |
| log_s: float, | |
| max_tox_prob: float, | |
| has_pains: bool, | |
| prefer_bbb: bool = False, | |
| bbb_prob: float = 0.5 | |
| ) -> float: | |
| """ | |
| Compute a composite Drug Desirability Score between 0.0 and 100.0. | |
| """ | |
| # 1. QED component (0 - 30 pts) | |
| score_qed = float(np.clip(qed, 0.0, 1.0)) * 30.0 | |
| # 2. Lipinski compliance (0 or 20 pts) | |
| score_lipinski = 20.0 if lipinski_pass else 0.0 | |
| # 3. Solubility (0 - 20 pts) - optimal between -1 and -3 | |
| if log_s >= -2.0: | |
| score_sol = 20.0 | |
| elif log_s >= -4.0: | |
| score_sol = 15.0 | |
| elif log_s >= -6.0: | |
| score_sol = 8.0 | |
| else: | |
| score_sol = 2.0 | |
| # 4. Toxicity Safety (0 - 20 pts) | |
| score_tox = max(0.0, (1.0 - max_tox_prob) * 20.0) | |
| # 5. Assay Interference penalty (-10 pts if PAINS alert) | |
| penalty_pains = -10.0 if has_pains else 0.0 | |
| # 6. BBB adjustment (0 - 10 pts) | |
| if prefer_bbb: | |
| score_bbb = bbb_prob * 10.0 | |
| else: | |
| # Non-CNS drug preference (lower BBB is better) | |
| score_bbb = (1.0 - bbb_prob) * 10.0 | |
| total = max(0.0, min(100.0, score_qed + score_lipinski + score_sol + score_tox + score_bbb + penalty_pains)) | |
| return round(total, 1) | |
| def screen_batch_molecules( | |
| smiles_list: List[str], | |
| compound_names: Optional[List[str]] = None, | |
| prefer_bbb: bool = False, | |
| progress_callback: Optional[Any] = None | |
| ) -> pd.DataFrame: | |
| """ | |
| Screen a collection of SMILES strings and return a ranked DataFrame. | |
| """ | |
| records = [] | |
| total = len(smiles_list) | |
| for idx, smi in enumerate(smiles_list): | |
| smi = str(smi).strip() | |
| if not smi: | |
| continue | |
| name = compound_names[idx] if compound_names and idx < len(compound_names) else f"Compound_{idx+1}" | |
| mol = Chem.MolFromSmiles(smi) | |
| if mol is None: | |
| continue | |
| try: | |
| admet = compute_admet_profile(mol) | |
| preds = predict_all_endpoints(smi) | |
| esol_val = float(preds['esol']['mean'][0]) | |
| bbbp_prob = float(preds['bbbp']['mean'][0]) | |
| tox_probs = preds['tox21']['mean'] | |
| max_tox = float(np.max(tox_probs)) | |
| mean_tox = float(np.mean(tox_probs)) | |
| high_risk_tox_count = int(np.sum(tox_probs >= 0.5)) | |
| qed_score = admet['descriptors']['QED'] | |
| lip_pass = admet['filters']['Lipinski']['passed'] | |
| pains_alert = admet['alerts']['PAINS']['has_alerts'] | |
| desirability = calculate_desirability_score( | |
| qed=qed_score, | |
| lipinski_pass=lip_pass, | |
| log_s=esol_val, | |
| max_tox_prob=max_tox, | |
| has_pains=pains_alert, | |
| prefer_bbb=prefer_bbb, | |
| bbb_prob=bbbp_prob | |
| ) | |
| records.append({ | |
| "Compound ID": name, | |
| "SMILES": smi, | |
| "Desirability Score": desirability, | |
| "QED": round(qed_score, 3), | |
| "Solubility logS": round(esol_val, 2), | |
| "BBB Probability": f"{bbbp_prob*100:.1f}%", | |
| "Max Tox Risk": f"{max_tox*100:.1f}%", | |
| "High Risk Assays": high_risk_tox_count, | |
| "MW (Da)": round(admet['descriptors']['Molecular Weight'], 1), | |
| "LogP": round(admet['descriptors']['LogP'], 2), | |
| "TPSA (Ų)": round(admet['descriptors']['TPSA'], 1), | |
| "Lipinski": "PASS" if lip_pass else f"FAIL ({admet['filters']['Lipinski']['num_violations']})", | |
| "PAINS Alert": "YES" if pains_alert else "Clean" | |
| }) | |
| except Exception: | |
| continue | |
| if progress_callback: | |
| progress_callback((idx + 1) / total) | |
| df = pd.DataFrame(records) | |
| if not df.empty: | |
| df = df.sort_values(by="Desirability Score", ascending=False).reset_index(drop=True) | |
| return df | |