File size: 4,857 Bytes
72d1407 | 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 | #!/usr/bin/env python
"""Inference for the DPP-4 Kd LTN (Delower/dpp4-kd-ltn).
Predicts pKd for human DPP-4 (UniProt P27487 / CHEMBL284) from a SMILES string.
Features are exactly those used in training: Morgan fingerprint (r=2, 2048 bits)
concatenated with the fixed DPP-4 ESM-2 650M embedding (1280 dims) = 3328, no scaling.
Usage:
python predict_kd.py "Cn1c(=O)cc(N2CCC[C@@H](N)C2)n(Cc2ccccc2C#N)c1=O"
python predict_kd.py --csv compounds.csv --smiles-col canonical_smiles --out preds.csv
Outputs per compound: predicted pKd, Kd in nM (= 10^(9-pKd)),
deltaG in kcal/mol (= -1.3642 * pKd), and potency class
(potent pKd>=7 / moderate 6-7 / weak <6).
Requires: torch, rdkit, numpy. Files ltn_stage2.pt, esm2_dpp4.npz and config.json
must sit next to this script (or pass --model-dir).
"""
import argparse, csv, json, os, sys
import numpy as np
import torch
import torch.nn as nn
HERE = os.path.dirname(os.path.abspath(__file__))
class LtnNet(nn.Module):
"""Same architecture as training (train_kd_ltn.py)."""
def __init__(self, in_dim=3328):
super().__init__()
self.trunk = nn.Sequential(
nn.Linear(in_dim, 768), nn.GELU(), nn.Dropout(0.1),
nn.Linear(768, 384), nn.GELU(), nn.Dropout(0.1),
nn.Linear(384, 192), nn.GELU())
self.reg = nn.Linear(192, 1)
self.clf = nn.Linear(192, 3)
def forward(self, x):
h = self.trunk(x)
return self.reg(h).squeeze(-1), self.clf(h)
def band_name(pkd):
return "potent" if pkd >= 7.0 else ("moderate" if pkd >= 6.0 else "weak")
def fingerprints(smiles_list):
from rdkit import Chem
from rdkit.Chem import rdFingerprintGenerator
gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
out = np.zeros((len(smiles_list), 2048), dtype=np.float32)
for i, s in enumerate(smiles_list):
m = Chem.MolFromSmiles(s)
if m is not None:
out[i] = np.asarray(gen.GetFingerprint(m), dtype=np.float32)
return out
class KdLtnPredictor:
def __init__(self, model_dir=HERE):
with open(os.path.join(model_dir, "config.json")) as f:
self.config = json.load(f)
z = np.load(os.path.join(model_dir, "esm2_dpp4.npz"), allow_pickle=True)
keys = z["target_keys"].tolist()
key = self.config["target"]["esm2_target_key"]
self.protein_vec = z["embeddings"][keys.index(key)].astype(np.float32)
self.net = LtnNet(self.config["features"]["total_dim"])
sd = torch.load(os.path.join(model_dir, self.config["weights_file"]),
map_location="cpu", weights_only=True)
self.net.load_state_dict(sd)
self.net.eval()
def predict(self, smiles_list):
fp = fingerprints(smiles_list)
X = np.concatenate([fp, np.tile(self.protein_vec, (len(smiles_list), 1))], axis=1)
with torch.no_grad():
pkd, _ = self.net(torch.tensor(X, dtype=torch.float32))
out = []
for s, p in zip(smiles_list, pkd.numpy()):
p = float(p)
out.append({
"smiles": s,
"predicted_pKd": round(p, 4),
"predicted_Kd_nM": 10 ** (9 - p),
"predicted_deltaG_kcal_mol": round(-1.3642 * p, 4),
"potency_class": band_name(p),
})
return out
def main():
ap = argparse.ArgumentParser(description="DPP-4 Kd LTN inference (human DPP-4, P27487)")
ap.add_argument("smiles", nargs="*", help="one or more SMILES strings")
ap.add_argument("--csv", help="input CSV with a SMILES column")
ap.add_argument("--smiles-col", default="smiles", help="SMILES column name for --csv")
ap.add_argument("--out", help="write results to this CSV instead of stdout")
ap.add_argument("--model-dir", default=HERE)
args = ap.parse_args()
smiles = list(args.smiles)
if args.csv:
with open(args.csv, newline="") as f:
smiles += [r[args.smiles_col] for r in csv.DictReader(f)]
if not smiles:
ap.error("give at least one SMILES or --csv")
rows = KdLtnPredictor(args.model_dir).predict(smiles)
fields = ["smiles", "predicted_pKd", "predicted_Kd_nM", "predicted_deltaG_kcal_mol", "potency_class"]
if args.out:
with open(args.out, "w", newline="") as f:
w = csv.DictWriter(f, fieldnames=fields)
w.writeheader()
for r in rows:
w.writerow({**r, "predicted_Kd_nM": f"{r['predicted_Kd_nM']:.4g}"})
print(f"wrote {len(rows)} predictions to {args.out}")
else:
for r in rows:
print(f"pKd {r['predicted_pKd']:7.4f} | Kd {r['predicted_Kd_nM']:12.4g} nM | "
f"dG {r['predicted_deltaG_kcal_mol']:8.4f} kcal/mol | {r['potency_class']:8s} | {r['smiles']}")
if __name__ == "__main__":
main()
|