File size: 16,261 Bytes

f12d3b3

"""
PROTAC Analyzer
Analyzes PROTAC linkers for structural validity

The linker is the generated component; warhead and E3 ligand are fixed inputs.

Hard constraints (validity = 1 if all pass, 0 if any fails):
1. Attachment site availability: linker has reactive termini for both warhead and E3 ends
2. Linker feasibility: ring-closure constraints, long carbon chains
3. Linker atom count: 3-20 heavy atoms in the linker
4. Assembled PROTAC MW: full PROTAC within 700-1200 Da (requires warhead + E3)
5. Representational validity: RDKit sanitization, valence, round-trip
6. Synthetic accessibility: SA score ≤ threshold (default 6.0)

All methods accept Chem.Mol objects — the molecular graph produced by DeFoG.
SMILES conversion happens once at the DeFoG graph → RDKit Mol boundary upstream;
this class never calls MolFromSmiles on method inputs.

Note: Bond-forming chemistry (which atom connects to which) is in protac_builder.py.
"""

from __future__ import annotations

import os
import sys
import torch
from rdkit import Chem, RDConfig
from rdkit.Chem import Descriptors, AllChem, rdMolDescriptors
from typing import Optional

sys.path.insert(0, os.path.join(RDConfig.RDContribDir, "SA_Score"))
import sascorer


LINKER_ATOM_MIN = 3
LINKER_ATOM_MAX = 25

PROTAC_MW_MIN = 700.0
PROTAC_MW_MAX = 1250.0

SA_SCORE_MAX = 6.0  # 1 (easy) → 10 (hard); ≤6 keeps synthetically accessible linkers

_LONG_CHAIN_SMARTS = "[CH2]" + "[CH2]" * 15   # 16 consecutive CH2

# Atoms allowed in PROTAC linkers — excludes halogens (I, Br optionally) and P
# which appear as nonsensical bridge atoms when generated from GuacaMol atom decoder.


