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