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