class PROTACAnalyzer:
    """
    Validates PROTAC linkers against hard chemical and structural constraints.

    All check_* methods and the batch interface accept Chem.Mol objects only.
    The DeFoG graph → Chem.Mol conversion happens once upstream; this class
    never parses SMILES strings.

    Parameters
    ----------
    warhead : Chem.Mol
        Fixed warhead molecule. Required for MW check.
    e3 : Chem.Mol
        Fixed E3-ligand molecule. Required for MW check.
    warhead_attachment_smarts : str
        SMARTS for the reactive group the linker must have on the warhead-facing end.
    e3_attachment_smarts : str
        SMARTS for the reactive group the linker must have on the E3-facing end.
    linker_atom_min / linker_atom_max : int
        Heavy-atom count bounds for the linker (default 3–20).
    protac_mw_min / protac_mw_max : float
        Assembled PROTAC MW window in Da (default 700–1200).
    check_mw : bool
        Enforce the MW window. Requires warhead and e3.
    sa_score_max : float
        Maximum SA score for the linker (default 6.0). Set to None to disable.
    calculate_properties : bool
        Compute optional descriptors in get_detailed_scores.
    """

    def __init__(
        self,
        warhead: Chem.Mol,
        e3: Chem.Mol,
        warhead_attachment_smarts: str,
        e3_attachment_smarts: str,
        linker_atom_min: int = LINKER_ATOM_MIN,
        linker_atom_max: int = LINKER_ATOM_MAX,
        protac_mw_min: float = PROTAC_MW_MIN,
        protac_mw_max: float = PROTAC_MW_MAX,
        check_mw: bool = True,
        sa_score_max: Optional[float] = SA_SCORE_MAX,
        calculate_properties: bool = True,
    ):
        if warhead is None:
            raise ValueError("warhead must be a Chem.Mol object")
        if e3 is None:
            raise ValueError("e3 must be a Chem.Mol object")

        self._warhead_mol = warhead
        self._e3_mol = e3
        self.linker_atom_min = linker_atom_min
        self.linker_atom_max = linker_atom_max
        self.protac_mw_min = protac_mw_min
        self.protac_mw_max = protac_mw_max
        self.check_mw = check_mw
        self.sa_score_max = sa_score_max
        self.calculate_properties = calculate_properties

        self._warhead_pat = Chem.MolFromSmarts(warhead_attachment_smarts)
        self._e3_pat = Chem.MolFromSmarts(e3_attachment_smarts)
        if self._warhead_pat is None:
            raise ValueError(f"Invalid warhead_attachment_smarts: {warhead_attachment_smarts!r}")
        if self._e3_pat is None:
            raise ValueError(f"Invalid e3_attachment_smarts: {e3_attachment_smarts!r}")

        self._long_chain_pat = Chem.MolFromSmarts(_LONG_CHAIN_SMARTS)

        self.property_flags = {
            "mf": True, "mw": True, "polararea": True, "xlogp": True,
            "heavycnt": True, "hbonddonor": True, "hbondacc": True, "rotbonds": True,
        }

    # ------------------------------------------------------------------
    # Property flag helpers
    # ------------------------------------------------------------------

    def set_property_flags(self, **kwargs):
        for k, v in kwargs.items():
            if k in self.property_flags:
                self.property_flags[k] = bool(v)
            else:
                raise ValueError(f"Unknown property: {k}. Valid: {list(self.property_flags)}")

    def enable_all_properties(self):
        for k in self.property_flags:
            self.property_flags[k] = True

    def disable_all_properties(self):
        for k in self.property_flags:
            self.property_flags[k] = False

    # ------------------------------------------------------------------
    # Hard constraint checks — all accept Chem.Mol
    # ------------------------------------------------------------------

    def check_attachment_sites(self, mol: Chem.Mol) -> bool:
        """Constraint 1: linker must have two *distinct* attachment atoms —
        one for the warhead end and one for the E3 end. Requiring the same
        atom to satisfy both patterns (e.g. a single CHO) is rejected."""
        warhead_atoms = {m[0] for m in mol.GetSubstructMatches(self._warhead_pat)}
        e3_atoms      = {m[0] for m in mol.GetSubstructMatches(self._e3_pat)}
        return any(w != e for w in warhead_atoms for e in e3_atoms)

    def check_linker_atom_count(self, mol: Chem.Mol) -> bool:
        """Constraint 3: heavy-atom count within [min, max]."""
        n = mol.GetNumHeavyAtoms()
        return self.linker_atom_min <= n <= self.linker_atom_max

    def check_protac_mw(self, linker_mol: Chem.Mol) -> bool:
        """
        Constraint 4: MW(warhead) + MW(linker) + MW(E3) - 36 Da
        (subtracts two amide-bond condensation water losses).
        """
        if not self.check_mw:
            return True
        mw = (Descriptors.MolWt(self._warhead_mol)
              + Descriptors.MolWt(linker_mol)
              + Descriptors.MolWt(self._e3_mol)
              - 36.0)
        return self.protac_mw_min <= mw <= self.protac_mw_max

    def _find_linker_rings(self, mol: Chem.Mol):
        linker_rings = []
        for ring in mol.GetRingInfo().AtomRings():
            ring_atoms = [mol.GetAtomWithIdx(i) for i in ring]
            if all(a.GetIsAromatic() for a in ring_atoms):
                continue
            size = len(ring)
            symbols = [a.GetSymbol() for a in ring_atoms]
            if size == 5 and "N" in symbols:
                c_count = sum(
                    1 for idx in ring
                    for nb in mol.GetAtomWithIdx(idx).GetNeighbors()
                    if nb.GetSymbol() == "O"
                    and nb.GetIdx() not in ring
                    and mol.GetBondBetweenAtoms(idx, nb.GetIdx()).GetBondType()
                    == Chem.BondType.DOUBLE
                )
                if c_count >= 2:
                    linker_rings.append(ring)
                    continue
            linker_rings.append(ring)
        return linker_rings

    def check_ring_closure_constraints(self, mol: Chem.Mol) -> bool:
        """
        Constraint 2a:
        - 3-ring with C=C → reject (severe strain)
        - 4-ring with >2 sp2 centers → reject
        - ring > 12 atoms → reject
        """
        for ring in self._find_linker_rings(mol):
            size = len(ring)
            if size < 3:
                return False
            if size == 3:
                for idx in ring:
                    atom = mol.GetAtomWithIdx(idx)
                    if (atom.GetSymbol() == "C"
                            and atom.GetHybridization() == Chem.HybridizationType.SP2):
                        for nb in atom.GetNeighbors():
                            bond = mol.GetBondBetweenAtoms(idx, nb.GetIdx())
                            if bond.GetBondType() == Chem.BondType.DOUBLE and nb.GetSymbol() == "C":
                                return False
            if size == 4:
                sp2 = sum(1 for idx in ring
                          if mol.GetAtomWithIdx(idx).GetHybridization()
                          == Chem.HybridizationType.SP2)
                if sp2 > 2:
                    return False
            if size > 12:
                return False
        return True

    def check_long_chain(self, mol: Chem.Mol) -> bool:
        """Constraint 2b: reject ≥16 consecutive CH2 groups."""
        if self._long_chain_pat is None:
            return True
        return not mol.HasSubstructMatch(self._long_chain_pat)

    def check_linker_feasibility(self, mol: Chem.Mol) -> bool:
        return self.check_ring_closure_constraints(mol) and self.check_long_chain(mol)

    def check_representational_validity(self, mol: Chem.Mol) -> bool:
        """Constraint 5: sanitization, round-trip SMILES, valence, 2D coords."""
        try:
            Chem.SanitizeMol(mol)
        except Exception:
            return False
        try:
            if Chem.MolFromSmiles(Chem.MolToSmiles(mol)) is None:
                return False
        except Exception:
            return False
        try:
            for prob in Chem.DetectChemistryProblems(mol):
                if prob.GetType() in ("AtomValenceException", "KekulizeException"):
                    return False
        except Exception:
            return False
        try:
            AllChem.Compute2DCoords(mol)
        except Exception:
            return False
        return True

    # ------------------------------------------------------------------
    # Validation entry point
    # ------------------------------------------------------------------

    def check_sa_score(self, mol: Chem.Mol) -> bool:
        """Constraint 6: SA score ≤ sa_score_max (skipped if sa_score_max is None)."""
        if self.sa_score_max is None:
            return True
        score = sascorer.calculateScore(mol)
        if score is None:
            return False
        return score <= self.sa_score_max

    def _run_checks(self, mol: Chem.Mol) -> dict:
        return {
            "connected":         len(Chem.GetMolFrags(mol)) == 1,
            "attachment_sites":  self.check_attachment_sites(mol),
            "ring_closure":      self.check_ring_closure_constraints(mol),
            "long_chain":        self.check_long_chain(mol),
            "linker_atom_count": self.check_linker_atom_count(mol),
            "protac_mw":         self.check_protac_mw(mol),
            "representational":  self.check_representational_validity(mol),
            "sa_score":          self.check_sa_score(mol),
        }

    def validate(self, linker_mol: Chem.Mol) -> bool:
        """Run all hard constraints. Returns True only if every check passes."""
        return all(self._run_checks(linker_mol).values())

    # ------------------------------------------------------------------
    # Optional property computation
    # ------------------------------------------------------------------

    def calculate_linker_properties(self, mol: Chem.Mol) -> dict:
        if not self.calculate_properties:
            return {}
        props = {}
        if self.property_flags.get("mf"):
            props["mf"] = rdMolDescriptors.CalcMolFormula(mol)
        if self.property_flags.get("mw"):
            props["mw"] = round(Descriptors.MolWt(mol), 2)
        if self.property_flags.get("polararea"):
            props["polararea"] = round(Descriptors.TPSA(mol), 2)
        if self.property_flags.get("xlogp"):
            props["xlogp"] = round(Descriptors.MolLogP(mol), 2)
        if self.property_flags.get("heavycnt"):
            props["heavycnt"] = mol.GetNumHeavyAtoms()
        if self.property_flags.get("hbonddonor"):
            props["hbonddonor"] = Descriptors.NumHDonors(mol)
        if self.property_flags.get("hbondacc"):
            props["hbondacc"] = Descriptors.NumHAcceptors(mol)
        if self.property_flags.get("rotbonds"):
            props["rotbonds"] = Descriptors.NumRotatableBonds(mol)
        return props

    def estimate_protac_mw(self, linker_mol: Chem.Mol) -> Optional[float]:
        if not (self._warhead_mol and self._e3_mol):
            return None
        return (Descriptors.MolWt(self._warhead_mol)
                + Descriptors.MolWt(linker_mol)
                + Descriptors.MolWt(self._e3_mol)
                - 36.0)

    # ------------------------------------------------------------------
    # Batch interface — list[Chem.Mol]
    # ------------------------------------------------------------------

    def get_scores(self, linker_mols: list[Chem.Mol]) -> list[int]:
        """
        Validate a batch of linker Mol objects.
        Returns list of 0/1 scores (1 = all constraints pass).
        """
        return [
            1 if all(self._run_checks(mol).values()) else 0
            for mol in linker_mols
        ]

    def get_detailed_scores(self, linker_mols: list[Chem.Mol]) -> list[dict]:
        """
        Per-constraint breakdown for each linker Mol.
        Keys: valid, smiles (canonical), checks, properties, estimated_protac_mw.
        """
        detailed = []
        for mol in linker_mols:
            checks = self._run_checks(mol)
            entry = {
                "valid":      1 if all(checks.values()) else 0,
                "smiles":     Chem.MolToSmiles(mol),
                "checks":     checks,
                "properties": self.calculate_linker_properties(mol),
            }
            est_mw = self.estimate_protac_mw(mol)
            if est_mw is not None:
                entry["estimated_protac_mw"] = round(est_mw, 2)
            detailed.append(entry)
        return detailed

    # ------------------------------------------------------------------
    # Tensor interface for guidance loop
    # ------------------------------------------------------------------

    def __call__(self, linker_mols: list[Chem.Mol]) -> torch.Tensor:
        """Return float tensor of validity scores (0 or 1) for the guidance loop."""
        return torch.tensor(self.get_scores(linker_mols), dtype=torch.float32)


