| """Feature extraction shared by train.py and model.py (must be identical at train and inference time).""" |
| from itertools import product |
|
|
| import numpy as np |
| import pandas as pd |
| from rdkit import Chem, RDLogger |
| from rdkit.Chem import Descriptors, MACCSkeys, rdFingerprintGenerator |
|
|
| RDLogger.DisableLog('rdApp.*') |
| TARGETS = ['log10_kcat', 'log10_km'] |
|
|
| |
| AA = 'ACDEFGHIKLMNPQRSTVWY' |
| DIP = {a + b: i for i, (a, b) in enumerate(product(AA, AA))} |
| KD = dict(A=1.8, R=-4.5, N=-3.5, D=-3.5, C=2.5, Q=-3.5, E=-3.5, G=-0.4, H=-3.2, I=4.5, L=3.8, K=-3.9, |
| M=1.9, F=2.8, P=-1.6, S=-0.8, T=-0.7, W=-0.9, Y=-1.3, V=4.2) |
| AA_GROUPS = ['AVLIMFWC', 'STNQGPY', 'KRH', 'DE', 'FWYH', 'AGSCTP'] |
|
|
|
|
| def seq_feats(seq): |
| s = ''.join(c for c in seq.upper() if c in AA) or 'A' |
| L = len(s) |
| aac = np.array([s.count(a) for a in AA]) / L |
| dc = np.zeros(400) |
| for i in range(L - 1): |
| dc[DIP[s[i:i + 2]]] += 1 |
| dc /= max(1, L - 1) |
| kd = np.array([KD[c] for c in s]) |
| grp = [sum(s.count(c) for c in g) / L for g in AA_GROUPS] |
| thirds = [np.array([p.count(a) for a in AA]) / max(1, len(p)) |
| for p in (s[:L // 3], s[L // 3:2 * L // 3], s[2 * L // 3:])] |
| return np.concatenate([[np.log(L), kd.mean(), kd.std()], grp, aac, dc, *thirds]).astype(np.float32) |
|
|
|
|
| |
| _MORGAN = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048) |
| N_SMI = 2048 + 167 + len(Descriptors._descList) |
|
|
|
|
| def smi_feats(smi): |
| m = Chem.MolFromSmiles(smi) |
| if m is None: |
| return np.full(N_SMI, np.nan, dtype=np.float32) |
| fp = np.asarray(_MORGAN.GetCountFingerprintAsNumPy(m), dtype=np.float32) |
| mac = np.array(list(MACCSkeys.GenMACCSKeys(m)), dtype=np.float32) |
| d = [] |
| for _, f in Descriptors._descList: |
| try: |
| d.append(f(m)) |
| except Exception: |
| d.append(np.nan) |
| d = np.nan_to_num(np.array(d, dtype=np.float64), nan=np.nan, posinf=1e6, neginf=-1e6).clip(-1e6, 1e6) |
| return np.concatenate([fp, mac, d]).astype(np.float32) |
|
|
|
|
| def base_features(sequences, smiles): |
| """Enzyme + substrate features, computed once per unique string.""" |
| S = {s: seq_feats(s) for s in pd.unique(sequences)} |
| M = {s: smi_feats(s) for s in pd.unique(smiles)} |
| return np.hstack([np.stack([S[s] for s in sequences]), np.stack([M[s] for s in smiles])]) |
|
|
|
|
| |
| PRIOR_K = 5 |
|
|
|
|
| def prior_table(smiles, Y): |
| """Per-substrate label statistics from training rows -> DataFrame indexed by SMILES.""" |
| d = pd.DataFrame(Y, columns=TARGETS) |
| g = d.groupby(np.asarray(smiles)) |
| m, n = g.mean(), g.size() |
| gm = Y.mean(0) |
| tab = (m.mul(n, axis=0) + gm * PRIOR_K).div(n + PRIOR_K, axis=0).add_prefix('shrunk_') |
| tab['log_count'] = np.log1p(n) |
| for agg in ['std', 'max', 'min', 'median']: |
| tab = tab.join(getattr(g, agg)().add_prefix(agg + '_')) |
| tab.attrs['global_mean'] = gm.tolist() |
| return tab |
|
|
|
|
| def prior_features(tab, smiles, global_mean): |
| """Look up prior features; unseen substrates get the global mean, count 0 and NaN spread stats.""" |
| out = tab.reindex(np.asarray(smiles)).to_numpy(dtype=np.float64).copy() |
| unseen = np.isnan(out[:, 0]) |
| out[unseen, 0], out[unseen, 1], out[unseen, 2] = global_mean[0], global_mean[1], 0.0 |
| return out |
|
|
|
|
| |
| def load_esm(model_dir, device='cpu', quantize=False, layer=None): |
| """layer=k keeps only the first k transformer layers and drops the final LayerNorm, so last_hidden_state equals |
| hidden_states[k] of the full model (an intermediate layer) and inference is faster.""" |
| import torch |
| from transformers import AutoTokenizer, EsmModel |
| tok = AutoTokenizer.from_pretrained(model_dir) |
| model = EsmModel.from_pretrained(model_dir, add_pooling_layer=False) |
| if layer: |
| model.encoder.layer = model.encoder.layer[:layer] |
| model.encoder.emb_layer_norm_after = None |
| model.contact_head = None |
| model.config.num_hidden_layers = layer |
| model = model.to(device).eval() |
| if quantize: |
| model = torch.ao.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8) |
| return tok, model |
|
|
|
|
| def esm_embed(seqs, tok, model, device='cpu', batch_tokens=2000, window=1022): |
| """Mean-pooled last-layer embedding per sequence (special tokens excluded). |
| Sequences longer than `window` are split into windows whose embeddings are averaged.""" |
| import torch |
| pieces, owner = [], [] |
| for i, s in enumerate(seqs): |
| s = s.upper() |
| for j in range(0, max(1, len(s)), window): |
| pieces.append(s[j:j + window]); owner.append(i) |
| order = np.argsort([len(p) for p in pieces])[::-1] |
| out = np.zeros((len(pieces), model.config.hidden_size), dtype=np.float32) |
| k = 0 |
| with torch.no_grad(): |
| while k < len(order): |
| n = max(1, batch_tokens // (len(pieces[order[k]]) + 2)) |
| idx = order[k:k + n]; k += n |
| enc = tok([pieces[i] for i in idx], return_tensors='pt', padding=True).to(device) |
| h = model(**enc).last_hidden_state.float() |
| m = enc['attention_mask'].clone(); m[:, 0] = 0 |
| m[torch.arange(len(idx)), enc['attention_mask'].sum(1) - 1] = 0 |
| out[idx] = ((h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True)).cpu().numpy() |
| owner = np.array(owner) |
| return np.stack([out[owner == i].mean(0) for i in range(len(seqs))]) |
|
|
