File size: 6,828 Bytes
b145940 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Shared feature code. train.py and predict.py BOTH import this file,
so features are guaranteed to be identical at training and inference time."""
import time
import zlib
from collections import Counter
import numpy as np
AA = "ACDEFGHIKLMNPQRSTVWY"
_LUT = np.full(256, -1, dtype=np.int16)
for _i, _a in enumerate(AA):
_LUT[ord(_a)] = _i
N_MORGAN = 1024
# ----------------------------------------------------------------- SMILES ---
def _smiles_vec_factory():
from rdkit import Chem, RDLogger
from rdkit.Chem import (Descriptors, MACCSkeys, rdFingerprintGenerator,
rdMolDescriptors)
RDLogger.DisableLog("rdApp.*")
gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=N_MORGAN)
n_desc = 12
dim = N_MORGAN + 167 + n_desc
def vec(smi):
out = np.zeros(dim, dtype=np.float32)
try:
mol = Chem.MolFromSmiles(smi)
if mol is None:
return out
out[:N_MORGAN] = gen.GetCountFingerprintAsNumPy(mol)
out[N_MORGAN:N_MORGAN + 167] = np.array(list(MACCSkeys.GenMACCSKeys(mol)), dtype=np.float32)
d = [
Descriptors.MolWt(mol), Descriptors.MolLogP(mol), Descriptors.TPSA(mol),
Descriptors.NumHDonors(mol), Descriptors.NumHAcceptors(mol),
Descriptors.NumRotatableBonds(mol), rdMolDescriptors.CalcNumRings(mol),
rdMolDescriptors.CalcNumAromaticRings(mol), mol.GetNumHeavyAtoms(),
Chem.GetFormalCharge(mol), rdMolDescriptors.CalcFractionCSP3(mol),
len(Chem.GetMolFrags(mol)),
]
out[N_MORGAN + 167:] = np.array(d, dtype=np.float32)
except Exception:
pass
return out
return vec, dim
def smiles_features(unique_smiles):
"""-> (matrix [n_unique, dim]) in the same order as unique_smiles"""
vec, dim = _smiles_vec_factory()
return np.stack([vec(s) for s in unique_smiles]).astype(np.float32)
# --------------------------------------------------------------- SEQUENCE ---
def _seq_vec(s):
a = _LUT[np.frombuffer(s.encode("ascii", "ignore"), dtype=np.uint8)]
a = a[a >= 0].astype(np.int64)
n = len(a)
comp = np.bincount(a, minlength=20) / max(n, 1)
if n > 1:
di = np.bincount(a[:-1] * 20 + a[1:], minlength=400) / (n - 1)
else:
di = np.zeros(400)
return np.concatenate([[len(s), np.log1p(len(s))], comp, di]).astype(np.float32)
def seq_features(unique_seqs):
return np.stack([_seq_vec(s) for s in unique_seqs]).astype(np.float32)
# -------------------------------------------------------------------- ESM ---
def esm_embed(seqs, model_dir, fallback=None, max_len=800, budget_s=None,
batch_tokens=4096, n_threads=2, device="cpu", log=print):
"""Mean-pooled ESM2 embedding for each sequence (list of str).
Short sequences are done first. If budget_s runs out, the remaining
sequences get `fallback` (the training-mean embedding) so we never time out."""
import torch
from transformers import AutoTokenizer, EsmModel
torch.set_num_threads(n_threads)
tok = AutoTokenizer.from_pretrained(model_dir)
model = EsmModel.from_pretrained(model_dir).to(device).eval()
dim = model.config.hidden_size
out = np.full((len(seqs), dim), np.nan, dtype=np.float32)
order = np.argsort([len(s) for s in seqs])
t0 = time.time()
i = 0
while i < len(order):
if budget_s is not None and time.time() - t0 > budget_s:
log(f"[esm] time budget hit after {i}/{len(order)} sequences; using fallback for the rest")
break
j, maxl = i, 0
while j < len(order):
l = min(len(seqs[order[j]]), max_len) + 2
if j > i and (j - i + 1) * max(maxl, l) > batch_tokens:
break
maxl = max(maxl, l)
j += 1
idx = order[i:j]
batch = [seqs[k][:max_len] for k in idx]
enc = tok(batch, return_tensors="pt", padding=True)
with torch.inference_mode():
h = model(**enc).last_hidden_state
m = enc["attention_mask"].clone().float()
lens = enc["attention_mask"].sum(1)
m[:, 0] = 0
m[torch.arange(len(idx)), lens - 1] = 0
e = (h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True).clamp(min=1)
out[idx] = e.cpu().numpy()
i = j
if (i // 50) != ((i - len(idx)) // 50):
log(f"[esm] {i}/{len(order)} done, {time.time() - t0:.0f}s")
missing = np.isnan(out[:, 0])
if missing.any():
out[missing] = fallback if fallback is not None else 0.0
return out
# ------------------------------------------------------ SEQUENCE CLUSTERS ---
def cluster_sequences(seqs, k=6, mod=2, min_frac=0.3):
"""Cheap grouping of similar sequences (mutants / close homologs).
Used ONLY to make an honest local validation split. Returns a group id per sequence."""
order = sorted(range(len(seqs)), key=lambda i: -len(seqs[i]))
kmer2leaders = {}
leader_of = [0] * len(seqs)
for i in order:
s = seqs[i]
kms = {s[j:j + k] for j in range(len(s) - k + 1)
if zlib.crc32(s[j:j + k].encode()) % mod == 0}
hits = Counter()
for km in kms:
for L in kmer2leaders.get(km, ()):
hits[L] += 1
best = None
if hits:
L, c = hits.most_common(1)[0]
if c >= min_frac * max(len(kms), 1):
best = L
if best is None:
leader_of[i] = i
for km in kms:
kmer2leaders.setdefault(km, []).append(i)
else:
leader_of[i] = best
return np.array(leader_of)
# ---------------------------------------- SUBSTRATE "AVERAGE" FEATURES ---
def te_stats(df):
"""per-substrate: how many training rows, mean log10_kcat, mean log10_km"""
g = df.groupby("smiles").agg(n=("log10_kcat", "size"),
kc=("log10_kcat", "mean"),
km=("log10_km", "mean"))
return g
def te_apply(smiles, stats):
s = stats.reindex(smiles)
return np.column_stack([s["n"].fillna(0).values, s["kc"].values, s["km"].values]).astype(np.float32)
# ----------------------------------------------------------------- ASSEMBLE ---
def assemble(df, smi_unique, smi_mat, seq_unique, seq_mat, esm_mat=None, te=None):
smi_idx = {s: i for i, s in enumerate(smi_unique)}
seq_idx = {s: i for i, s in enumerate(seq_unique)}
si = df["smiles"].map(smi_idx).values
qi = df["sequence"].map(seq_idx).values
parts = [smi_mat[si], seq_mat[qi]]
if esm_mat is not None:
parts.append(esm_mat[qi])
if te is not None:
parts.append(te)
return np.hstack(parts).astype(np.float32)
|