StaMP v2 release: 5-seed ensemble weights, ProteinMPNN backbone, inference, honest model card
2f40071 verified | """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 | |
| 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) | |
| 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 --------------------------------------------------------------------- | |
| 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() | |