| """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) |
| enc = {k: v.to(device) for k, v in enc.items()} |
| 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), device=m.device), 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 cluster_by_cosine(seqs, thr=0.35, max_frac=0.2, log=print): |
| """Group sequences that are even remotely related (3-mer cosine similarity >= thr), |
| so a local holdout looks like the real validation set (whose enzymes are unrelated to training). |
| If one giant group appears, the threshold is raised automatically.""" |
| from sklearn.feature_extraction.text import CountVectorizer |
| from sklearn.preprocessing import normalize |
| from scipy.sparse import coo_matrix |
| from scipy.sparse.csgraph import connected_components |
| cv = CountVectorizer(analyzer="char", ngram_range=(3, 3), lowercase=False, dtype=np.float32) |
| A = normalize(cv.fit_transform(seqs)).tocsr() |
| AT = A.T.tocsr() |
| n = A.shape[0] |
| rows, cols, vals = [], [], [] |
| for i in range(0, n, 500): |
| S = (A[i:i + 500] @ AT).toarray() |
| r, c = np.nonzero(S >= 0.30) |
| keep = (r + i) < c |
| rows.append(r[keep] + i); cols.append(c[keep]); vals.append(S[r[keep], c[keep]]) |
| rows, cols, vals = np.concatenate(rows), np.concatenate(cols), np.concatenate(vals) |
| t = thr |
| while True: |
| m = vals >= t |
| g = coo_matrix((np.ones(int(m.sum())), (rows[m], cols[m])), shape=(n, n)) |
| _, lab = connected_components(g, directed=False) |
| big = np.bincount(lab).max() / n |
| if big <= max_frac or t >= 0.9: |
| break |
| t += 0.05 |
| log(f"grouping threshold used: {t:.2f} | groups: {len(set(lab))} | biggest group: {big:.1%} of sequences") |
| return lab |
|
|
|
|
| |
| def te_stats(df): |
| """per-substrate: row count, mean and std of log10_kcat / log10_km in the training rows""" |
| g = df.groupby("smiles").agg(n=("log10_kcat", "size"), |
| kc=("log10_kcat", "mean"), km=("log10_km", "mean"), |
| kc_sd=("log10_kcat", "std"), km_sd=("log10_km", "std")) |
| 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, |
| s["kc_sd"].values, s["km_sd"].values]).astype(np.float32) |
|
|
|
|
| def sub_neighbor_feats(q_bits, ref_bits, ref_vals, k=5): |
| """For each query substrate: how similar are the k most similar TRAINING substrates |
| (Tanimoto on Morgan bits) and what are their weighted-mean kcat / km. |
| This helps for the ~1/3 of validation substrates never seen in training.""" |
| nq = q_bits.shape[0] |
| out = np.full((nq, 4), np.nan, dtype=np.float32) |
| rs = ref_bits.sum(1) |
| k = min(k, ref_bits.shape[0] - 1) |
| for i in range(0, nq, 512): |
| qb = q_bits[i:i + 512] |
| inter = qb @ ref_bits.T |
| union = qb.sum(1)[:, None] + rs[None, :] - inter |
| sim = inter / np.maximum(union, 1) |
| idx = np.argpartition(-sim, k, axis=1)[:, :k] |
| s = np.take_along_axis(sim, idx, 1) |
| w = s ** 2 + 1e-6 |
| vals = ref_vals[idx] |
| mean = (vals * w[:, :, None]).sum(1) / w.sum(1)[:, None] |
| out[i:i + 512, 0] = s.max(1) |
| out[i:i + 512, 1:3] = mean |
| out[i:i + 512, 3] = s.mean(1) |
| return out |
|
|
|
|
| def te_full(smiles_values, smi_index, bits_u, stats, ref_bits): |
| """9 substrate-level features per row: exact-substrate stats (5) + similar-substrate stats (4). |
| smi_index: dict smiles -> row of bits_u; ref_bits aligned with stats.index.""" |
| exact = te_apply(smiles_values, stats) |
| uq = list(dict.fromkeys(smiles_values)) |
| q = bits_u[[smi_index[s] for s in uq]] |
| ref_vals = stats[["kc", "km"]].values.astype(np.float32) |
| nb = sub_neighbor_feats(q, ref_bits, ref_vals) |
| pos = {s: i for i, s in enumerate(uq)} |
| nb_rows = nb[[pos[s] for s in smiles_values]] |
| return np.hstack([exact, nb_rows]).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) |
|
|
|
|
| |
| def _knn_one(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, k=5, temp=5.0): |
| """For each query row, look at training enzymes that ESM thinks are similar. |
| cols 0-2: among training rows with the SAME substrate: best similarity, weighted mean kcat, weighted mean km |
| cols 3-5: among ALL training enzymes (per-enzyme mean targets): best similarity, mean kcat, mean km |
| q_si / ref_si = row index of each row's sequence inside q_embn / ref_embn (centered + length-normalised).""" |
| from collections import defaultdict |
| nq = len(q_smiles) |
| out = np.full((nq, 6), np.nan, dtype=np.float32) |
| q_si = np.asarray(q_si); ref_si = np.asarray(ref_si); ref_y = np.asarray(ref_y, dtype=np.float32) |
| ref_by = defaultdict(list) |
| for j, s in enumerate(ref_smiles): |
| ref_by[s].append(j) |
| q_by = defaultdict(list) |
| for i, s in enumerate(q_smiles): |
| q_by[s].append(i) |
| for s, qi in q_by.items(): |
| rj = ref_by.get(s) |
| if not rj: |
| continue |
| rj = np.array(rj); qi = np.array(qi) |
| sim = q_embn[q_si[qi]] @ ref_embn[ref_si[rj]].T |
| kk = min(k, len(rj)) |
| idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk] |
| sv = np.take_along_axis(sim, idx, 1) |
| w = np.exp(temp * sv) |
| y = ref_y[rj][idx] |
| out[qi, 0] = sv.max(1) |
| out[qi, 1:3] = (y * w[:, :, None]).sum(1) / w.sum(1)[:, None] |
| |
| n_ref = ref_embn.shape[0] |
| ysum = np.zeros((n_ref, 2)); cnt = np.zeros(n_ref) |
| np.add.at(ysum, ref_si, ref_y); np.add.at(cnt, ref_si, 1) |
| have = np.nonzero(cnt)[0] |
| ref_seq_y = (ysum[have] / cnt[have][:, None]).astype(np.float32) |
| R = ref_embn[have] |
| uq = np.unique(q_si) |
| G = np.full((len(uq), 3), np.nan, dtype=np.float32) |
| kk = min(k, len(have)) |
| for i in range(0, len(uq), 1024): |
| sim = q_embn[uq[i:i + 1024]] @ R.T |
| idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk] |
| sv = np.take_along_axis(sim, idx, 1) |
| w = np.exp(temp * sv) |
| y = ref_seq_y[idx] |
| G[i:i + 1024, 0] = sv.max(1) |
| G[i:i + 1024, 1:3] = (y * w[:, :, None]).sum(1) / w.sum(1)[:, None] |
| pos = {int(s): i for i, s in enumerate(uq)} |
| out[:, 3:6] = G[[pos[int(s)] for s in q_si]] |
| return out |
|
|
|
|
| def normalize_emb(emb, mean): |
| c = emb - mean |
| return (c / np.maximum(np.linalg.norm(c, axis=1, keepdims=True), 1e-6)).astype(np.float32) |
|
|
|
|
| def knn_feats(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, ks=(5,), temp=5.0): |
| """6 columns per value of k (see _knn_one).""" |
| return np.hstack([_knn_one(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, k=k, temp=temp) |
| for k in ks]) |
|
|