# ---------------------------------------------------------------------------
# Self-test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys, os
    sys.path.insert(0, os.path.dirname(__file__))
    from protac.ErlotinibVHL import EGFR_WARHEAD_SMILES, VHL_E3_SMILES

    analyzer = PROTACAnalyzer(
        warhead=Chem.MolFromSmiles(EGFR_WARHEAD_SMILES),
        e3=Chem.MolFromSmiles(VHL_E3_SMILES),
        warhead_attachment_smarts="[NX3;H2]",
        e3_attachment_smarts="[NX3;H2]",
    )

    test_cases = [
        ("NCCOCCOCCN",           "Valid PEG-diamine"),
        ("NCCCCN",               "Valid butane-1,4-diamine"),
        ("NCCCCCCCCCCCCCCCCCN",  "Invalid: too long + long chain"),
        ("NCC1=CC1N",            "Invalid: strained 3-ring"),
        ("N",                    "Invalid: too short"),
    ]

    print("=" * 70)
    print("PROTAC Linker Analyzer — Self-Test")
    print("=" * 70)

    for smiles, desc in test_cases:
        mol = Chem.MolFromSmiles(smiles)   # SMILES used only here in the test
        if mol is None:
            print(f"\n[INVALID] {desc} — unparseable SMILES")
            continue
        result = analyzer.get_detailed_scores([mol])[0]
        tag = "VALID" if result["valid"] == 1 else "INVALID"
        print(f"\n[{tag}] {desc}")
        for k, v in result["checks"].items():
            print(f"  {'✓' if v else '✗'} {k}")
        if "estimated_protac_mw" in result:
            print(f"  estimated PROTAC MW: {result['estimated_protac_mw']} Da")