stamp/predict.py

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