File size: 8,759 Bytes
f12d3b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | #!/usr/bin/env python3
"""Train a target-specific PROTAC linker classifier.
For a given protac_module (e.g. ErlotinibVHL, BRD9VHL, CDK6CRBN):
1. Load PROTAC-DB entries matching the target + E3 pair
2. Extract linker SMILES from each (strip warhead + E3 substructures)
3. Train LogisticRegression on Morgan FP (radius=2, 1024 bits) of
target-specific linkers (positives) vs the existing generic negatives
4. Save as protac_linker_clf_<module>.pkl
DeFoG generation can then load this classifier (when --protac_module matches)
and steer linker generation toward the chemistry distribution that's actually
represented in the target's PROTAC-DB active subset.
Usage:
python train_target_specific_clf.py --protac-module BRD9VHL
python train_target_specific_clf.py --protac-module CDK6CRBN
python train_target_specific_clf.py --protac-module ErlotinibVHL
"""
import argparse
import importlib
import json
import pickle
import sys
from pathlib import Path
import numpy as np
from rdkit import Chem, RDLogger
from rdkit.Chem import AllChem, DataStructs
from sklearn.linear_model import LogisticRegression
RDLogger.DisableLog("rdApp.*")
# Resolve paths relative to the repo layout (scripts/ -> ../src/...).
_REPO_ROOT = Path(__file__).resolve().parent.parent
CLF_DIR = _REPO_ROOT / "src" / "classifier"
PROTACDB_JSON = _REPO_ROOT / "src" / "degrademaster" / "data" / "PROTAC" / "name.json"
_TARGET_FILTER = {
"ErlotinibVHL": ("EGFR", "VHL"),
"BRD9VHL": ("BRD9", "VHL"),
"BRD4VHL": ("BRD4", "VHL"),
"KRASVHL": ("KRAS", "VHL"),
"ARCRBN": ("AR-V7", "CRBN"),
"ARCRBNv2": ("AR-V7", "CRBN"),
"BRD4CRBN": ("BRD4", "CRBN"),
"CDK6CRBN": ("CDK6", "CRBN"),
}
_TERMINI_FOR_NEGATIVES = {
# Use the matching termini's existing negatives so the classifier learns
# to distinguish target-specific linker chemistry from generic junk.
"ErlotinibVHL": "protac_linker_clf.pkl", # cooh
"BRD9VHL": "protac_linker_clf.pkl", # cooh
"BRD4VHL": "protac_linker_clf_amino_acid.pkl",
"KRASVHL": "protac_linker_clf.pkl",
"ARCRBN": "protac_linker_clf_amine.pkl",
"ARCRBNv2": "protac_linker_clf.pkl", # methyl uses generic
"BRD4CRBN": "protac_linker_clf_amine.pkl",
"CDK6CRBN": "protac_linker_clf.pkl",
}
def _morgan_fp(smi):
mol = Chem.MolFromSmiles(smi)
if mol is None:
return None
fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=1024)
arr = np.zeros(1024, dtype=np.float32)
DataStructs.ConvertToNumpyArray(fp, arr)
return arr
def _load_protac_module(module_name):
sys.path.insert(0, str(CLF_DIR / "protac"))
mod = importlib.import_module(module_name)
# Prefer the loose SMARTS (WARHEAD_PAT / E3_PAT) which match the core scaffold
# — full WARHEAD_MOL only matches one specific variant, missing PROTAC-DB
# diversity (e.g., gefitinib vs erlotinib both = "EGFR" but fail strict match).
wh = getattr(mod, "WARHEAD_PAT", None) or mod.WARHEAD_MOL
e3 = getattr(mod, "E3_PAT", None) or mod.E3_MOL
return wh, e3
def _extract_linker_smi(protac_mol, warhead_pat, e3_pat):
"""Get all atoms in PROTAC NOT belonging to either warhead-pattern or E3-pattern matches.
Take the largest connected fragment after stripping. Returns None on ambiguity."""
wh_atoms = set()
for match in protac_mol.GetSubstructMatches(warhead_pat):
wh_atoms.update(match)
e3_atoms = set()
for match in protac_mol.GetSubstructMatches(e3_pat):
e3_atoms.update(match)
if not wh_atoms or not e3_atoms:
return None
lk_atoms = set(range(protac_mol.GetNumAtoms())) - wh_atoms - e3_atoms
if not lk_atoms:
return None
# Bonds at the linker boundary (exactly one endpoint in lk_atoms)
bond_ids = [
b.GetIdx() for b in protac_mol.GetBonds()
if (b.GetBeginAtomIdx() in lk_atoms) != (b.GetEndAtomIdx() in lk_atoms)
]
if not bond_ids:
return None
try:
broken = Chem.FragmentOnBonds(protac_mol, bond_ids, addDummies=False)
# Find the fragment containing the linker atoms
frag_atom_lists = Chem.GetMolFrags(broken, asMols=False)
for atom_list in frag_atom_lists:
if set(atom_list).issubset(lk_atoms):
frag_atoms = list(atom_list)
if len(frag_atoms) >= 3: # min linker size
sub_mol = Chem.PathToSubmol(broken, [b.GetIdx() for b in broken.GetBonds()
if b.GetBeginAtomIdx() in frag_atoms
and b.GetEndAtomIdx() in frag_atoms])
try:
Chem.SanitizeMol(sub_mol)
smi = Chem.MolToSmiles(sub_mol)
if smi and '.' not in smi:
return smi
except Exception:
pass
# Fallback: take all fragments, find the one with exact linker atom count
frags = Chem.GetMolFrags(broken, asMols=True)
for frag in frags:
if frag.GetNumHeavyAtoms() == len(lk_atoms):
Chem.SanitizeMol(frag)
return Chem.MolToSmiles(frag)
except Exception:
return None
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--protac-module", required=True, choices=list(_TARGET_FILTER.keys()))
ap.add_argument("--out", default=None,
help="Output .pkl path (default: protac_linker_clf_<module>.pkl)")
args = ap.parse_args()
module = args.protac_module
target, e3 = _TARGET_FILTER[module]
out_path = Path(args.out) if args.out else CLF_DIR / f"protac_linker_clf_{module}.pkl"
print(f"=== Training target-specific classifier for {module} ({target}+{e3}) ===")
warhead_mol, e3_mol = _load_protac_module(module)
# 1. Load PROTAC-DB target+E3 entries (positives)
with open(PROTACDB_JSON) as f:
name = json.load(f)
target_entries = [n for n in name.values()
if target in n.get("tar_path", "")
and e3 in n.get("e3_ligase_path", "")
and "smiles" in n]
print(f" PROTAC-DB {target}+{e3} entries: {len(target_entries)}")
# 2. Extract linkers
linker_smis = set()
for n in target_entries:
m = Chem.MolFromSmiles(n["smiles"])
if m is None:
continue
lk_smi = _extract_linker_smi(m, warhead_mol, e3_mol)
if lk_smi:
linker_smis.add(lk_smi)
pos_smiles = sorted(linker_smis)
print(f" Extracted unique target-specific linkers (positives): {len(pos_smiles)}")
if len(pos_smiles) < 10:
print(f" ⚠ Only {len(pos_smiles)} positives — classifier may overfit. Consider "
f"augmenting with similar linkers or using termini-default classifier.")
# 3. Load negatives from existing termini-matched classifier
neg_pkl = CLF_DIR / _TERMINI_FOR_NEGATIVES[module]
with open(neg_pkl, "rb") as f:
existing = pickle.load(f)
neg_smiles = list(existing["neg_smiles"])
print(f" Negatives (from {neg_pkl.name}): {len(neg_smiles)}")
# 4. Featurize + train
X_pos = [_morgan_fp(s) for s in pos_smiles]
X_pos = [x for x in X_pos if x is not None]
X_neg = [_morgan_fp(s) for s in neg_smiles]
X_neg = [x for x in X_neg if x is not None]
X = np.vstack(X_pos + X_neg)
y = np.concatenate([np.ones(len(X_pos), dtype=int),
np.zeros(len(X_neg), dtype=int)])
print(f" Training set: {len(X_pos)} pos / {len(X_neg)} neg, {X.shape[1]}-dim FP")
# Class-balanced LogReg (since pos is much smaller than neg)
clf = LogisticRegression(class_weight="balanced", max_iter=2000)
clf.fit(X, y)
train_acc = clf.score(X, y)
print(f" Training accuracy: {train_acc*100:.1f}%")
# Sanity: predict prob on a few positives + negatives
sample_pos = clf.predict_proba(X_pos[:5])[:, 1]
sample_neg = clf.predict_proba(X_neg[:5])[:, 1]
print(f" Sample pos probs: {[f'{p:.3f}' for p in sample_pos]}")
print(f" Sample neg probs: {[f'{p:.3f}' for p in sample_neg]}")
# 5. Save
with open(out_path, "wb") as f:
pickle.dump({
"clf": clf,
"pos_smiles": pos_smiles,
"neg_smiles": neg_smiles,
"module": module,
"target": target,
"e3": e3,
"n_pos": len(X_pos),
"n_neg": len(X_neg),
}, f)
print(f" Saved → {out_path}")
if __name__ == "__main__":
main()
|