| |
| """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() |
|
|
