"""StaMP — Stability from Message-Passing. A small, fast structure-conditioned predictor of single-point mutation stability (ddG, kcal/mol) built on frozen ProteinMPNN features. Positive ddG = destabilizing. Given a PDB/mmCIF backbone, StaMP scores any single amino-acid substitution, or scans every possible substitution at every residue (deep mutational scan) in a few seconds on CPU. Pipeline (all frozen except the small head): 1. ProteinMPNN v_48_020 (MIT, Dauparas et al. 2022) is run once per chain as a feature extractor: 3 encoder-layer outputs (3x128), the W_s(wt) embedding (128), and conditional log-probs averaged over 8 random decoding orders (21). -> R[L, 533] 2. For each residue, the 48 nearest structural neighbours (ProteinMPNN's Ca k-NN graph) are mean-pooled and concatenated with the centre features and two burial counts. 3. A 592k-param gated head emits one score per amino acid; ddG = o[mut] - o[wt], so ddG(x->x) = 0 exactly. The released model is a 5-seed ensemble (mean of the heads). Usage: python predict.py --pdb 1abc.pdb --chain A --mut A42G python predict.py --pdb 1abc.pdb --chain A --scan # full single-mutation scan See README.md for benchmarks and caveats. """ import os, argparse, glob, numpy as np, torch from protein_mpnn_utils import ProteinMPNN ALPHABET = 'ACDEFGHIKLMNPQRSTVWYX' HERE = os.path.dirname(os.path.abspath(__file__)) THREE = {'ALA': 'A', 'ARG': 'R', 'ASN': 'N', 'ASP': 'D', 'CYS': 'C', 'GLN': 'Q', 'GLU': 'E', 'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LEU': 'L', 'LYS': 'K', 'MET': 'M', 'PHE': 'F', 'PRO': 'P', 'SER': 'S', 'THR': 'T', 'TRP': 'W', 'TYR': 'Y', 'VAL': 'V'} def parse_chain(path, chain): """Minimal backbone parser for one chain: returns (seq, resnums, X[L,4,3]). Keeps only residues with all of N, CA, C, O. resnums are the PDB residue labels (resseq+icode) in order, so --mut can be given in PDB numbering with --pdb_numbering.""" res = {} for line in open(path): if not line.startswith('ATOM'): continue if chain and line[21] != chain: continue if line[16] not in (' ', 'A'): continue name = line[12:16].strip() if name not in ('N', 'CA', 'C', 'O'): continue rn = line[22:27].strip() aa = THREE.get(line[17:20].strip()) if aa is None: continue xyz = [float(line[30:38]), float(line[38:46]), float(line[46:54])] res.setdefault(rn, {'aa': aa, 'at': {}})['at'][name] = xyz keys = [k for k in res if len(res[k]['at']) == 4] keys.sort(key=lambda k: (int(''.join(c for c in k if c.isdigit() or c == '-')), k)) seq = ''.join(res[k]['aa'] for k in keys) X = np.array([[res[k]['at'][a] for a in ('N', 'CA', 'C', 'O')] for k in keys], np.float32) return seq, keys, X # ----- frozen ProteinMPNN feature extractor ------------------------------------------- def load_mpnn(weights, device): ckpt = torch.load(weights, map_location='cpu', weights_only=False) m = ProteinMPNN(ca_only=False, num_letters=21, node_features=128, edge_features=128, hidden_dim=128, num_encoder_layers=3, num_decoder_layers=3, k_neighbors=ckpt['num_edges'], augment_eps=0.0) m.load_state_dict(ckpt['model_state_dict']) m.to(device).eval() for p in m.parameters(): p.requires_grad = False return m @torch.no_grad() def encode(model, X, S, mask, chain_M, residue_idx, chain_enc): E, E_idx = model.features(X, mask, residue_idx, chain_enc) h_V = torch.zeros((E.shape[0], E.shape[1], E.shape[-1]), device=E.device, dtype=E.dtype) h_E = model.W_e(E) ma = torch.gather(mask, 1, E_idx.view(E_idx.shape[0], -1)).view(E_idx.shape) ma = mask.unsqueeze(-1) * ma outs = [] for layer in model.encoder_layers: h_V, h_E = layer(h_V, h_E, E_idx, mask, ma) outs.append(h_V) return torch.cat(outs, -1), model.W_s(S) @torch.no_grad() def cond_logprobs(model, X, S, mask, chain_M, residue_idx, chain_enc, n_orders, gen): acc = None for _ in range(n_orders): randn = torch.randn(chain_M.shape, device=X.device, generator=gen) lp = model.conditional_probs(X, S, mask, chain_M * mask, residue_idx, chain_enc, randn) lp = lp[:, :, 0, :] if lp.dim() == 4 else lp acc = lp if acc is None else torch.logaddexp(acc, lp) return acc - np.log(n_orders) def knn(X, k=48): Ca = torch.as_tensor(X)[:, 1].float() D = torch.sqrt(((Ca[:, None] - Ca[None]) ** 2).sum(-1) + 1e-6) k = min(k, Ca.shape[0]) d, idx = torch.topk(D, k, dim=-1, largest=False) return idx.numpy().astype(np.int64), d.numpy().astype(np.float32) # ----- trained head ------------------------------------------------------------------- import torch.nn as nn class MSHead(nn.Module): def __init__(self, in_dim, gate_dim=256, hidden=(128, 64), vocab=21, dropout=0.0): super().__init__() self.W1 = nn.Linear(in_dim, gate_dim) self.W2 = nn.Linear(in_dim, gate_dim) self.norm = nn.LayerNorm(gate_dim) layers, d = [], gate_dim for h in hidden: layers += [nn.Linear(d, h), nn.ReLU()] if dropout: layers += [nn.Dropout(dropout)] d = h layers += [nn.Linear(d, vocab)] self.mlp = nn.Sequential(*layers) def forward(self, x, wt_idx, mut_idx): h = self.norm(self.W1(x) * torch.sigmoid(self.W2(x))) o = self.mlp(h) return o.gather(1, mut_idx[:, None]).squeeze(1) - o.gather(1, wt_idx[:, None]).squeeze(1) def scores(self, x): return self.mlp(self.norm(self.W1(x) * torch.sigmoid(self.W2(x)))) class MSNbrHead(nn.Module): """Released config: mode='mean', k=48. (attn path kept for completeness.)""" def __init__(self, xc_dim, r_dim, mode='mean', k=48, gate_dim=256, hidden=(128, 64), dropout=0.0, heads=4, dh=32, ctx_dim=128): super().__init__() self.mode, self.k = mode, k in_dim = xc_dim + r_dim + 2 if mode == 'mean' else xc_dim self.head = MSHead(in_dim, gate_dim, hidden, dropout=dropout) def context(self, xc, Rn, dist, valid): burial = torch.stack([((dist < 8) & valid).float().sum(1), ((dist < 12) & valid).float().sum(1)], -1) / 10.0 Rk, vk = Rn[:, 1:self.k + 1], valid[:, 1:self.k + 1] w = vk.float()[..., None] m = (Rk * w).sum(1) / w.sum(1).clamp(min=1) return torch.cat([xc, m, burial], -1) def forward(self, xc, Rn, dist, valid, wt_idx, mut_idx): return self.head(self.context(xc, Rn, dist, valid), wt_idx, mut_idx) def scores(self, xc, Rn, dist, valid): return self.head.scores(self.context(xc, Rn, dist, valid)) def load_ensemble(ckpt_dir, device): members = [] for cp in sorted(glob.glob(os.path.join(ckpt_dir, 'stamp_s*.pt'))): ck = torch.load(cp, map_location='cpu', weights_only=False) c = ck['cfg'] h = MSNbrHead(ck['mu'].shape[1], ck['r_mu'].shape[1], c['mode'], c['k'], c['gate_dim'], tuple(int(x) for x in c['hidden'].split(',')), c['dropout']) h.load_state_dict(ck['state_dict']); h.to(device).eval() members.append((h, ck['mu'], ck['sd'], ck['r_mu'], ck['r_sd'])) if not members: raise FileNotFoundError(f'no stamp_s*.pt checkpoints in {ckpt_dir}') return members # ----- end-to-end --------------------------------------------------------------------- @torch.no_grad() def featurize_chain(mpnn, pdb_path, chain, device): seq, resnums, Xnp = parse_chain(pdb_path, chain) if len(seq) < 1: raise ValueError(f'no complete-backbone residues for chain {chain!r} in {pdb_path}') L = len(seq) X = torch.tensor(Xnp, device=device)[None] S = torch.tensor([[ALPHABET.index(a) if a in ALPHABET else 20 for a in seq]], device=device) mask = torch.ones((1, L), device=device) chain_M = torch.ones((1, L), device=device) residue_idx = torch.arange(L, device=device)[None] chain_enc = torch.ones((1, L), dtype=torch.long, device=device) enc, hseq = encode(mpnn, X, S, mask, chain_M, residue_idx, chain_enc) gen = torch.Generator(device=device).manual_seed(0) lp = cond_logprobs(mpnn, X, S, mask, chain_M, residue_idx, chain_enc, 8, gen) enc, hseq, lp = enc[0].float().cpu().numpy(), hseq[0].float().cpu().numpy(), lp[0].float().cpu().numpy() R = np.concatenate([enc, hseq, lp], -1).astype(np.float32) # [L, 533] E, D = knn(Xnp, 48) return seq, resnums, R, E, D def predict_muts(members, seq, R, E, D, muts, device): """muts: list of (wt, pos0, mut). Returns ensemble-mean ddG for each.""" xs, Rns, ds, wts, mus, valids = [], [], [], [], [], [] for wt, i, mut in muts: wl, ml = R[i][-21:][ALPHABET.index(wt)], R[i][-21:][ALPHABET.index(mut)] xs.append(np.concatenate([R[i], [wl, ml, ml - wl]]).astype(np.float32)) Rns.append(R[E[i]]); ds.append(D[i]); valids.append(D[i] < 98) wts.append(ALPHABET.index(wt)); mus.append(ALPHABET.index(mut)) X = np.stack(xs); Rn = np.stack(Rns); Dn = np.stack(ds) wt_i = torch.tensor(wts, device=device); mu_i = torch.tensor(mus, device=device) valid = torch.tensor(np.stack(valids), device=device); Dt = torch.tensor(Dn, device=device) preds = [] with torch.no_grad(): for h, m, s, rm, rs in members: preds.append(h(torch.tensor((X - m) / s).float().to(device), torch.tensor((Rn - rm[None]) / rs[None]).float().to(device), Dt, valid, wt_i, mu_i).cpu().numpy()) return np.mean(preds, 0) def main(): ap = argparse.ArgumentParser(description='StaMP single-mutation ddG predictor') ap.add_argument('--pdb', required=True) ap.add_argument('--chain', default='A') ap.add_argument('--mut', help='e.g. A42G (wt, position, mut). Position is 1-based along the ' 'chain by default, or the PDB residue number with --pdb_numbering.') ap.add_argument('--pdb_numbering', action='store_true', help='interpret --mut position as the PDB residue number, not 1-based index') ap.add_argument('--scan', action='store_true', help='score every single substitution') ap.add_argument('--ckpt_dir', default=os.path.join(HERE, 'weights')) ap.add_argument('--mpnn', default=os.path.join(HERE, 'weights', 'v_48_020.pt')) ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu') ap.add_argument('--out', default=None) a = ap.parse_args() dev = torch.device(a.device) mpnn = load_mpnn(a.mpnn, dev) members = load_ensemble(a.ckpt_dir, dev) seq, resnums, R, E, D = featurize_chain(mpnn, a.pdb, a.chain, dev) if a.mut: wt, posstr, mut = a.mut[0], a.mut[1:-1], a.mut[-1] if a.pdb_numbering: assert posstr in resnums, f'residue {posstr} not found in chain {a.chain}' i = resnums.index(posstr) else: i = int(posstr) - 1 assert 0 <= i < len(seq), f'position {posstr} out of range (len {len(seq)})' assert seq[i] == wt, f'expected {wt} at position {posstr}, chain has {seq[i]}' ddg = predict_muts(members, seq, R, E, D, [(wt, i, mut)], dev)[0] print(f'{a.mut}\tddG = {ddg:+.3f} kcal/mol\t({"destabilizing" if ddg>0 else "stabilizing"})') elif a.scan: muts, labels = [], [] for i, wt in enumerate(seq): if wt not in ALPHABET[:20]: continue for mut in ALPHABET[:20]: if mut == wt: continue muts.append((wt, i, mut)); labels.append(f'{wt}{i+1}{mut}') ddg = predict_muts(members, seq, R, E, D, muts, dev) rows = sorted(zip(labels, ddg), key=lambda r: r[1]) out = a.out or (os.path.splitext(os.path.basename(a.pdb))[0] + f'_{a.chain}_scan.tsv') with open(out, 'w') as f: f.write('mutation\tddG_kcal_mol\n') for lab, d in sorted(zip(labels, ddg), key=lambda r: r[0]): f.write(f'{lab}\t{d:.4f}\n') print(f'scanned {len(muts)} mutations over {len(seq)} residues -> {out}') print('\nMost stabilizing:') for lab, d in rows[:10]: print(f' {lab}\t{d:+.3f}') print('Most destabilizing:') for lab, d in rows[-10:]: print(f' {lab}\t{d:+.3f}') else: ap.error('give --mut A42G or --scan') if __name__ == '__main__': main()