""" PROTAC classifiers and guidance scorers for the DeFoG guidance loop. All three classes accept DeFoG's molecule_list format and share _graph_to_mol. They are called together on the same batch in each guidance loop iteration: PROTACClassifier → 0/1 hard validity mask EGFRPROPKScorer → continuous PK guidance (T½, CL, Vss) PROTACRDKitScorer → continuous property guidance (QED, logP, TPSA, Fsp3, HBD, rotbonds) DeFoG molecule_list format: list of [atom_types, edge_types] atom_types : 1D int tensor, length n — index into atom_decoder edge_types : 2D int tensor [n, n] — 0=none 1=single 2=double 3=triple 4=aromatic ≥5=virtual """ from __future__ import annotations import os import pickle from typing import Optional import numpy as np import torch from rdkit import Chem, RDLogger from rdkit.Chem import rdchem from protac_analyzer import PROTACAnalyzer RDLogger.DisableLog("rdApp.*") # --------------------------------------------------------------------------- # Shared graph → Mol conversion # --------------------------------------------------------------------------- _BOND_TYPES = [ None, rdchem.BondType.SINGLE, rdchem.BondType.DOUBLE, rdchem.BondType.TRIPLE, rdchem.BondType.AROMATIC, ] def _graph_to_mol(atom_types: torch.Tensor, edge_types: torch.Tensor, atom_decoder: list[str]) -> Chem.RWMol: """Convert DeFoG (atom_types, edge_types) tensors to a Chem.RWMol.""" mol = Chem.RWMol() for idx in atom_types: mol.AddAtom(Chem.Atom(atom_decoder[idx.item()])) edge_types = torch.triu(edge_types) edge_types[edge_types >= 5] = 0 for i, j in torch.nonzero(edge_types).tolist(): if i != j: mol.AddBond(i, j, _BOND_TYPES[edge_types[i, j].item()]) return mol def prebuild_mols(molecule_list: list, atom_decoder: list[str]) -> list: """Convert + sanitize each (atom_types, edge_types) once for shared scorer use. Returns list of sanitized RWMol or None if conversion/sanitization failed.""" mols = [] for atom_types, edge_types in molecule_list: try: m = _graph_to_mol(atom_types, edge_types, atom_decoder) Chem.SanitizeMol(m) mols.append(m) except Exception: mols.append(None) return mols # --------------------------------------------------------------------------- # Hard validity classifier # --------------------------------------------------------------------------- class PROTACClassifier: """ Hard validity scorer for PROTAC linker graphs produced by DeFoG. Returns 0/1 tensor — 1 only when: - PROTACAnalyzer structural checks pass (atom count, ring closure, MW, SA, repr.) - builder.assemble() succeeds (actual RDKit reaction with warhead + E3) Using assembly as the attachment check means guidance is steered by real chemistry — NH2, reactive NH, or any N that actually reacts — rather than a SMARTS proxy. Parameters ---------- analyzer : PROTACAnalyzer Structural checks only (attachment_sites check is bypassed here). builder : PROTAC_Builder Used for the assembly-based attachment check. atom_decoder : list[str] Maps atom-type index → element symbol. """ def __init__(self, analyzer: PROTACAnalyzer, builder, atom_decoder: list[str]): self.analyzer = analyzer self.builder = builder self.atom_decoder = atom_decoder # SMARTS for site counting. Requiring an alpha carbon on the carbonyl # rules out carbamic acids (N-C(=O)OH), which can't form stable amides. self._cooh_pat = Chem.MolFromSmarts("[#6][CX3](=[OX1])[OX2H1]") self._cho_pat = Chem.MolFromSmarts("[#6][CX3H1](=[OX1])") self._methyl_pat = Chem.MolFromSmarts("[CX4H3]") self._amine_pat = Chem.MolFromSmarts("[NX3;!H0;!$(NC=O);!$(NS=O);!$(N=O)]") self._nh2_pat = Chem.MolFromSmarts("[NX3H2;D1;!$(NC=O);!$(NS=O);!$(N=O)]") # Hard rejection patterns. self._alerts = [ Chem.MolFromSmarts('[NX3][NX3]'), # hydrazide Chem.MolFromSmarts('[#7]=[#7]'), # azo Chem.MolFromSmarts('[N+][O-]'), # N-oxide Chem.MolFromSmarts('[#7]-[OX2]'), # N-O single bond Chem.MolFromSmarts('[NX3]C(=[OX1])[OX2H1]'), # carbamate ] def get_scores(self, molecule_list: list, fix_termini=False, mols=None) -> torch.Tensor: """Binary validity (0/1) over a list of molecules. Returns 1 if the mol passes (a) sanitize + SMILES round-trip, (b) no liability SMARTS, and (c) >= 2 reactive termini for the chosen fix_termini mode (cooh / methyl / amine / amino_acid). Returns 0 otherwise. Pass `mols` to skip graph_to_mol + sanitize. """ methyl_mode = (fix_termini == 'methyl') amine_mode = (fix_termini == 'amine') amino_acid_mode = (fix_termini == 'amino_acid') if mols is None: mols = prebuild_mols(molecule_list, self.atom_decoder) scores = [] for mol in mols: if mol is None: scores.append(0.0) continue try: # sanitize + SMILES round-trip try: Chem.SanitizeMol(mol) smi = Chem.MolToSmiles(mol) if not smi or Chem.MolFromSmiles(smi) is None: scores.append(0.0) continue except Exception: scores.append(0.0) continue # liability filter if any(mol.HasSubstructMatch(p) for p in self._alerts): scores.append(0.0) continue # reactive termini count if amine_mode: n_sites = len(mol.GetSubstructMatches(self._nh2_pat)) elif amino_acid_mode: n_sites = (len(mol.GetSubstructMatches(self._nh2_pat)) + len(mol.GetSubstructMatches(self._cooh_pat))) else: n_sites = (len(mol.GetSubstructMatches(self._cooh_pat)) + len(mol.GetSubstructMatches(self._cho_pat))) if methyl_mode: n_sites += len(mol.GetSubstructMatches(self._methyl_pat)) scores.append(1.0 if n_sites >= 2 else 0.0) except Exception: scores.append(0.0) return torch.tensor(scores, dtype=torch.float32) def __call__(self, molecule_list: list, mols=None) -> torch.Tensor: return self.get_scores(molecule_list, fix_termini=getattr(self, 'fix_termini', False), mols=mols) # --------------------------------------------------------------------------- # EGFR-PROPK soft PK guidance scorer # --------------------------------------------------------------------------- class EGFRPROPKScorer: """ Soft PK guidance using fine-tuned CatBoost models from Zhang et al. 2026 (EGFR-PROPK). Predicts T½, CL, Vss on the assembled PROTAC. score = sigmoid(t½_raw) − sigmoid(CL_raw) GIN features (300 dims) are zero-padded — DGL not available in this env. The remaining 2563 structural features still provide useful guidance. Parameters ---------- builder : PROTAC_Builder Used to assemble linker + warhead + E3 into a full PROTAC. model_dir : str Directory containing t0.5_10_finetune.pkl, CL_1_finetune.pkl, Vss_1_finetune.pkl. atom_decoder : list[str] Maps atom-type index → element symbol. half_life_weight, clearance_weight : float Relative weights for the composite score. """ # Feature vector layout matching EGFR-PROPK training (maplight_gnn.py) _N_MORGAN = 1024 _N_AVALON = 1024 _N_ERG = 315 _N_GIN = 300 # zero-padded — DGL unavailable _RDKIT_DESC_NAMES = [ "BalabanJ", "BertzCT", "Chi0", "Chi0n", "Chi0v", "Chi1", "Chi1n", "Chi1v", "Chi2n", "Chi2v", "Chi3n", "Chi3v", "Chi4n", "Chi4v", "EState_VSA1", "EState_VSA10", "EState_VSA11", "EState_VSA2", "EState_VSA3", "EState_VSA4", "EState_VSA5", "EState_VSA6", "EState_VSA7", "EState_VSA8", "EState_VSA9", "ExactMolWt", "FpDensityMorgan1", "FpDensityMorgan2", "FpDensityMorgan3", "FractionCSP3", "HallKierAlpha", "HeavyAtomCount", "HeavyAtomMolWt", "Ipc", "Kappa1", "Kappa2", "Kappa3", "LabuteASA", "MaxAbsEStateIndex", "MaxAbsPartialCharge", "MaxEStateIndex", "MaxPartialCharge", "MinAbsEStateIndex", "MinAbsPartialCharge", "MinEStateIndex", "MinPartialCharge", "MolLogP", "MolMR", "MolWt", "NHOHCount", "NOCount", "NumAliphaticCarbocycles", "NumAliphaticHeterocycles", "NumAliphaticRings", "NumAromaticCarbocycles", "NumAromaticHeterocycles", "NumAromaticRings", "NumHAcceptors", "NumHDonors", "NumHeteroatoms", "NumRadicalElectrons", "NumRotatableBonds", "NumSaturatedCarbocycles", "NumSaturatedHeterocycles", "NumSaturatedRings", "NumValenceElectrons", "PEOE_VSA1", "PEOE_VSA10", "PEOE_VSA11", "PEOE_VSA12", "PEOE_VSA13", "PEOE_VSA14", "PEOE_VSA2", "PEOE_VSA3", "PEOE_VSA4", "PEOE_VSA5", "PEOE_VSA6", "PEOE_VSA7", "PEOE_VSA8", "PEOE_VSA9", "RingCount", "SMR_VSA1", "SMR_VSA10", "SMR_VSA2", "SMR_VSA3", "SMR_VSA4", "SMR_VSA5", "SMR_VSA6", "SMR_VSA7", "SMR_VSA8", "SMR_VSA9", "SlogP_VSA1", "SlogP_VSA10", "SlogP_VSA11", "SlogP_VSA12", "SlogP_VSA2", "SlogP_VSA3", "SlogP_VSA4", "SlogP_VSA5", "SlogP_VSA6", "SlogP_VSA7", "SlogP_VSA8", "SlogP_VSA9", "TPSA", "VSA_EState1", "VSA_EState10", "VSA_EState2", "VSA_EState3", "VSA_EState4", "VSA_EState5", "VSA_EState6", "VSA_EState7", "VSA_EState8", "VSA_EState9", "fr_Al_COO", "fr_Al_OH", "fr_Al_OH_noTert", "fr_ArN", "fr_Ar_COO", "fr_Ar_N", "fr_Ar_NH", "fr_Ar_OH", "fr_COO", "fr_COO2", "fr_C_O", "fr_C_O_noCOO", "fr_C_S", "fr_HOCCN", "fr_Imine", "fr_NH0", "fr_NH1", "fr_NH2", "fr_N_O", "fr_Ndealkylation1", "fr_Ndealkylation2", "fr_Nhpyrrole", "fr_SH", "fr_aldehyde", "fr_alkyl_carbamate", "fr_alkyl_halide", "fr_allylic_oxid", "fr_amide", "fr_amidine", "fr_aniline", "fr_aryl_methyl", "fr_azide", "fr_azo", "fr_barbitur", "fr_benzene", "fr_benzodiazepine", "fr_bicyclic", "fr_diazo", "fr_dihydropyridine", "fr_epoxide", "fr_ester", "fr_ether", "fr_furan", "fr_guanido", "fr_halogen", "fr_hdrzine", "fr_hdrzone", "fr_imidazole", "fr_imide", "fr_isocyan", "fr_isothiocyan", "fr_ketone", "fr_ketone_Topliss", "fr_lactam", "fr_lactone", "fr_methoxy", "fr_morpholine", "fr_nitrile", "fr_nitro", "fr_nitro_arom", "fr_nitro_arom_nonortho", "fr_nitroso", "fr_oxazole", "fr_oxime", "fr_para_hydroxylation", "fr_phenol", "fr_phenol_noOrthoHbond", "fr_phos_acid", "fr_phos_ester", "fr_piperdine", "fr_piperzine", "fr_priamide", "fr_prisulfonamd", "fr_pyridine", "fr_quatN", "fr_sulfide", "fr_sulfonamd", "fr_sulfone", "fr_term_acetylene", "fr_tetrazole", "fr_thiazole", "fr_thiocyan", "fr_thiophene", "fr_unbrch_alkane", "fr_urea", "qed", ] def __init__(self, builder, model_dir: str, atom_decoder: list[str], half_life_weight: float = 1.0, clearance_weight: float = 1.0): from rdkit.Chem import DataStructs from rdkit.Chem.rdMolDescriptors import GetHashedMorganFingerprint from rdkit.Avalon.pyAvalonTools import GetAvalonCountFP from rdkit.Chem import rdReducedGraphs from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator self.builder = builder self.atom_decoder = atom_decoder self.half_life_weight = half_life_weight self.clearance_weight = clearance_weight self._DataStructs = DataStructs self._GetHashedMorganFP = GetHashedMorganFingerprint self._GetAvalonFP = GetAvalonCountFP self._GetErGFP = rdReducedGraphs.GetErGFingerprint self._desc_calc = MolecularDescriptorCalculator(self._RDKIT_DESC_NAMES) def _load(name): with open(os.path.join(model_dir, name), "rb") as f: return pickle.load(f) self._model_t_half = _load("t0.5_10_finetune.pkl") self._model_cl = _load("CL_1_finetune.pkl") self._model_vss = _load("Vss_1_finetune.pkl") def _assemble(self, linker_mol: Chem.Mol) -> Optional[Chem.Mol]: amide_only = getattr(self, 'fix_termini', None) != 'methyl' smi = Chem.MolToSmiles(linker_mol) if smi == getattr(self, '_asm_smi', None): return self._asm_mol try: result = self.builder.assemble(linker_mol, return_all=False, amide_only=amide_only) smiles = result if isinstance(result, str) else result.get("protac_smiles") mol = Chem.MolFromSmiles(smiles) if smiles else None except Exception: mol = None self._asm_smi, self._asm_mol = smi, mol return mol def _featurize(self, mol: Chem.Mol) -> np.ndarray: def _fp_arr(fp): a = np.zeros((0,), dtype=np.int8) self._DataStructs.ConvertToNumpyArray(fp, a) return a vec = np.concatenate([ _fp_arr(self._GetHashedMorganFP(mol, nBits=self._N_MORGAN, radius=2)), _fp_arr(self._GetAvalonFP(mol, nBits=self._N_AVALON)), np.array(self._GetErGFP(mol)), np.array(self._desc_calc.CalcDescriptors(mol)), np.zeros(self._N_GIN, dtype=np.float32), ]) return np.nan_to_num(vec, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32) def _score_protac(self, protac_mol: Chem.Mol) -> float: X = self._featurize(protac_mol).reshape(1, -1) t_half = float(self._model_t_half.predict(X)[0]) cl = float(self._model_cl.predict(X)[0]) sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x)) return self.half_life_weight * sigmoid(t_half) - self.clearance_weight * sigmoid(cl) def predict_pk(self, protac_mol: Chem.Mol) -> dict: """Return raw EGFR-PROPK predictions for the assembled PROTAC. Keys: t_half (hours, log-scale), cl (clearance), vss (volume of distribution), plus the composite score (sigmoid(t_half) − sigmoid(cl)). Use this for post-hoc validation / reporting; use _score_protac for guidance scalars. """ X = self._featurize(protac_mol).reshape(1, -1) t_half = float(self._model_t_half.predict(X)[0]) cl = float(self._model_cl.predict(X)[0]) vss = float(self._model_vss.predict(X)[0]) sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x)) return { "t_half": t_half, "cl": cl, "vss": vss, "score": self.half_life_weight * sigmoid(t_half) - self.clearance_weight * sigmoid(cl), } def get_scores(self, molecule_list: list) -> torch.Tensor: scores = [] for atom_types, edge_types in molecule_list: try: linker = _graph_to_mol(atom_types, edge_types, self.atom_decoder) Chem.SanitizeMol(linker) protac = self._assemble(linker) scores.append(self._score_protac(protac) if protac is not None else 0.0) except Exception: scores.append(0.0) return torch.tensor(scores, dtype=torch.float32) def __call__(self, molecule_list: list) -> torch.Tensor: return self.get_scores(molecule_list) # --------------------------------------------------------------------------- # RDKit property soft guidance scorer # --------------------------------------------------------------------------- class PROTACRDKitScorer: """ Soft property guidance using RDKit descriptors on the assembled PROTAC. Returns weighted mean of per-property scores, each in [0, 1]. All windows are derived from PROTAC-DB (n=6,110 unique assembled PROTACs): flat top = [mean − 1σ, mean + 1σ] (score = 1) outer = [mean − 2σ, mean + 2σ] (score ramps 0 → 1) Parameters ---------- builder : PROTAC_Builder Used to assemble linker + warhead + E3 into a full PROTAC. atom_decoder : list[str] Maps atom-type index → element symbol. weights : dict[str, float] | None Per-property weight. Keys: qed, logp, tpsa, fsp3, hbd, rotbonds. active : set[str] | None Subset of properties to score. Defaults to all. """ # Trapezoidal: (lo, lo_ramp, hi_ramp, hi) # Score = 0 outside [lo, hi], linearly ramps to 1 in shoulders, flat 1 in [lo_ramp, hi_ramp]. _PROPERTY_WINDOWS = { "qed": (0.001, 0.015, 0.145, 0.210), "logp": (0.7, 3.1, 7.9, 10.3), "tpsa": (119.0, 165.0, 259.0, 306.0), "fsp3": (0.24, 0.33, 0.51, 0.60), "hbd": (0.7, 2.4, 6.0, 7.8), "rotbonds": (5.0, 12.0, 25.0, 32.0), } def __init__(self, builder, atom_decoder: list[str], weights: dict | None = None, active: set | None = None): from rdkit.Chem import Descriptors, QED, rdMolDescriptors self.builder = builder self.atom_decoder = atom_decoder self.active = active if active is not None else set(self._PROPERTY_WINDOWS) self.weights = {k: (weights or {}).get(k, 1.0) for k in self.active} self._prop_fns = { "qed": QED.qed, "logp": Descriptors.MolLogP, "tpsa": Descriptors.TPSA, "fsp3": rdMolDescriptors.CalcFractionCSP3, "hbd": lambda m: float(Descriptors.NumHDonors(m)), "rotbonds": lambda m: float(Descriptors.NumRotatableBonds(m)), } @staticmethod def _trapezoid(v, lo, lo_ramp, hi_ramp, hi) -> float: if v <= lo or v >= hi: return 0.0 if v <= lo_ramp: return (v - lo) / (lo_ramp - lo) if v >= hi_ramp: return (hi - v) / (hi - hi_ramp) return 1.0 def _score_property(self, name: str, mol: Chem.Mol) -> float: return self._trapezoid(self._prop_fns[name](mol), *self._PROPERTY_WINDOWS[name]) def _assemble(self, linker_mol: Chem.Mol) -> Optional[Chem.Mol]: amide_only = getattr(self, 'fix_termini', None) != 'methyl' smi = Chem.MolToSmiles(linker_mol) if smi == getattr(self, '_asm_smi', None): return self._asm_mol try: result = self.builder.assemble(linker_mol, return_all=False, amide_only=amide_only) smiles = result if isinstance(result, str) else result.get("protac_smiles") mol = Chem.MolFromSmiles(smiles) if smiles else None except Exception: mol = None self._asm_smi, self._asm_mol = smi, mol return mol def _score_protac(self, protac_mol: Chem.Mol) -> float: per_prop = {k: self._score_property(k, protac_mol) for k in self.active} total_w = sum(self.weights[k] for k in per_prop) return sum(self.weights[k] * v for k, v in per_prop.items()) / total_w if total_w else 0.0 def get_scores(self, molecule_list: list) -> torch.Tensor: scores = [] for atom_types, edge_types in molecule_list: try: linker = _graph_to_mol(atom_types, edge_types, self.atom_decoder) Chem.SanitizeMol(linker) protac = self._assemble(linker) scores.append(self._score_protac(protac) if protac is not None else 0.0) except Exception: scores.append(0.0) return torch.tensor(scores, dtype=torch.float32) def __call__(self, molecule_list: list) -> torch.Tensor: return self.get_scores(molecule_list) # --------------------------------------------------------------------------- # Connectivity soft guidance scorer # --------------------------------------------------------------------------- class PROTACLinkerConnectivityScorer: """ Soft signal pushing the linker toward a single connected molecule with both anchors in the same fragment. Returns score in [0, 1]: 1.0 — connected (single fragment) frag_frac / num_fragments — disconnected with anchors in same fragment frag_frac / num_fragments / 2 — disconnected with anchors in different fragments 0.0 — sanitization fails Scaling by 1/num_fragments penalizes "almost-connected" (95%-largest-fragment + 1 stray atom) much more harshly than the older bare frag_frac, which previously rewarded such candidates at ~0.95 and gave guidance no incentive to merge the stray. With this scaling, 2-fragment cases peak at 0.5 even when the largest fragment is 99% of atoms; only true single-fragment graphs reach 1.0. Anchor halving stacks on top to favour bridging the warhead/E3 anchors specifically. """ def __init__(self, atom_decoder: list[str]): self.atom_decoder = atom_decoder def get_scores(self, molecule_list: list, mols=None) -> torch.Tensor: anchor_nodes = getattr(self, 'anchor_nodes', None) if mols is None: mols = prebuild_mols(molecule_list, self.atom_decoder) scores = [] for mol in mols: if mol is None: scores.append(0.0) continue try: frags = Chem.GetMolFrags(mol) num_frags = len(frags) if num_frags == 1: scores.append(1.0) continue frag_frac = max(len(f) for f in frags) / mol.GetNumAtoms() # Harsher: scale by 1/num_fragments so any extra fragment caps the # score below 0.5, regardless of how small the stray fragment is. score = frag_frac / num_frags if anchor_nodes is not None and len(anchor_nodes) >= 2 and mol.GetNumAtoms() > max(anchor_nodes): anchor_frag = next((f for f in frags if anchor_nodes[0] in f), None) if anchor_frag is None or any(a not in anchor_frag for a in anchor_nodes[1:]): score *= 0.5 scores.append(score) except Exception: scores.append(0.0) return torch.tensor(scores, dtype=torch.float32) def __call__(self, molecule_list: list, mols=None) -> torch.Tensor: return self.get_scores(molecule_list, mols=mols) # --------------------------------------------------------------------------- # SA score soft guidance scorer (scores the linker, not the assembled PROTAC) # --------------------------------------------------------------------------- class PROTACLinkerSAScorer: """ Soft guidance scorer based on synthetic accessibility of the linker itself. Maps SA score [1, sa_threshold] → [1.0, 0.0] linearly (soft gradient). SA > sa_threshold hard-zeros to 0.0, acting as both a guidance penalty and a hard validity gate. Scoring the linker (not the assembled PROTAC) keeps the signal focused on what the generator actually controls. """ def __init__(self, atom_decoder: list[str], sa_threshold: float = 6.0): import sys from rdkit.Chem import RDConfig sys.path.insert(0, os.path.join(RDConfig.RDContribDir, "SA_Score")) import sascorer as _sascorer self._sascorer = _sascorer self.atom_decoder = atom_decoder self.sa_threshold = sa_threshold def get_scores(self, molecule_list: list, mols=None) -> torch.Tensor: if mols is None: mols = prebuild_mols(molecule_list, self.atom_decoder) scores = [] for mol in mols: if mol is None: scores.append(0.0) continue try: sa = self._sascorer.calculateScore(mol) scores.append(max(0.0, (self.sa_threshold - sa) / (self.sa_threshold - 1.0))) except Exception: scores.append(0.0) return torch.tensor(scores, dtype=torch.float32) def __call__(self, molecule_list: list, mols=None) -> torch.Tensor: return self.get_scores(molecule_list, mols=mols)